mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 08:57:57 +00:00
feat(gateway): pass through exact upstream token usage in streams
Streaming responses now carry the upstream's real token usage instead of gateway estimates: - UnifiedChunk gains an optional Usage field; adapters (opencode, openai) extract usage from upstream stream chunks (including the final chunk with empty choices) and pass it through. - standardSSEChunk preserves usage for passthrough adapters. - Gateway emits the exact usage in the final stream chunk when available, falling back to estimates only when the upstream provided none. Non-streaming usage was already fixed to emit OpenAI-standard keys.
This commit is contained in:
@ -901,20 +901,52 @@ func standardSSEChunk(data string) string {
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
UpstreamUsage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &raw); err != nil {
|
||||
return ""
|
||||
}
|
||||
if len(raw.Choices) == 0 {
|
||||
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
|
||||
return ""
|
||||
}
|
||||
var usage *types.TokenUsage
|
||||
pu := raw.UpstreamUsage
|
||||
if pu.Total > 0 || pu.TotalTokens > 0 {
|
||||
usage = &types.TokenUsage{
|
||||
Prompt: pickFirst(pu.PromptTokens, pu.Prompt),
|
||||
Completion: pickFirst(pu.CompletionTokens, pu.Completion),
|
||||
Total: pickFirst(pu.TotalTokens, pu.Total),
|
||||
}
|
||||
}
|
||||
out, _ := json.Marshal(types.UnifiedChunk{
|
||||
Content: raw.Choices[0].Delta.Content,
|
||||
Done: raw.Choices[0].FinishReason != nil,
|
||||
Content: func() string {
|
||||
if len(raw.Choices) > 0 {
|
||||
return raw.Choices[0].Delta.Content
|
||||
}
|
||||
return ""
|
||||
}(),
|
||||
Done: len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil,
|
||||
Usage: usage,
|
||||
})
|
||||
return string(out)
|
||||
}
|
||||
|
||||
// pickFirst returns a if non-zero, else b (for usage keys that may appear in
|
||||
// either standard *_tokens or legacy short form).
|
||||
func pickFirst(a, b int) int {
|
||||
if a != 0 {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
|
||||
Reference in New Issue
Block a user