nub/internal/llm/openai/stream.go

package openai

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"

	"nub/internal/llm"
)

// partialCall akkumuliert die fragmentierten tool_calls-Deltas. `index` ist
// die einzige verlässliche Korrelation (5.1 Punkt 1); id/name kommen
// typischerweise nur im ersten Chunk.
type partialCall struct {
	id       string
	name     string
	ourIndex int
	started  bool
}

// decoder liest einen OpenAI-kompatiblen SSE-Chatstream und synthetisiert
// unser internes Event-Schema. Tolerant gegenüber Server-Eigenheiten:
// kompletter Call in einem Chunk, falsch gezählter index, nachträglicher name.
type decoder struct {
	scanner    *bufio.Scanner
	calls      map[int]*partialCall
	nextIndex  int
	textIndex  int
	textOpen   bool
	thinkIndex int
	thinkOpen  bool
	usage      wireUsage
	haveUsage  bool
	finishStop string
	done       bool
}

func newDecoder(body io.Reader) *decoder {
	s := bufio.NewScanner(body)
	s.Buffer(make([]byte, 64*1024), 8*1024*1024)
	return &decoder{scanner: s, calls: make(map[int]*partialCall)}
}

// next liest den nächsten SSE-Frame und liefert die daraus resultierenden
// Events. Ein Frame kann 0, 1 oder mehrere Events erzeugen (z.B. BlockStart
// + BlockDelta beim ersten Tool-Call-Fragment). more=false signalisiert das
// Ende des Streams (nach dem letzten Done-Event).
func (d *decoder) next() (events []llm.Event, more bool, err error) {
	if d.done {
		return nil, false, nil
	}
	for d.scanner.Scan() {
		line := d.scanner.Text()
		if line == "" || bytes.HasPrefix([]byte(line), []byte(":")) {
			continue // Leerzeile oder SSE-Kommentar/Keep-Alive
		}
		data, ok := cutPrefix(line, "data: ")
		if !ok {
			data, ok = cutPrefix(line, "data:")
		}
		if !ok {
			continue
		}
		data = trimSpace(data)
		if data == "[DONE]" {
			d.done = true
			return d.finalEvents(), true, nil
		}

		var streamErr wireStreamError
		if err := json.Unmarshal([]byte(data), &streamErr); err == nil && streamErr.Error != nil {
			d.done = true
			return nil, false, fmt.Errorf("stream error: %s (%s)", streamErr.Error.Message, streamErr.Error.Type)
		}

		var chunk wireChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			d.done = true
			return nil, false, fmt.Errorf("decode chunk: %w", err)
		}
		return d.applyChunk(chunk), true, nil
	}
	if err := d.scanner.Err(); err != nil {
		d.done = true
		return nil, false, err
	}
	// Stream endete ohne explizites [DONE] (manche Server tun das) -> Finalize.
	d.done = true
	return d.finalEvents(), true, nil
}

func (d *decoder) applyChunk(chunk wireChunk) []llm.Event {
	var events []llm.Event
	if chunk.Usage != nil {
		d.usage = *chunk.Usage
		d.haveUsage = true
	}
	if len(chunk.Choices) == 0 {
		return events
	}
	choice := chunk.Choices[0]

	if choice.Delta.Content != nil && *choice.Delta.Content != "" {
		if !d.textOpen {
			d.textIndex = d.nextIndex
			d.nextIndex++
			d.textOpen = true
			events = append(events, llm.BlockStart{Index: d.textIndex, Block: llm.Block{Kind: llm.KindText}})
		}
		events = append(events, llm.BlockDelta{Index: d.textIndex, Text: *choice.Delta.Content})
	}

	if choice.Delta.ReasoningContent != nil && *choice.Delta.ReasoningContent != "" {
		if !d.thinkOpen {
			d.thinkIndex = d.nextIndex
			d.nextIndex++
			d.thinkOpen = true
			events = append(events, llm.BlockStart{Index: d.thinkIndex, Block: llm.Block{Kind: llm.KindThinking}})
		}
		events = append(events, llm.BlockDelta{Index: d.thinkIndex, Text: *choice.Delta.ReasoningContent})
	}

	for _, tc := range choice.Delta.ToolCalls {
		pc, ok := d.calls[tc.Index]
		if !ok {
			pc = &partialCall{}
			d.calls[tc.Index] = pc
		}
		if tc.ID != nil {
			pc.id = *tc.ID
		}
		if tc.Function != nil && tc.Function.Name != nil {
			pc.name += *tc.Function.Name
		}
		if !pc.started && pc.id != "" && pc.name != "" {
			pc.started = true
			pc.ourIndex = d.nextIndex
			d.nextIndex++
			events = append(events, llm.BlockStart{
				Index: pc.ourIndex,
				Block: llm.Block{Kind: llm.KindToolUse, ID: pc.id, Name: pc.name},
			})
		}
		if tc.Function != nil && tc.Function.Arguments != nil && pc.started {
			events = append(events, llm.BlockDelta{Index: pc.ourIndex, PartialJSON: *tc.Function.Arguments})
		}
	}

	if choice.FinishReason != nil {
		d.finishStop = *choice.FinishReason
	}

	return events
}

// finalEvents schließt alle offenen Blöcke und liefert das abschließende
// Done-Event.
func (d *decoder) finalEvents() []llm.Event {
	var events []llm.Event
	if d.textOpen {
		events = append(events, llm.BlockStop{Index: d.textIndex})
		d.textOpen = false
	}
	if d.thinkOpen {
		events = append(events, llm.BlockStop{Index: d.thinkIndex})
		d.thinkOpen = false
	}
	for _, pc := range d.calls {
		if pc.started {
			events = append(events, llm.BlockStop{Index: pc.ourIndex})
		}
	}

	stop := mapStopReason(d.finishStop)
	usage := llm.Usage{}
	if d.haveUsage {
		usage.InputTokens = d.usage.PromptTokens
		usage.OutputTokens = d.usage.CompletionTokens
		if d.usage.PromptTokensDetails != nil {
			usage.CacheReadTokens = d.usage.PromptTokensDetails.CachedTokens
		}
	}
	events = append(events, llm.Done{Stop: stop, Usage: usage})
	return events
}

func mapStopReason(finish string) llm.StopReason {
	switch finish {
	case "tool_calls":
		return llm.StopToolUse
	case "length":
		return llm.StopMaxTokens
	default:
		return llm.StopEnd
	}
}

func cutPrefix(s, prefix string) (string, bool) {
	if len(s) < len(prefix) || s[:len(prefix)] != prefix {
		return "", false
	}
	return s[len(prefix):], true
}

func trimSpace(s string) string {
	for len(s) > 0 && (s[0] == ' ' || s[0] == '\t') {
		s = s[1:]
	}
	for len(s) > 0 && (s[len(s)-1] == '\r' || s[len(s)-1] == '\n') {
		s = s[:len(s)-1]
	}
	return s
}