Cyrene_v2/internal/runtime/model/anthropic_stream.go

package model

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

// anthropicStreamEvent 是 /v1/messages（stream: true）的一个 SSE 负载。
// Anthropic 既在 SSE 的 `event:` 字段里、也在 JSON 的 `type` 字段里
// 携带事件种类，所以解析器按 JSON 负载的 type 取值。
type anthropicStreamEvent struct {
	Type    string                `json:"type"`
	Message *anthropicStartBlock  `json:"message,omitempty"`
	Index   *int                  `json:"index,omitempty"`
	Content *anthropicBlock       `json:"content_block,omitempty"`
	Delta   *anthropicDelta       `json:"delta,omitempty"`
	Usage   *anthropicUsage       `json:"usage,omitempty"`
	Error   *anthropicStreamError `json:"error,omitempty"`
}

// anthropicStartBlock 是 message_start 事件的负载。
type anthropicStartBlock struct {
	Model string          `json:"model"`
	Usage *anthropicUsage `json:"usage,omitempty"`
}

// anthropicDelta 是 content_block_delta 事件里的增量负载。
// Anthropic 把 message_delta 事件的终止 stop_reason 复用了同一字段名
// （end_turn / max_tokens / tool_use），所以这里也一并解码。
type anthropicDelta struct {
	Type        string `json:"type"` // 取值：text_delta | input_json_delta | thinking_delta
	Text        string `json:"text,omitempty"`
	PartialJSON string `json:"partial_json,omitempty"`
	Thinking    string `json:"thinking,omitempty"`
	StopReason  string `json:"stop_reason,omitempty"`
}

// anthropicStreamError 是 error 事件的负载。
type anthropicStreamError struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

// streamBlock 跟踪流式过程中单个内容块的部分状态。
type streamBlock struct {
	index       int
	typ         string
	id          string
	name        string
	text        strings.Builder
	input       strings.Builder
	initialArgs map[string]any
}

// scanAnthropicSSE 增量解码 Anthropic 的内容块流。文本
// 增量以 text_delta 事件转发（片段由
// text_start/text_end 包起来）；tool_use 的 input_json_delta 片段以
// tool_call_delta 事件转发（片段由 tool_call_start/tool_call_end 包起来），
// 并按块累积，以便块结束时重新拼装成结构化的 Args map。
// 流必须以 message_stop 事件终止，
// 否则返回 ErrStream。upstream 只用于给流内 error 事件标注「打到哪」。
func scanAnthropicSSE(upstream string, r io.Reader, onEvent func(StreamEvent)) (Response, error) {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
	agg := Response{}
	blocks := map[int]*streamBlock{}
	var order []int
	var text strings.Builder
	sawMessageStop := false

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
		if data == "" {
			continue
		}
		var evt anthropicStreamEvent
		if err := json.Unmarshal([]byte(data), &evt); err != nil {
			return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
		}
		if err := applyAnthropicEvent(&agg, blocks, &order, &text, evt, onEvent, upstream); err != nil {
			return Response{}, err
		}
		if evt.Type == "message_stop" {
			sawMessageStop = true
		}
	}
	if err := scanner.Err(); err != nil {
		return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
	}
	if !sawMessageStop {
		return Response{}, ErrStream
	}
	finalizeAnthropicStream(&agg, blocks, order, text.String())
	return agg, nil
}

