mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 17:07:59 +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:
@ -629,7 +629,11 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
|||||||
}) {
|
}) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
var lastUsage *types.TokenUsage
|
||||||
for ck := range chunks {
|
for ck := range chunks {
|
||||||
|
if ck.Usage != nil {
|
||||||
|
lastUsage = ck.Usage
|
||||||
|
}
|
||||||
chunk := ChatChunk{
|
chunk := ChatChunk{
|
||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||||
}
|
}
|
||||||
@ -657,16 +661,24 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
|||||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||||
})
|
})
|
||||||
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||||
if rec.Prompt+rec.Compl > 0 {
|
// Prefer the upstream's exact usage if the stream carried it; fall back to
|
||||||
usage := types.TokenUsage{
|
// the gateway's estimate otherwise.
|
||||||
|
var tut *types.TokenUsage
|
||||||
|
if lastUsage != nil {
|
||||||
|
tut = lastUsage
|
||||||
|
} else if rec.Prompt+rec.Compl > 0 {
|
||||||
|
u := types.TokenUsage{
|
||||||
Prompt: int(rec.Prompt),
|
Prompt: int(rec.Prompt),
|
||||||
Completion: int(rec.Compl),
|
Completion: int(rec.Compl),
|
||||||
Total: int(rec.Prompt + rec.Compl),
|
Total: int(rec.Prompt + rec.Compl),
|
||||||
}
|
}
|
||||||
|
tut = &u
|
||||||
|
}
|
||||||
|
if tut != nil {
|
||||||
send(ChatChunk{
|
send(ChatChunk{
|
||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||||
Choices: []ChunkChoice{},
|
Choices: []ChunkChoice{},
|
||||||
Usage: &usage,
|
Usage: tut,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
@ -787,7 +799,11 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
|||||||
}) {
|
}) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
var lastUsage *types.TokenUsage
|
||||||
for ck := range chunks {
|
for ck := range chunks {
|
||||||
|
if ck.Usage != nil {
|
||||||
|
lastUsage = ck.Usage
|
||||||
|
}
|
||||||
chunk := ChatChunk{
|
chunk := ChatChunk{
|
||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||||
}
|
}
|
||||||
@ -815,16 +831,24 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
|||||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||||
})
|
})
|
||||||
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||||
if rec.Prompt+rec.Compl > 0 {
|
// Prefer the upstream's exact usage if the stream carried it; fall back to
|
||||||
usage := types.TokenUsage{
|
// the gateway's estimate otherwise.
|
||||||
|
var tut *types.TokenUsage
|
||||||
|
if lastUsage != nil {
|
||||||
|
tut = lastUsage
|
||||||
|
} else if rec.Prompt+rec.Compl > 0 {
|
||||||
|
u := types.TokenUsage{
|
||||||
Prompt: int(rec.Prompt),
|
Prompt: int(rec.Prompt),
|
||||||
Completion: int(rec.Compl),
|
Completion: int(rec.Compl),
|
||||||
Total: int(rec.Prompt + rec.Compl),
|
Total: int(rec.Prompt + rec.Compl),
|
||||||
}
|
}
|
||||||
|
tut = &u
|
||||||
|
}
|
||||||
|
if tut != nil {
|
||||||
send(ChatChunk{
|
send(ChatChunk{
|
||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||||
Choices: []ChunkChoice{},
|
Choices: []ChunkChoice{},
|
||||||
Usage: &usage,
|
Usage: tut,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
|||||||
@ -67,7 +67,24 @@ function adapter.transform_stream_chunk(raw_chunk)
|
|||||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||||
if not ok then return "" end
|
if not ok then return "" end
|
||||||
|
|
||||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||||
|
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||||
|
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||||
|
local uses = nil
|
||||||
|
if type(chunk.usage) == "table" then
|
||||||
|
uses = {
|
||||||
|
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||||
|
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||||
|
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||||
|
}
|
||||||
|
end
|
||||||
|
|
||||||
|
if not chunk.choices or #chunk.choices == 0 then
|
||||||
|
if uses ~= nil then
|
||||||
|
return json.encode({ usage = uses, done = (chunk.usage ~= nil) })
|
||||||
|
end
|
||||||
|
return ""
|
||||||
|
end
|
||||||
local delta = chunk.choices[1].delta or {}
|
local delta = chunk.choices[1].delta or {}
|
||||||
local fr = chunk.choices[1].finish_reason
|
local fr = chunk.choices[1].finish_reason
|
||||||
|
|
||||||
@ -75,11 +92,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
|||||||
content = delta.content or "",
|
content = delta.content or "",
|
||||||
done = (fr ~= nil)
|
done = (fr ~= nil)
|
||||||
}
|
}
|
||||||
|
if uses ~= nil then
|
||||||
|
unified.usage = uses
|
||||||
|
end
|
||||||
if delta.reasoning_content then
|
if delta.reasoning_content then
|
||||||
unified.reasoning_content = delta.reasoning_content
|
unified.reasoning_content = delta.reasoning_content
|
||||||
end
|
end
|
||||||
if delta.tool_calls then
|
if delta.tool_calls then
|
||||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
|
||||||
unified.tool_calls = delta.tool_calls
|
unified.tool_calls = delta.tool_calls
|
||||||
end
|
end
|
||||||
return json.encode(unified)
|
return json.encode(unified)
|
||||||
|
|||||||
@ -112,7 +112,25 @@ function adapter.transform_stream_chunk(raw_chunk)
|
|||||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||||
if not ok then return "" end
|
if not ok then return "" end
|
||||||
|
|
||||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||||
|
-- (the final usage chunk). Preserve it; the gateway emits it as the
|
||||||
|
-- terminal usage chunk. Note: keys must match Go's TokenUsage json tags
|
||||||
|
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||||
|
local uses = nil
|
||||||
|
if type(chunk.usage) == "table" then
|
||||||
|
uses = {
|
||||||
|
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||||
|
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||||
|
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||||
|
}
|
||||||
|
end
|
||||||
|
|
||||||
|
if not chunk.choices or #chunk.choices == 0 then
|
||||||
|
if uses ~= nil then
|
||||||
|
return json.encode({ usage = uses, done = (chunk.usage ~= nil) })
|
||||||
|
end
|
||||||
|
return ""
|
||||||
|
end
|
||||||
local delta = chunk.choices[1].delta or {}
|
local delta = chunk.choices[1].delta or {}
|
||||||
local fr = chunk.choices[1].finish_reason
|
local fr = chunk.choices[1].finish_reason
|
||||||
|
|
||||||
@ -120,6 +138,9 @@ function adapter.transform_stream_chunk(raw_chunk)
|
|||||||
content = delta.content or "",
|
content = delta.content or "",
|
||||||
done = (fr ~= nil)
|
done = (fr ~= nil)
|
||||||
}
|
}
|
||||||
|
if uses ~= nil then
|
||||||
|
unified.usage = uses
|
||||||
|
end
|
||||||
if delta.reasoning_content then
|
if delta.reasoning_content then
|
||||||
unified.reasoning_content = delta.reasoning_content
|
unified.reasoning_content = delta.reasoning_content
|
||||||
end
|
end
|
||||||
|
|||||||
@ -901,20 +901,52 @@ func standardSSEChunk(data string) string {
|
|||||||
} `json:"delta"`
|
} `json:"delta"`
|
||||||
FinishReason *string `json:"finish_reason"`
|
FinishReason *string `json:"finish_reason"`
|
||||||
} `json:"choices"`
|
} `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 {
|
if err := json.Unmarshal([]byte(data), &raw); err != nil {
|
||||||
return ""
|
return ""
|
||||||
}
|
}
|
||||||
if len(raw.Choices) == 0 {
|
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
|
||||||
return ""
|
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{
|
out, _ := json.Marshal(types.UnifiedChunk{
|
||||||
Content: raw.Choices[0].Delta.Content,
|
Content: func() string {
|
||||||
Done: raw.Choices[0].FinishReason != nil,
|
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)
|
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 {
|
func truncate(s string, n int) string {
|
||||||
if len(s) <= n {
|
if len(s) <= n {
|
||||||
return s
|
return s
|
||||||
|
|||||||
@ -138,6 +138,11 @@ type UnifiedChunk struct {
|
|||||||
Done bool `json:"done"`
|
Done bool `json:"done"`
|
||||||
ToolCalls json.RawMessage `json:"tool_calls,omitempty"`
|
ToolCalls json.RawMessage `json:"tool_calls,omitempty"`
|
||||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||||
|
// Usage carries the upstream token usage when the stream chunk provides
|
||||||
|
// it (OpenAI-style streams attach usage to some chunks; the final one
|
||||||
|
// often has empty choices). Gateway uses it to emit exact usage in the
|
||||||
|
// final stream chunk instead of estimates.
|
||||||
|
Usage *TokenUsage `json:"usage,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Meta passed to Lua build_headers hook
|
// Meta passed to Lua build_headers hook
|
||||||
|
|||||||
Reference in New Issue
Block a user