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
}

// 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)
	agg := Response{}
	var text strings.Builder
	textStarted := false
	thinkingStarted := false
	toolCalls := map[int]*openAIToolAcc{}
	var toolOrder []int

	endThinking := func() {
		if onEvent != nil && thinkingStarted {
			thinkingStarted = false
			onEvent(StreamEvent{Kind: StreamKindThinkingEnd})
		}
	}

	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)
		}
		if u := chunk.Usage; u != nil {
			agg.PromptTokens = u.PromptTokens
			agg.CompletionTokens = u.CompletionTokens
			agg.TotalTokens = u.TotalTokens
		}
		if len(chunk.Choices) > 0 && chunk.Choices[0].FinishReason != "" {
			agg.StopReason = chunk.Choices[0].FinishReason
		}
		if len(chunk.Choices) == 0 || chunk.Choices[0].Delta == nil {
			continue
		}
		delta := chunk.Choices[0].Delta

		// Reasoning / thinking segment (OpenAI o-series, DeepSeek, DashScope).
		if delta.ReasoningContent != "" {
			if onEvent != nil && !thinkingStarted {
				thinkingStarted = true
				onEvent(StreamEvent{Kind: StreamKindThinkingStart})
			}
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindThinkingDelta, Delta: delta.ReasoningContent})
			}
			continue
		}

		// Streaming tool-call fragments.
		if len(delta.ToolCalls) > 0 {
			endThinking()
			for _, tc := range delta.ToolCalls {
				acc := toolCalls[tc.Index]
				if acc == nil {
					acc = &openAIToolAcc{}
					toolCalls[tc.Index] = acc
					toolOrder = append(toolOrder, tc.Index)
				}
				if acc.id == "" {
					acc.id = tc.ID
				}
				if acc.name == "" {
					acc.name = tc.Function.Name
				}
				if onEvent != nil && !acc.started && acc.idOrName() != "" {
					acc.started = true
					onEvent(StreamEvent{Kind: StreamKindToolCallStart, Index: tc.Index, ID: acc.idOrName()})
				}
				if tc.Function.Arguments != "" {
					acc.args.WriteString(tc.Function.Arguments)
					if onEvent != nil {
						onEvent(StreamEvent{Kind: StreamKindToolCallDelta, Index: tc.Index, ID: acc.idOrName(), Delta: tc.Function.Arguments})
					}
				}
			}
			continue
		}

		// Body-text segment.
		if delta.Content == "" {
			continue
		}
		endThinking()
		text.WriteString(delta.Content)
		if onEvent != nil {
			if !textStarted {
				textStarted = true
				onEvent(StreamEvent{Kind: StreamKindTextStart})
			}
			onEvent(StreamEvent{Kind: StreamKindTextDelta, Delta: delta.Content})
		}
	}
	if err := scanner.Err(); err != nil {
		return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
	}

	if onEvent != nil {
		if thinkingStarted {
			onEvent(StreamEvent{Kind: StreamKindThinkingEnd})
		}
		if textStarted {
			onEvent(StreamEvent{Kind: StreamKindTextEnd})
		}
		// OpenAI has no per-block stop event, so each tool_call block is closed
		// when the stream terminates.
		for _, idx := range toolOrder {
			acc := toolCalls[idx]
			if acc == nil || !acc.started {
				continue
			}
			onEvent(StreamEvent{Kind: StreamKindToolCallEnd, Index: idx, ID: acc.idOrName()})
		}
	}

	for _, idx := range toolOrder {
		acc := 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{}
			}
		}
		agg.ToolCalls = append(agg.ToolCalls, ToolCall{Name: acc.name, ID: acc.id, Args: args})
	}
	agg.Content = text.String()
	agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: agg.Content}}}
	return agg, nil
}