// applyAnthropicEvent 把单个流事件折叠进正在运行的聚合状态。upstream 只用于给
// error 事件标注「打到哪」，不参与任何成功路径。
func applyAnthropicEvent(
	agg *Response,
	blocks map[int]*streamBlock,
	order *[]int,
	text *strings.Builder,
	evt anthropicStreamEvent,
	onEvent func(StreamEvent),
	upstream string,
) error {
	switch evt.Type {
	case "message_start":
		if evt.Message != nil {
			agg.Model = evt.Message.Model
			// 输入侧用量（含缓存读/写两档）在 message_start 就齐了，
			// 因此在这里一次折进去 —— 它与 message_delta 的输出侧用量
			// 合起来构成这次调用的完整成本面。
			applyAnthropicUsage(agg, evt.Message.Usage)
		}
	case "content_block_start":
		if evt.Index == nil || evt.Content == nil {
			return nil
		}
		b := &streamBlock{
			index:       *evt.Index,
			typ:         evt.Content.Type,
			id:          evt.Content.ID,
			name:        evt.Content.Name,
			initialArgs: evt.Content.Input,
		}
		switch evt.Content.Type {
		case "text":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextStart, Index: *evt.Index})
			}
			if evt.Content.Text != "" {
				b.text.WriteString(evt.Content.Text)
				text.WriteString(evt.Content.Text)
				if onEvent != nil {
					onEvent(StreamEvent{Kind: StreamKindTextDelta, Delta: evt.Content.Text, Index: *evt.Index})
				}
			}
		case "tool_use":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallStart, Index: *evt.Index, ID: evt.Content.ID})
			}
		case "thinking":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindThinkingStart, Index: *evt.Index})
			}
			if evt.Content.Text != "" && onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindThinkingDelta, Delta: evt.Content.Text, Index: *evt.Index})
			}
		}
		if _, ok := blocks[*evt.Index]; !ok {
			*order = append(*order, *evt.Index)
		}
		blocks[*evt.Index] = b
	case "content_block_delta":
		if evt.Index == nil || evt.Delta == nil {
			return nil
		}
		b := blocks[*evt.Index]
		if b == nil {
			return nil
		}
		switch evt.Delta.Type {
		case "text_delta":
			b.text.WriteString(evt.Delta.Text)
			text.WriteString(evt.Delta.Text)
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextDelta, Delta: evt.Delta.Text, Index: *evt.Index})
			}
		case "input_json_delta":
			b.input.WriteString(evt.Delta.PartialJSON)
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallDelta, Delta: evt.Delta.PartialJSON, Index: *evt.Index, ID: b.id})
			}
		case "thinking_delta":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindThinkingDelta, Delta: evt.Delta.Thinking, Index: *evt.Index})
			}
		}
	case "content_block_stop":
		if evt.Index == nil {
			return nil
		}
		b := blocks[*evt.Index]
		if b == nil {
			return nil
		}
		if b.typ == "text" {
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextEnd, Index: *evt.Index})
			}
		} else if b.typ == "tool_use" {
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallEnd, Index: *evt.Index, ID: b.id})
			}
		} else if b.typ == "thinking" {
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindThinkingEnd, Index: *evt.Index})
			}
		}
	case "message_delta":
		if evt.Delta != nil {
			agg.StopReason = evt.Delta.StopReason
		}
		if evt.Usage != nil {
			agg.CompletionTokens = evt.Usage.OutputTokens
		}
	case "error":
		if evt.Error != nil {
			return streamEventError(upstream, evt.Error)
		}
	}
	return nil
}

// finalizeAnthropicStream 由累积的内容块拼装出聚合后的 Response：
// 拼接的文本，加上按块导出的 tool_use ToolCall。
func finalizeAnthropicStream(agg *Response, blocks map[int]*streamBlock, order []int, content string) {
	agg.Content = content
	agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: content}}}
	agg.TotalTokens = agg.PromptTokens + agg.CompletionTokens
	for _, idx := range order {
		b := blocks[idx]
		if b == nil || b.typ != "tool_use" {
			continue
		}
		args := b.initialArgs
		if b.input.Len() > 0 {
			decoded := map[string]any{}
			if err := json.Unmarshal([]byte(b.input.String()), &decoded); err == nil {
				args = decoded
			} else {
				args = map[string]any{}
			}
		}
		agg.ToolCalls = append(agg.ToolCalls, ToolCall{Name: b.name, ID: b.id, Args: args})
	}
}