jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 00:48:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(gateway): pass through exact upstream token usage in streams

Browse Source 
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:
valuelesser
2026-08-18 19:04:44 +08:00
parent 6f1c806591
commit 83e6d88813
5 changed files with 113 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
internal/gateway/chat.go
View File

@ -629,7 +629,11 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
}) {
return
}
var lastUsage *types.TokenUsage
for ck := range chunks {
if ck.Usage != nil {
lastUsage = ck.Usage
}
chunk := ChatChunk{
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}},
})
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
if rec.Prompt+rec.Compl > 0 {
usage := types.TokenUsage{
// Prefer the upstream's exact usage if the stream carried it; fall back to
// 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),
Completion: int(rec.Compl),
Total: int(rec.Prompt + rec.Compl),
}
tut = &u
}
if tut != nil {
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
Choices: []ChunkChoice{},
Usage: &usage,
Usage: tut,
})
}
fmt.Fprintf(w, "data: [DONE]\n\n")
@ -787,7 +799,11 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
}) {
return
}
var lastUsage *types.TokenUsage
for ck := range chunks {
if ck.Usage != nil {
lastUsage = ck.Usage
}
chunk := ChatChunk{
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}},
})
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
if rec.Prompt+rec.Compl > 0 {
usage := types.TokenUsage{
// Prefer the upstream's exact usage if the stream carried it; fall back to
// 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),
Completion: int(rec.Compl),
Total: int(rec.Prompt + rec.Compl),
}
tut = &u
}
if tut != nil {
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
Choices: []ChunkChoice{},
Usage: &usage,
Usage: tut,
})
}
fmt.Fprintf(w, "data: [DONE]\n\n")