Cyrene_v2/internal/runtime/model/openai_stream.go

package model

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"strings"
)

// openAIToolAcc accumulates the partial state of one streaming tool call (keyed
// by the provider's tool_calls[].index) so its arguments re-assemble across the
// delta fragments. started guards a single ToolCallStart emission per block.
type openAIToolAcc struct {
	id      string
	name    string
	args    strings.Builder
	started bool
}

// idOrName resolves the block identifier used to correlate start/delta/end: the
// provider tool-call id when present, else the function name.
func (a *openAIToolAcc) idOrName() string {
	if a.id != "" {
		return a.id
	}
	return a.name
}

// openAIStreamState is the running aggregation of one streamed
// /chat/completions response: the assembled Response, the bracketing flags of the
// currently open segment, and the per-index tool-call accumulators. It is the
// OpenAI counterpart of anthropic's streamBlock bookkeeping, so both adapters
// share one parse → apply → finalize shape. emit is the nil-tolerant event sink
// (a nil onEvent simply drops every event).
type openAIStreamState struct {
	agg             Response
	text            strings.Builder
	textStarted     bool
	thinkingStarted bool
	toolCalls       map[int]*openAIToolAcc
	toolOrder       []int
	onEvent         func(StreamEvent)
}

// scanSSE incrementally decodes stream chunks, accumulates the full body text
// with a strings.Builder (never reading the whole body at once) and captures
// usage from the final chunk. It splits each delta into its segment and brackets
// them: reasoning_content as thinking_start/delta/end, tool_calls fragments as
// tool_call_start/delta/end, and body text as text_start/delta/end. The
// accumulated tool-call arguments re-assemble across fragments into a structured
// Args map, and are attached to the returned Response alongside the text.
func scanSSE(r io.Reader, onEvent func(StreamEvent)) (Response, error) {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
	st := &openAIStreamState{toolCalls: map[int]*openAIToolAcc{}, onEvent: onEvent}

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
		if data == sseDone {
			break
		}
		var chunk openAIChatResponse
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
		}
		st.applyChunk(chunk)
	}
	if err := scanner.Err(); err != nil {
		return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
	}
	return finalizeOpenAIStream(st), nil
}

// emit forwards one structured event, tolerating a nil sink.
func (st *openAIStreamState) emit(ev StreamEvent) {
	if st.onEvent != nil {
		st.onEvent(ev)
	}
}

// applyChunk folds a single decoded stream chunk into the running aggregation:
// usage and stop reason are recorded even on a choice-less terminal chunk, and
// the segment delta (if any) is routed to its handler.
func (st *openAIStreamState) applyChunk(chunk openAIChatResponse) {
	if u := chunk.Usage; u != nil {
		st.agg.PromptTokens = u.PromptTokens
		st.agg.CompletionTokens = u.CompletionTokens
		st.agg.TotalTokens = u.TotalTokens
	}
	if len(chunk.Choices) == 0 {
		return
	}
	if chunk.Choices[0].FinishReason != "" {
		st.agg.StopReason = chunk.Choices[0].FinishReason
	}
	if chunk.Choices[0].Delta == nil {
		return
	}
	st.applyDelta(chunk.Choices[0].Delta)
}

// applyDelta routes one delta fragment to its segment handler. The precedence is
// the provider's: a reasoning fragment wins over a tool-call fragment, which wins
// over body text, so a delta never contributes to two segments at once.
func (st *openAIStreamState) applyDelta(delta *openAIDelta) {
	switch {
	case delta.ReasoningContent != "":
		st.applyThinking(delta.ReasoningContent)
	case len(delta.ToolCalls) > 0:
		st.endThinking()
		for _, tc := range delta.ToolCalls {
			st.applyToolCall(tc)
		}
	case delta.Content != "":
		st.endThinking()
		st.applyText(delta.Content)
	}
}

