internal/gateway/generation_stream.go
Rama main · réplica de lectura del repositorio original.
package gateway
import (
"bytes"
"encoding/json"
"slices"
"strings"
"github.com/enterpilot/gomodel/internal/core"
"github.com/enterpilot/gomodel/internal/usage"
)
// generationStreamObserver rebuilds a streamed response from the canonical
// SSE events and finishes the request's generation trace when the stream
// closes, so the provider-call span ends with the stream and carries its
// usage. Message content is accumulated only when the trace captures it.
type generationStreamObserver struct {
trace *core.GenerationTrace
request any
responses bool
content bool
closed bool
chat core.ChatResponse
choices map[int]*chatChoiceState
response *core.ResponsesResponse
}
type chatChoiceState struct {
role string
text cappedText
toolCalls map[int]*streamToolCall
finishReason string
}
type streamToolCall struct {
id string
name string
arguments cappedText
}
// cappedText accumulates streamed text up to one byte past
// core.GenerationContentLimit, so memory stays bounded however long the
// stream runs while the exporter still sees the text as too long and marks
// it truncated.
type cappedText struct {
strings.Builder
}
func (t *cappedText) append(text string) {
if room := core.GenerationContentLimit + 1 - t.Len(); room > 0 {
t.WriteString(text[:min(len(text), room)])
}
}
func newChatGenerationObserver(trace *core.GenerationTrace, req *core.ChatRequest) *generationStreamObserver {
if trace == nil {
return nil
}
return &generationStreamObserver{trace: trace, request: req, content: trace.CapturesContent()}
}
func newResponsesGenerationObserver(trace *core.GenerationTrace, req *core.ResponsesRequest) *generationStreamObserver {
if trace == nil {
return nil
}
return &generationStreamObserver{trace: trace, request: req, responses: true, content: trace.CapturesContent()}
}
var (
finishReasonKey = []byte(`"finish_reason"`)
responsesTerminalLiterals = [][]byte{[]byte(`"response.completed"`), []byte(`"response.incomplete"`), []byte(`"response.failed"`)}
)
// WantsJSONEvent keeps per-chunk decoding off unless content is captured: a
// chat stream is read for its usage and finish chunks, a Responses stream
// for its terminal event, which carries the whole response.
func (o *generationStreamObserver) WantsJSONEvent(raw []byte) bool {
if o.responses {
return slices.ContainsFunc(responsesTerminalLiterals, func(literal []byte) bool {
return bytes.Contains(raw, literal)
})
}
return o.content || usage.HasUsageObject(raw) || hasStringMember(raw, finishReasonKey)
}
// hasStringMember reports whether raw has key with a string value, so chunks
// carrying "finish_reason":null are skipped. Passthrough streams keep the
// upstream's encoding, which may put spaces around the colon.
func hasStringMember(raw, key []byte) bool {
for rest := raw; ; {
i := bytes.Index(rest, key)
if i < 0 {
return false
}
rest = bytes.TrimLeft(rest[i+len(key):], " \t\r\n")
if len(rest) > 0 && rest[0] == ':' {
if value := bytes.TrimLeft(rest[1:], " \t\r\n"); len(value) > 0 && value[0] == '"' {
return true
}
}
}
}
func (o *generationStreamObserver) OnJSONEvent(event map[string]any) {
if o.responses {
o.observeResponsesEvent(event)
return
}
o.observeChatChunk(event)
}
func (o *generationStreamObserver) OnStreamClose() {
if o.closed {
return
}
o.closed = true
outcome := core.GenerationOutcome{Request: o.request}
if o.responses {
if o.response != nil {
outcome.Response = o.response
}
} else if o.chat.ID != "" || o.chat.Model != "" || len(o.choices) > 0 {
outcome.Response = o.chatResponse()
}
o.trace.Finish(outcome)
}
func (o *generationStreamObserver) observeResponsesEvent(event map[string]any) {
eventType, _ := event["type"].(string)
switch eventType {
case "response.completed", "response.incomplete", "response.failed":
default:
return
}
encoded, err := json.Marshal(event["response"])
if err != nil {
return
}
var resp core.ResponsesResponse
if json.Unmarshal(encoded, &resp) == nil {
o.response = &resp
}
}
func (o *generationStreamObserver) observeChatChunk(event map[string]any) {
if id, ok := event["id"].(string); ok && id != "" {
o.chat.ID = id
}
if model, ok := event["model"].(string); ok && model != "" {
o.chat.Model = model
}
if raw, ok := event["usage"].(map[string]any); ok {
o.chat.Usage.PromptTokens = intValue(raw["prompt_tokens"])
o.chat.Usage.CompletionTokens = intValue(raw["completion_tokens"])
o.chat.Usage.TotalTokens = intValue(raw["total_tokens"])
}
choices, _ := event["choices"].([]any)
for _, raw := range choices {
choice, ok := raw.(map[string]any)
if !ok {
continue
}
state := o.choice(intValue(choice["index"]))
if reason, ok := choice["finish_reason"].(string); ok && reason != "" {
state.finishReason = reason
}
if delta, ok := choice["delta"].(map[string]any); ok && o.content {
state.observeDelta(delta)
}
}
}
func (o *generationStreamObserver) choice(index int) *chatChoiceState {
if o.choices == nil {
o.choices = make(map[int]*chatChoiceState)
}
state, ok := o.choices[index]
if !ok {
state = &chatChoiceState{}
o.choices[index] = state
}
return state
}
func (s *chatChoiceState) observeDelta(delta map[string]any) {
if role, ok := delta["role"].(string); ok && role != "" {
s.role = role
}
if text, ok := delta["content"].(string); ok {
s.text.append(text)
}
calls, _ := delta["tool_calls"].([]any)
for _, raw := range calls {
call, ok := raw.(map[string]any)
if !ok {
continue
}
index := intValue(call["index"])
if s.toolCalls == nil {
s.toolCalls = make(map[int]*streamToolCall)
}
state, ok := s.toolCalls[index]
if !ok {
state = &streamToolCall{}
s.toolCalls[index] = state
}
if id, ok := call["id"].(string); ok && id != "" {
state.id = id
}
if function, ok := call["function"].(map[string]any); ok {
if name, ok := function["name"].(string); ok && name != "" {
state.name = name
}
if arguments, ok := function["arguments"].(string); ok {
state.arguments.append(arguments)
}
}
}
}
func (o *generationStreamObserver) chatResponse() *core.ChatResponse {
resp := o.chat
for _, index := range sortedKeys(o.choices) {
state := o.choices[index]
message := core.ResponseMessage{Role: state.role}
if text := state.text.String(); text != "" {
message.Content = text
}
for _, callIndex := range sortedKeys(state.toolCalls) {
call := state.toolCalls[callIndex]
message.ToolCalls = append(message.ToolCalls, core.ToolCall{
ID: call.id,
Type: "function",
Function: core.FunctionCall{Name: call.name, Arguments: call.arguments.String()},
})
}
resp.Choices = append(resp.Choices, core.Choice{Index: index, Message: message, FinishReason: state.finishReason})
}
return &resp
}
func sortedKeys[V any](m map[int]V) []int {
keys := make([]int, 0, len(m))
for key := range m {
keys = append(keys, key)
}
slices.Sort(keys)
return keys
}
// intValue reads a JSON number decoded into an interface.
func intValue(value any) int {
number, _ := value.(float64)
return int(number)
}