mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
fix(gateway): emit OpenAI-standard token usage in responses
- TokenUsage.MarshalJSON now emits both standard (prompt_tokens, completion_tokens, total_tokens) and legacy (prompt, completion, total) keys, so OpenAI-compatible clients (DSH, DevEco Code, etc.) can read token usage. - ChatChunk gains an optional Usage field; stream responses now send a final usage chunk (empty choices) before [DONE]. Refs: usage not visible in clients because the gateway serialized only the internal short keys and never emitted a streaming usage chunk.
This commit is contained in:
@ -53,11 +53,14 @@ type RespMessage struct {
|
||||
}
|
||||
|
||||
type ChatChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
// Usage is sent in the final chunk of a stream (empty choices) so
|
||||
// OpenAI-compatible clients can read token usage.
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type ChunkChoice struct {
|
||||
@ -653,6 +656,19 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
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{
|
||||
Prompt: int(rec.Prompt),
|
||||
Completion: int(rec.Compl),
|
||||
Total: int(rec.Prompt + rec.Compl),
|
||||
}
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: &usage,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
@ -798,6 +814,19 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
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{
|
||||
Prompt: int(rec.Prompt),
|
||||
Completion: int(rec.Compl),
|
||||
Total: int(rec.Prompt + rec.Compl),
|
||||
}
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: &usage,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
|
||||
@ -85,6 +85,28 @@ type TokenUsage struct {
|
||||
Total int `json:"total"`
|
||||
}
|
||||
|
||||
// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
|
||||
// by the internal unified representation and older clients) and the OpenAI
|
||||
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||
return json.Marshal(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"`
|
||||
}{
|
||||
PromptTokens: t.Prompt,
|
||||
CompletionTokens: t.Completion,
|
||||
TotalTokens: t.Total,
|
||||
Prompt: t.Prompt,
|
||||
Completion: t.Completion,
|
||||
Total: t.Total,
|
||||
})
|
||||
}
|
||||
|
||||
type ImageData struct {
|
||||
B64JSON string `json:"b64_json,omitempty"`
|
||||
URL string `json:"url,omitempty"`
|
||||
|
||||
Reference in New Issue
Block a user