// applyThinking forwards one reasoning fragment, opening the thinking segment on
// the first fragment only.
func (st *openAIStreamState) applyThinking(fragment string) {
	if !st.thinkingStarted {
		st.thinkingStarted = true
		st.emit(StreamEvent{Kind: StreamKindThinkingStart})
	}
	st.emit(StreamEvent{Kind: StreamKindThinkingDelta, Delta: fragment})
}

// endThinking closes the open thinking segment (no-op when none is open), so a
// following tool-call or text segment starts a fresh one.
func (st *openAIStreamState) endThinking() {
	if !st.thinkingStarted {
		return
	}
	st.thinkingStarted = false
	st.emit(StreamEvent{Kind: StreamKindThinkingEnd})
}

// applyToolCall folds one streaming tool-call fragment into its per-index
// accumulator, opening the block's start event once its id/name is known and
// forwarding the argument fragments as they arrive.
func (st *openAIStreamState) applyToolCall(tc openAIToolCallDelta) {
	acc := st.toolCalls[tc.Index]
	if acc == nil {
		acc = &openAIToolAcc{}
		st.toolCalls[tc.Index] = acc
		st.toolOrder = append(st.toolOrder, tc.Index)
	}
	if acc.id == "" {
		acc.id = tc.ID
	}
	if acc.name == "" {
		acc.name = tc.Function.Name
	}
	if !acc.started && acc.idOrName() != "" {
		acc.started = true
		st.emit(StreamEvent{Kind: StreamKindToolCallStart, Index: tc.Index, ID: acc.idOrName()})
	}
	if tc.Function.Arguments == "" {
		return
	}
	acc.args.WriteString(tc.Function.Arguments)
	st.emit(StreamEvent{Kind: StreamKindToolCallDelta, Index: tc.Index, ID: acc.idOrName(), Delta: tc.Function.Arguments})
}

// applyText appends one body-text fragment, opening the text segment on the first
// fragment only.
func (st *openAIStreamState) applyText(fragment string) {
	st.text.WriteString(fragment)
	if !st.textStarted {
		st.textStarted = true
		st.emit(StreamEvent{Kind: StreamKindTextStart})
	}
	st.emit(StreamEvent{Kind: StreamKindTextDelta, Delta: fragment})
}

// closeSegments emits the end event of every segment still open at stream
// termination. OpenAI has no per-block stop event, so each started tool_call
// block is closed here.
func (st *openAIStreamState) closeSegments() {
	if st.thinkingStarted {
		st.emit(StreamEvent{Kind: StreamKindThinkingEnd})
	}
	if st.textStarted {
		st.emit(StreamEvent{Kind: StreamKindTextEnd})
	}
	for _, idx := range st.toolOrder {
		acc := st.toolCalls[idx]
		if acc == nil || !acc.started {
			continue
		}
		st.emit(StreamEvent{Kind: StreamKindToolCallEnd, Index: idx, ID: acc.idOrName()})
	}
}

// assembleToolCalls decodes each accumulated tool-call argument buffer into its
// structured Args map, in provider order. Unparseable fragments degrade to an
// empty map (the call is still surfaced).
func (st *openAIStreamState) assembleToolCalls() {
	for _, idx := range st.toolOrder {
		acc := st.toolCalls[idx]
		if acc == nil {
			continue
		}
		args := map[string]any{}
		if acc.args.Len() > 0 {
			if err := json.Unmarshal([]byte(acc.args.String()), &args); err != nil {
				args = map[string]any{}
			}
		}
		st.agg.ToolCalls = append(st.agg.ToolCalls, ToolCall{Name: acc.name, ID: acc.id, Args: args})
	}
}

// finalizeOpenAIStream closes the open segments and assembles the aggregated
// Response: the concatenated body text (reasoning never leaks into it) plus the
// re-assembled tool calls.
func finalizeOpenAIStream(st *openAIStreamState) Response {
	st.closeSegments()
	st.assembleToolCalls()
	st.agg.Content = st.text.String()
	st.agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: st.agg.Content}}}
	return st.agg
}