Cyrene_v2/internal/runtime/model/openai_stream.go

package model

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

// openAIToolAcc 累积单个流式工具调用的部分状态（按
// 提供商的 tool_calls[].index 作键），以便其参数跨
// 增量片段重新拼装。started 保证每块只发出一次 ToolCallStart。
type openAIToolAcc struct {
	id      string
	name    string
	args    strings.Builder
	started bool
}

// idOrName 解出用于关联 start/delta/end 的块标识：
// 有提供商工具调用 id 时用它，否则用函数名。
func (a *openAIToolAcc) idOrName() string {
	if a.id != "" {
		return a.id
	}
	return a.name
}

// openAIStreamState 是一次流式 /chat/completions 响应的
// 运行中聚合状态：拼装好的 Response、当前打开片段的
// 包裹标志，以及按索引的工具调用累加器。它是
// Anthropic 的 streamBlock 记账的对位实现，因此两个适配器
// 共用同一套「解析 → 应用 → 收尾」形态。emit 是容忍 nil 的事件槽
// （onEvent 为 nil 时只是丢弃每个事件）。
type openAIStreamState struct {
	agg             Response
	text            strings.Builder
	textStarted     bool
	thinkingStarted bool
	toolCalls       map[int]*openAIToolAcc
	toolOrder       []int
	onEvent         func(StreamEvent)
}

// scanSSE 增量解码流块，用 strings.Builder 累积完整正文
// 文本（绝不一次性读完整个响应体），并从最后一个块
// 捕获用量。它把每个增量拆到其片段并按对包裹：
// reasoning_content 作为 thinking_start/delta/end，tool_calls 片段作为
// tool_call_start/delta/end，正文文本作为 text_start/delta/end。
// 累积的工具调用参数会跨片段重新拼装成结构化的
// Args map，并与文本一起挂到返回的 Response 上。
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 转发一个结构化事件，容忍 nil 事件槽。
func (st *openAIStreamState) emit(ev StreamEvent) {
	if st.onEvent != nil {
		st.onEvent(ev)
	}
}

// applyChunk 把单个已解码的流块折叠进运行中的聚合状态：
// 用量与停止原因即使在无 choice 的终止块上也会被记录，
// 片段增量（若有）则被路由到其处理器。
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 把一个增量片段路由到其片段处理器。优先级
// 取自提供商：推理片段压过工具调用片段，后者又压过
// 正文文本，因此一个增量绝不会同时进入两个片段。
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 转发一个推理片段，只在第一个
// 片段上打开 thinking 片段。
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 关闭已打开的 thinking 片段（未打开时是空操作），
// 这样随后的工具调用或文本片段会开一个新的。
func (st *openAIStreamState) endThinking() {
	if !st.thinkingStarted {
		return
	}
	st.thinkingStarted = false
	st.emit(StreamEvent{Kind: StreamKindThinkingEnd})
}

// applyToolCall 把一个流式工具调用片段折叠进它按索引的
// 累加器，在 id/name 已知时打开该块的 start 事件，
// 并随着参数片段到达而转发它们。
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 追加一个正文文本片段，只在第一个
// 片段上打开文本片段。
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 在流终止时发出每个仍打开片段的结束
// 事件。OpenAI 没有按块的停止事件，所以每个已开始的 tool_call
// 块都在此关闭。
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 按提供商顺序把每个累积的工具调用参数缓冲
// 解码成结构化的 Args map。无法解析的片段降级为
// 空 map（该调用仍会浮出）。
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 关闭打开的片段并拼装聚合后的
// Response：拼接的正文文本（推理永不泄漏进去）加上
// 重新拼装的工具调用。
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
}