ituoke/internal/agent/stream.go

package agent

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

func completeStream(ctx context.Context, c ModelConfig, key string, messages []modelMessage, onDelta func(string), onReasoning ...func(string)) (string, string, error) {
	res, err := openModelRequest(ctx, c, key, http.MethodPost, "/chat/completions", completionBody(c, messages, true))
	if err != nil {
		return "", "", err
	}
	defer res.Body.Close()
	if !strings.HasPrefix(strings.ToLower(res.Header.Get("Content-Type")), "text/event-stream") {
		return "", "", fmt.Errorf("模型服务没有返回流式响应，请确认所选接口支持 stream=true")
	}
	step, err := readToolStream(res.Body, key, onDelta, onReasoning...)
	text, reason := step.Text, step.Reason
	if ctx.Err() != nil {
		return text, reason, ctx.Err()
	}
	return text, reason, err
}

// SSE data events may span TCP reads or multiple data: lines. Completion markers
// distinguish a finished answer from a connection that silently ended mid-answer.
func readCompletionStream(reader io.Reader, key string, onDelta func(string)) (string, string, error) {
	result, err := readToolStream(reader, key, onDelta)
	return result.Text, result.Reason, err
}

type completionStep struct {
	Text      string
	Reason    string
	Reasoning string
	Calls     []modelToolCall
}

func readToolStream(reader io.Reader, key string, onDelta func(string), onReasoning ...func(string)) (completionStep, error) {
	scanner := bufio.NewScanner(io.LimitReader(reader, 8*1024*1024+1))
	scanner.Buffer(make([]byte, 4096), 512*1024)
	var text strings.Builder
	var reasoning strings.Builder
	var data []string
	finishReason := ""
	done := false
	calls := map[int]*modelToolCall{}
	result := func() completionStep {
		step := completionStep{Text: text.String(), Reason: finishReason, Reasoning: reasoning.String()}
		for i := 0; i < len(calls); i++ {
			if calls[i] != nil {
				step.Calls = append(step.Calls, *calls[i])
			}
		}
		return step
	}
	dispatch := func() error {
		payload := strings.Join(data, "\n")
		data = nil
		if payload == "" {
			return nil
		}
		if payload == "[DONE]" {
			done = true
			return nil
		}
		var chunk struct {
			Error   json.RawMessage `json:"error"`
			Choices []struct {
				Index int `json:"index"`
				Delta struct {
					ReasoningContent string `json:"reasoning_content"`
					Reasoning        string `json:"reasoning"`
					Content          string `json:"content"`
					Refusal          string `json:"refusal"`
					ToolCalls        []struct {
						Index    int    `json:"index"`
						ID       string `json:"id"`
						Type     string `json:"type"`
						Function struct {
							Name      string `json:"name"`
							Arguments string `json:"arguments"`
						} `json:"function"`
					} `json:"tool_calls"`
				} `json:"delta"`
				FinishReason string `json:"finish_reason"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			return fmt.Errorf("模型流式响应格式无效，已保留收到的内容")
		}
		if len(chunk.Error) > 0 && string(chunk.Error) != "null" {
			return modelHTTPError(200, []byte(payload), key)
		}
		for _, choice := range chunk.Choices {
			if choice.Index != 0 {
				continue
			}
			for _, delta := range choice.Delta.ToolCalls {
				if delta.Index < 0 {
					return fmt.Errorf("模型工具调用索引无效")
				}
				call := calls[delta.Index]
				if call == nil {
					call = &modelToolCall{Type: "function"}
					calls[delta.Index] = call
				}
				if delta.ID != "" {
					call.ID = delta.ID
				}
				if delta.Type != "" {
					call.Type = delta.Type
				}
				call.Function.Name += delta.Function.Name
				call.Function.Arguments += delta.Function.Arguments
				if len(call.Function.Name) > 128 || len(call.Function.Arguments) > 64*1024 {
					return fmt.Errorf("工具调用参数过大")
				}
			}
			delta := choice.Delta.Content
			if delta == "" {
				delta = choice.Delta.Refusal
			}
			thought := choice.Delta.ReasoningContent
			if thought == "" {
				thought = choice.Delta.Reasoning
			}
			if text.Len()+reasoning.Len()+len(delta)+len(thought) > 4*1024*1024 {
				return fmt.Errorf("模型输出过大，已停止接收")
			}
			if thought != "" {
				reasoning.WriteString(thought)
				for _, notify := range onReasoning {
					if notify != nil {
						notify(thought)
					}
				}
			}
			if delta != "" {
				text.WriteString(delta)
				if onDelta != nil {
					onDelta(delta)
				}
			}
			if choice.FinishReason != "" {
				finishReason = choice.FinishReason
			}
		}
		return nil
	}
	for scanner.Scan() {
		line := strings.TrimPrefix(scanner.Text(), "\ufeff")
		if line == "" {
			if err := dispatch(); err != nil {
				return result(), err
			}
			if done {
				break
			}
		} else if value, ok := strings.CutPrefix(line, "data:"); ok {
			data = append(data, strings.TrimPrefix(value, " "))
		}
	}
	if err := scanner.Err(); err != nil {
		return result(), fmt.Errorf("模型连接中断或超时，已保留收到的内容")
	}
	if len(data) > 0 {
		if err := dispatch(); err != nil {
			return result(), err
		}
	}
	if !done && finishReason == "" {
		return result(), fmt.Errorf("模型连接提前结束，回答可能不完整")
	}
	if finishReason == "content_filter" {
		return result(), fmt.Errorf("模型输出被服务商过滤，已保留收到的内容")
	}
	if len(calls) > 0 {
		for i := 0; i < len(calls); i++ {
			call := calls[i]
			if call == nil || call.ID == "" || call.Type != "function" || call.Function.Name == "" || !json.Valid([]byte(call.Function.Arguments)) {
				return result(), fmt.Errorf("模型工具调用不完整，未执行")
			}
		}
		if finishReason != "tool_calls" && finishReason != "stop" {
			return result(), fmt.Errorf("工具调用未正常结束，未执行")
		}
	}
	if text.Len() == 0 && len(calls) == 0 {
		return result(), fmt.Errorf("模型没有返回文字，请检查模型能力或输出上限")
	}
	return result(), nil
}