mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 08:57:57 +00:00
fix(gateway): pass through upstream token usage in streams for all adapters
The prior usage-passthrough fix only covered openai/opencode; the same empty-choices+usage drop bug remained in the 5 sibling OpenAI-compatible adapters, and non-OpenAI providers (anthropic/gemini/ollama) never surfaced streaming usage at all. - deepseek/github/groq/kimicode/mistral: preserve usage on empty-choices chunks and attach it to normal chunks (same pattern as openai.lua) - anthropic: emit usage from message_start (prompt) and message_delta (completion); gateway merges split usage additively - gemini: read usageMetadata in the stream path - ollama: fix non-streaming key (usage -> token_usage, matches UnifiedResponse json tag) and read prompt_eval_count/eval_count; surface counts from the done stream chunk - gateway: mergeUsage combines usage across chunks (non-zero fields win, total recomputed from prompt+completion) so split usage doesn't lose the prompt half; single-chunk case (OpenAI) preserved exactly - usage-only chunks: done=false (no redundant terminal stop), matching the Go fallback standardSSEChunk
This commit is contained in:
@ -580,6 +580,32 @@ func (g *Gateway) writeRec(rec *Req) {
|
||||
g.stats.Record(*rec)
|
||||
}
|
||||
|
||||
// mergeUsage combines token usage across stream chunks additively. Some
|
||||
// providers split usage across chunks (e.g. Anthropic reports prompt tokens
|
||||
// in message_start and the final completion tokens in message_delta); a plain
|
||||
// "last non-nil wins" would discard the prompt half. Non-zero fields from cur
|
||||
// override prev; total is recomputed from the merged parts so a partial later
|
||||
// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal
|
||||
// empty-choices+usage chunk) upstream totals are preserved exactly.
|
||||
func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
|
||||
if prev == nil {
|
||||
u := *cur
|
||||
if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) {
|
||||
u.Total = u.Prompt + u.Completion
|
||||
}
|
||||
return &u
|
||||
}
|
||||
out := *prev
|
||||
if cur.Prompt > 0 {
|
||||
out.Prompt = cur.Prompt
|
||||
}
|
||||
if cur.Completion > 0 {
|
||||
out.Completion = cur.Completion
|
||||
}
|
||||
out.Total = out.Prompt + out.Completion
|
||||
return &out
|
||||
}
|
||||
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
@ -632,7 +658,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
var lastUsage *types.TokenUsage
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = ck.Usage
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
@ -802,7 +828,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
var lastUsage *types.TokenUsage
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = ck.Usage
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
|
||||
Reference in New Issue
Block a user