lisone/ds2api
1
0
Fork 0
mirror of https://github.com/CJackHwang/ds2api.git synced 2026-05-21 16:37:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

Propagate upstream token usage across Gemini usage metadata

Browse Source
This commit is contained in:
CJACK.
2026-04-07 10:16:00 +08:00
parent 96b8587c5b
commit 5bcea3d727
5 changed files with 42 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
internal/adapter/gemini/handler_test.go
View File

@@ -296,6 +296,32 @@ func TestGenerateContentOpenAIProxyErrorUsesGeminiEnvelope(t *testing.T) {
}
}
func TestBuildGeminiUsageOverridesPromptAndOutputTokensWhenProvided(t *testing.T) {
usage := buildGeminiUsage("prompt", "thinking", "answer", 11, 29)
if got, _ := usage["promptTokenCount"].(int); got != 11 {
t.Fatalf("expected promptTokenCount=11, got %#v", usage["promptTokenCount"])
}
if got, _ := usage["candidatesTokenCount"].(int); got != 29 {
t.Fatalf("expected candidatesTokenCount=29, got %#v", usage["candidatesTokenCount"])
}
if got, _ := usage["totalTokenCount"].(int); got != 40 {
t.Fatalf("expected totalTokenCount=40, got %#v", usage["totalTokenCount"])
}
}
func TestBuildGeminiUsageFallsBackToEstimateWhenNoUpstreamUsage(t *testing.T) {
usage := buildGeminiUsage("abcdef", "", "ghijkl", 0, 0)
if got, _ := usage["promptTokenCount"].(int); got <= 0 {
t.Fatalf("expected positive promptTokenCount estimate, got %#v", usage["promptTokenCount"])
}
if got, _ := usage["candidatesTokenCount"].(int); got <= 0 {
t.Fatalf("expected positive candidatesTokenCount estimate, got %#v", usage["candidatesTokenCount"])
}
if got, _ := usage["totalTokenCount"].(int); got <= 0 {
t.Fatalf("expected positive totalTokenCount estimate, got %#v", usage["totalTokenCount"])
}
}
func extractGeminiSSEFrames(t *testing.T, body string) []map[string]any {
t.Helper()
scanner := bufio.NewScanner(strings.NewReader(body))