Agent-System/internal/agentx/llm/sse.go

package llm

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

type sseChoice struct {
	Delta struct {
		Content   string `json:"content"`
		Reasoning string `json:"reasoning_content"`
		ToolCalls []struct {
			Index    int    `json:"index"`
			ID       string `json:"id"`
			Function struct {
				Name      string `json:"name"`
				Arguments string `json:"arguments"`
			} `json:"function"`
		} `json:"tool_calls"`
	} `json:"delta"`
	FinishReason string `json:"finish_reason"`
}

type sseChunk struct {
	Choices []sseChoice `json:"choices"`
	Usage   *Usage      `json:"usage"`
	Error   *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// ParseSSE reads an OpenAI-compatible server-sent event body and emits
// deltas plus one terminal message_end or error event.
func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
	out := make(chan Event)
	go func() {
		defer close(out)
		emitParsed(ctx, r, out)
	}()
	return out
}

func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
	var content strings.Builder
	calls := map[int]*ToolCall{}
	order := []int{}
	var usage Usage
	finished := ""
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	for sc.Scan() {
		if ctx.Err() != nil {
			emitAbort(out, content.String(), ctx.Err())
			return
		}
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, ":") {
			continue
		}
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if data == "[DONE]" {
			break
		}
		var chunk sseChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			emitError(out, content.String(), false, "malformed stream chunk: "+err.Error())
			return
		}
		if chunk.Error != nil && chunk.Error.Message != "" {
			emitError(out, content.String(), false, chunk.Error.Message)
			return
		}
		if chunk.Usage != nil {
			usage = *chunk.Usage
		}
		if len(chunk.Choices) == 0 {
			continue
		}
		choice := chunk.Choices[0]
		if choice.FinishReason != "" {
			finished = choice.FinishReason
		}
		if choice.Delta.Content != "" {
			content.WriteString(choice.Delta.Content)
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{Kind: KindTextDelta, Text: choice.Delta.Content}:
			}
		}
		// reasoning_content（国产 OpenAI 兼容模型的思考链）与正文分流：
		// 只作为 reasoning_delta 事件外发，不计入 Message.Content（与 runner
		// stripThinkBlocks"正文无 think"语义对齐）。
		if choice.Delta.Reasoning != "" {
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{Kind: KindReasoningDelta, ReasoningDelta: choice.Delta.Reasoning}:
			}
		}
		for _, tc := range choice.Delta.ToolCalls {
			cur, ok := calls[tc.Index]
			if !ok {
				cur = &ToolCall{}
				calls[tc.Index] = cur
				order = append(order, tc.Index)
			}
			if tc.ID != "" {
				cur.ID = tc.ID
			}
			if tc.Function.Name != "" {
				cur.Name += tc.Function.Name
			}
			if tc.Function.Arguments != "" {
				cur.Arguments += tc.Function.Arguments
			}
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{
				Kind:           KindToolCallDelta,
				ToolIndex:      tc.Index,
				ToolID:         tc.ID,
				ToolName:       tc.Function.Name,
				ArgumentsDelta: tc.Function.Arguments,
			}:
			}
		}
	}
	if err := sc.Err(); err != nil {
		if ctx.Err() != nil {
			emitAbort(out, content.String(), ctx.Err())
			return
		}
		emitError(out, content.String(), true, err.Error())
		return
	}
	msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
	for _, idx := range order {
		msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
	}
	if msg.StopReason == StopToolCalls && len(msg.ToolCalls) == 0 {
		msg.StopReason = StopEnd
	}
	select {
	case <-ctx.Done():
		emitAbort(out, content.String(), ctx.Err())
	case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage}:
	}
}

func mapFinish(reason string, hasTools bool) StopReason {
	switch reason {
	case "tool_calls", "function_call":
		return StopToolCalls
	case "content_filter", "error":
		return StopError
	case "stop", "length", "":
		if hasTools {
			return StopToolCalls
		}
		return StopEnd
	default:
		if hasTools {
			return StopToolCalls
		}
		return StopEnd
	}
}

func emitError(out chan<- Event, partial string, retryable bool, message string) {
	out <- Event{
		Kind:         KindError,
		Retryable:    retryable,
		ErrorMessage: message,
		Err:          &StreamError{Message: message, Retryable: retryable},
		Message: Message{
			Role:       RoleAssistant,
			Content:    partial,
			StopReason: StopError,
		},
	}
}

func emitAbort(out chan<- Event, partial string, err error) {
	msg := "aborted"
	if err != nil {
		msg = err.Error()
	}
	out <- Event{
		Kind:         KindError,
		Retryable:    false,
		ErrorMessage: msg,
		Err:          err,
		Message: Message{
			Role:       RoleAssistant,
			Content:    partial,
			StopReason: StopAborted,
		},
	}
}

// StreamError is a provider or transport failure encoded on the event stream.
type StreamError struct {
	Message   string
	Retryable bool
}

func (e *StreamError) Error() string {
	if e == nil {
		return "stream error"
	}
	return e.Message
}