536 lines
15 KiB
Go
536 lines
15 KiB
Go
package llm
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"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 *usageWire `json:"usage"`
|
|
Error *struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
type usageWire struct {
|
|
Usage
|
|
Details struct {
|
|
Cached int `json:"cached_tokens"`
|
|
} `json:"prompt_tokens_details"`
|
|
}
|
|
|
|
const (
|
|
maxObservedToolFragments = 10000
|
|
maxObservedToolArgumentBytes = 4*1024*1024 + 1
|
|
maxObservedTailEvents = 10000
|
|
)
|
|
|
|
// sseFacts are structural stream facts the route contract needs. They are
|
|
// written by the parser goroutine before it sends the terminal event.
|
|
type sseFacts struct {
|
|
// finishChunks counts chunks with a non-empty finish_reason.
|
|
finishChunks int
|
|
finishReason string
|
|
// multiChoice is set when any chunk carried more than one choice.
|
|
multiChoice bool
|
|
// fragmentAfterFinish is set when a tool-call fragment arrived in a chunk
|
|
// after the one carrying finish_reason.
|
|
fragmentAfterFinish bool
|
|
}
|
|
|
|
// sseHooks are internal parser options. The zero value is ParseSSE.
|
|
type sseHooks struct {
|
|
facts *sseFacts
|
|
// decideEOF, when set, receives a tool stream whose terminal would be
|
|
// "tool call stream ended before [DONE]": success is the message_end the
|
|
// stream would end with had [DONE] arrived, failure is the event emitted
|
|
// today. It returns the event to send.
|
|
decideEOF func(success, failure Event) Event
|
|
}
|
|
|
|
// 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 {
|
|
return parseSSE(ctx, r, sseHooks{})
|
|
}
|
|
|
|
func parseSSE(ctx context.Context, r io.Reader, hooks sseHooks) <-chan Event {
|
|
out := make(chan Event)
|
|
go func() {
|
|
defer close(out)
|
|
emitParsed(ctx, r, out, hooks)
|
|
}()
|
|
return out
|
|
}
|
|
|
|
func emitParsed(ctx context.Context, r io.Reader, out chan<- Event, hooks sseHooks) {
|
|
facts := hooks.facts
|
|
if facts == nil {
|
|
facts = &sseFacts{}
|
|
}
|
|
var content strings.Builder
|
|
var reasoning strings.Builder
|
|
calls := map[int]*ToolCall{}
|
|
order := []int{}
|
|
var usage Usage
|
|
finished := ""
|
|
sawDone := false
|
|
sawUsage := false
|
|
readOutcome := "clean_eof"
|
|
lastFrameJSONComplete := false
|
|
toolFragments := 0
|
|
toolArgumentBytes := 0
|
|
sc := bufio.NewScanner(r)
|
|
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
|
reader := &sseEventReader{sc: sc}
|
|
observation := func() *TerminalObservation {
|
|
tailKind := reader.tailKind
|
|
if tailKind == "" {
|
|
tailKind = "unknown"
|
|
}
|
|
argumentsValid := len(order) > 0
|
|
for _, idx := range order {
|
|
var args map[string]json.RawMessage
|
|
if err := json.Unmarshal([]byte(calls[idx].Arguments), &args); err != nil || args == nil {
|
|
argumentsValid = false
|
|
break
|
|
}
|
|
}
|
|
return &TerminalObservation{
|
|
ReadOutcome: readOutcome, LastFrameJSONComplete: lastFrameJSONComplete,
|
|
UsagePresent: sawUsage, ToolFragments: toolFragments,
|
|
ToolArgumentBytes: toolArgumentBytes, ToolArgumentsValid: argumentsValid,
|
|
TailEventKind: tailKind, TailEventCount: reader.tailCount,
|
|
}
|
|
}
|
|
buildTerminalFailure := func(why string) Event {
|
|
message := Message{Role: RoleAssistant, Content: content.String(), StopReason: StopError}
|
|
if reasoning.Len() > 0 {
|
|
message.Thinking = []ThinkingBlock{{Thinking: reasoning.String()}}
|
|
}
|
|
failed := terminalStreamFailure(message, usage, finished, sawDone, why)
|
|
failed.TerminalObservation = observation()
|
|
return failed
|
|
}
|
|
sendTerminal := func(ev Event) {
|
|
select {
|
|
case <-ctx.Done():
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
case out <- ev:
|
|
}
|
|
}
|
|
sendTerminalFailure := func(why string) {
|
|
sendTerminal(buildTerminalFailure(why))
|
|
}
|
|
buildMessageEnd := func() Event {
|
|
msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
|
|
if reasoning.Len() > 0 {
|
|
msg.Thinking = []ThinkingBlock{{Thinking: reasoning.String()}}
|
|
}
|
|
for _, idx := range order {
|
|
msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
|
|
}
|
|
if msg.StopReason == StopToolCalls && len(msg.ToolCalls) == 0 {
|
|
msg.StopReason = StopEnd
|
|
}
|
|
return Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished, StreamComplete: sawDone, TerminalObservation: observation()}
|
|
}
|
|
for {
|
|
data, ok := reader.next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if ctx.Err() != nil {
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
return
|
|
}
|
|
if data == "[DONE]" {
|
|
reader.noteTail("done_literal")
|
|
sawDone = true
|
|
readOutcome = "done_marker"
|
|
break
|
|
}
|
|
var chunk sseChunk
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
reader.noteTail(badFrameKind(data, err))
|
|
readOutcome = "invalid_frame"
|
|
lastFrameJSONComplete = false
|
|
if sawUsage {
|
|
sendTerminalFailure("malformed stream chunk after reported usage")
|
|
return
|
|
}
|
|
emitError(out, content.String(), false, "malformed stream chunk: "+err.Error(), observation())
|
|
return
|
|
}
|
|
reader.noteTail("json_frame")
|
|
lastFrameJSONComplete = true
|
|
if chunk.Error != nil && chunk.Error.Message != "" {
|
|
readOutcome = "provider_error"
|
|
msg := chunk.Error.Message
|
|
if sawUsage {
|
|
sendTerminalFailure("provider error after reported usage: " + msg)
|
|
return
|
|
}
|
|
retryable := CapacityLimited(msg) && !ContextOverflow(msg)
|
|
emitError(out, content.String(), retryable, msg, observation())
|
|
return
|
|
}
|
|
if chunk.Usage != nil {
|
|
usage = chunk.Usage.Usage
|
|
sawUsage = true
|
|
if chunk.Usage.Details.Cached > 0 {
|
|
usage.CacheReadTokens = chunk.Usage.Details.Cached
|
|
}
|
|
}
|
|
if len(chunk.Choices) > 1 {
|
|
facts.multiChoice = true
|
|
}
|
|
finishedBefore := facts.finishChunks > 0
|
|
for _, c := range chunk.Choices {
|
|
if c.FinishReason != "" {
|
|
if facts.finishChunks < maxObservedTailEvents {
|
|
facts.finishChunks++
|
|
}
|
|
break
|
|
}
|
|
}
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
choice := chunk.Choices[0]
|
|
if choice.FinishReason != "" {
|
|
finished = choice.FinishReason
|
|
facts.finishReason = finished
|
|
}
|
|
if finishedBefore && len(choice.Delta.ToolCalls) > 0 {
|
|
facts.fragmentAfterFinish = true
|
|
}
|
|
if choice.Delta.Content != "" {
|
|
content.WriteString(choice.Delta.Content)
|
|
select {
|
|
case <-ctx.Done():
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
return
|
|
case out <- Event{Kind: KindTextDelta, Text: choice.Delta.Content}:
|
|
}
|
|
}
|
|
// reasoning_content（国产 OpenAI 兼容模型的思考链）与正文分流：
|
|
// 只作为 reasoning_delta 事件外发，不计入 Message.Content（与 runner
|
|
// stripThinkBlocks"正文无 think"语义对齐）。同时累计进 reasoning——
|
|
// 交错思考模型（MiniMax/DeepSeek R1 等）要求历史 assistant 原样带回
|
|
// 思考内容，丢弃会降低后续轮质量；捕获进 Message.Thinking 随会话
|
|
// 转录持久化，请求序列化侧按需回放。
|
|
if choice.Delta.Reasoning != "" {
|
|
reasoning.WriteString(choice.Delta.Reasoning)
|
|
select {
|
|
case <-ctx.Done():
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
return
|
|
case out <- Event{Kind: KindReasoningDelta, ReasoningDelta: choice.Delta.Reasoning}:
|
|
}
|
|
}
|
|
for _, tc := range choice.Delta.ToolCalls {
|
|
if toolFragments < maxObservedToolFragments {
|
|
toolFragments++
|
|
}
|
|
if n := len(tc.Function.Arguments); n > 0 && toolArgumentBytes < maxObservedToolArgumentBytes {
|
|
if n > maxObservedToolArgumentBytes-toolArgumentBytes {
|
|
toolArgumentBytes = maxObservedToolArgumentBytes
|
|
} else {
|
|
toolArgumentBytes += n
|
|
}
|
|
}
|
|
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():
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
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 {
|
|
readOutcome = "read_error"
|
|
if ctx.Err() != nil {
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
return
|
|
}
|
|
if sawUsage {
|
|
sendTerminalFailure("stream read failed after reported usage")
|
|
return
|
|
}
|
|
emitError(out, content.String(), true, err.Error(), observation())
|
|
return
|
|
}
|
|
// 取消必须压过正常结束。select 在两个分支都就绪时随机挑选，
|
|
// 竞态下会把 message_end 发给已经取消的调用方。
|
|
if ctx.Err() != nil {
|
|
readOutcome = "cancelled"
|
|
emitAbort(out, content.String(), ctx.Err(), observation())
|
|
return
|
|
}
|
|
// A finish reason alone does not prove the tool arguments arrived intact.
|
|
// Preserve returned usage when rejecting a partial terminal, so the host
|
|
// can charge this call and stop rather than treating it as a free retry.
|
|
if (finished == "tool_calls" || finished == "function_call") && len(order) == 0 {
|
|
sendTerminalFailure("tool call finish reason without a tool call")
|
|
return
|
|
}
|
|
if len(order) > 0 {
|
|
problem := ""
|
|
switch {
|
|
case !sawDone:
|
|
problem = "tool call stream ended before [DONE]"
|
|
case !toolFinishExecutable(finished, sawDone):
|
|
problem = "tool call ended without a supported finish reason"
|
|
default:
|
|
for _, idx := range order {
|
|
call := calls[idx]
|
|
var args map[string]json.RawMessage
|
|
if call.ID == "" || call.Name == "" {
|
|
problem = "tool call identity is incomplete"
|
|
break
|
|
}
|
|
if err := json.Unmarshal([]byte(call.Arguments), &args); err != nil || args == nil {
|
|
problem = "tool call arguments are not a JSON object"
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if problem != "" {
|
|
if !sawDone && hooks.decideEOF != nil {
|
|
sendTerminal(hooks.decideEOF(buildMessageEnd(), buildTerminalFailure(problem)))
|
|
return
|
|
}
|
|
sendTerminalFailure(problem)
|
|
return
|
|
}
|
|
}
|
|
sendTerminal(buildMessageEnd())
|
|
}
|
|
|
|
func terminalStreamFailure(message Message, usage Usage, reason string, complete bool, why string) Event {
|
|
err := &StreamError{Message: why, Retryable: false}
|
|
return Event{Kind: KindMessageEnd, Message: message, Usage: usage,
|
|
FinishReason: reason, StreamComplete: complete, Err: err,
|
|
ErrorMessage: why, Retryable: false}
|
|
}
|
|
|
|
// sseEventReader 按空行切事件，并把同一事件里的多行 data: 用换行拼回。
|
|
// 已经是完整 JSON 的 data 行单独成事件，这样单行流不用空行也能解析。
|
|
// tailKind/tailCount 只记最后完成的结构事件类别与饱和计数，不存帧内容。
|
|
type sseEventReader struct {
|
|
sc *bufio.Scanner
|
|
carry string
|
|
tailKind string
|
|
tailCount int
|
|
}
|
|
|
|
func (r *sseEventReader) noteTail(kind string) {
|
|
if r.tailCount < maxObservedTailEvents {
|
|
r.tailCount++
|
|
}
|
|
r.tailKind = kind
|
|
}
|
|
|
|
// badFrameKind classifies a data event that did not decode as a chunk using
|
|
// only the decoder's error category and position, never its text or bytes.
|
|
func badFrameKind(data string, err error) string {
|
|
if strings.TrimSpace(data) == "" {
|
|
return "empty_data"
|
|
}
|
|
var syntax *json.SyntaxError
|
|
if !errors.As(err, &syntax) {
|
|
return "unknown"
|
|
}
|
|
if syntax.Offset == int64(len(data)) {
|
|
var probe json.RawMessage
|
|
if errors.Is(json.NewDecoder(strings.NewReader(data)).Decode(&probe), io.ErrUnexpectedEOF) {
|
|
return "partial_frame"
|
|
}
|
|
}
|
|
return "unknown_non_json"
|
|
}
|
|
|
|
func (r *sseEventReader) next() (string, bool) {
|
|
var parts []string
|
|
nextPiece := func() (string, bool) {
|
|
if r.carry != "" {
|
|
piece := r.carry
|
|
r.carry = ""
|
|
return piece, true
|
|
}
|
|
if !r.sc.Scan() {
|
|
return "", false
|
|
}
|
|
return r.sc.Text(), true
|
|
}
|
|
for {
|
|
line, ok := nextPiece()
|
|
if !ok {
|
|
break
|
|
}
|
|
if strings.TrimSpace(line) == "" {
|
|
if len(parts) == 0 {
|
|
continue
|
|
}
|
|
return strings.Join(parts, "\n"), true
|
|
}
|
|
trimmed := strings.TrimSpace(line)
|
|
if strings.HasPrefix(trimmed, ":") {
|
|
r.noteTail("comment")
|
|
continue
|
|
}
|
|
if !strings.HasPrefix(trimmed, "data:") {
|
|
continue
|
|
}
|
|
piece := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
|
if len(parts) > 0 && json.Valid([]byte(strings.Join(parts, "\n"))) {
|
|
r.carry = line
|
|
return strings.Join(parts, "\n"), true
|
|
}
|
|
parts = append(parts, piece)
|
|
}
|
|
if len(parts) > 0 {
|
|
return strings.Join(parts, "\n"), true
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
// toolFinishExecutable 只在供应商明确以工具调用或正常 stop 结束时执行。
|
|
// length、缺失的结束原因，以及没有 [DONE] 的断开，都视为未完成。
|
|
func toolFinishExecutable(reason string, sawDone bool) bool {
|
|
if !sawDone {
|
|
return false
|
|
}
|
|
switch reason {
|
|
case "tool_calls", "function_call":
|
|
return true
|
|
case "stop":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func mapFinish(reason string, hasTools bool) StopReason {
|
|
switch reason {
|
|
case "tool_calls", "function_call":
|
|
return StopToolCalls
|
|
case "content_filter", "error":
|
|
return StopError
|
|
case "length":
|
|
return StopEnd
|
|
case "stop", "":
|
|
if hasTools {
|
|
return StopToolCalls
|
|
}
|
|
return StopEnd
|
|
default:
|
|
if hasTools {
|
|
return StopToolCalls
|
|
}
|
|
return StopEnd
|
|
}
|
|
}
|
|
|
|
func emitError(out chan<- Event, partial string, retryable bool, message string, observed ...*TerminalObservation) {
|
|
event := Event{
|
|
Kind: KindError,
|
|
Retryable: retryable,
|
|
ErrorMessage: message,
|
|
Err: &StreamError{Message: message, Retryable: retryable},
|
|
Message: Message{
|
|
Role: RoleAssistant,
|
|
Content: partial,
|
|
StopReason: StopError,
|
|
},
|
|
}
|
|
if len(observed) > 0 {
|
|
event.TerminalObservation = observed[0]
|
|
}
|
|
out <- event
|
|
}
|
|
|
|
func emitAbort(out chan<- Event, partial string, err error, observed ...*TerminalObservation) {
|
|
msg := "aborted"
|
|
if err != nil {
|
|
msg = err.Error()
|
|
}
|
|
event := Event{
|
|
Kind: KindError,
|
|
Retryable: false,
|
|
ErrorMessage: msg,
|
|
Err: err,
|
|
Message: Message{
|
|
Role: RoleAssistant,
|
|
Content: partial,
|
|
StopReason: StopAborted,
|
|
},
|
|
}
|
|
if len(observed) > 0 {
|
|
event.TerminalObservation = observed[0]
|
|
}
|
|
out <- event
|
|
}
|
|
|
|
// 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
|
|
}
|