eric/promptkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Preserve reported token usage presence

Browse Source
This commit is contained in:
Eric Rakestraw
2026-09-23 16:39:43 +00:00
parent ea94d66fc2
commit d31a563ca1
8 changed files with 166 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
internal/llm/openai_compatible_client_test.go
View File

@@ -881,6 +881,44 @@ func checkCacheUsageMapping(t *testing.T) {
}
}
func TestOpenAICompatibleUsagePresence(t *testing.T) {
for _, tc := range []struct {
name string
usage string
want domain.TokenUsage
presence domain.TokenUsagePresence
}{
{name: "absent"},
{name: "null", usage: `"usage":null,`},
{name: "partial nested", usage: `"usage":{"prompt_tokens":7,"prompt_tokens_details":{"cached_tokens":0}},`,
want: domain.TokenUsage{PromptTokens: 7}, presence: domain.TokenUsagePresence{PromptTokens: true, CachedTokens: true}},
{name: "explicit null fields", usage: `"usage":{"prompt_tokens":null,"completion_tokens":0,"prompt_tokens_details":{"cached_tokens":null}},`,
presence: domain.TokenUsagePresence{CompletionTokens: true}},
{name: "all reported", usage: `"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":4},"cache_write_tokens":5},`,
want: domain.TokenUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3, CachedTokens: 4, CacheWriteTokens: 5},
presence: domain.TokenUsagePresence{PromptTokens: true, CompletionTokens: true, TotalTokens: true, CachedTokens: true, CacheWriteTokens: true}},
} {
t.Run(tc.name, func(t *testing.T) {
provider := newRecordingProvider(t)
provider.respond(http.StatusOK, `{`+tc.usage+`"choices":[{"message":{"content":"ok"}}]}`)
client := newProviderClient(t, provider, OpenAICompatibleConfig{Model: "model"})
resp, err := client.Generate(context.Background(), ordinaryGenerateRequest())
if err != nil {
t.Fatal(err)
}
if resp.Usage != tc.want || resp.UsagePresence == nil || *resp.UsagePresence != tc.presence {
t.Fatalf("usage=%+v, presence=%+v", resp.Usage, resp.UsagePresence)
}
})
}
provider := newRecordingProvider(t)
provider.respond(http.StatusOK, `{"usage":{"prompt_tokens":"seven"},"choices":[{"message":{"content":"ok"}}]}`)
client := newProviderClient(t, provider, OpenAICompatibleConfig{Model: "model"})
if resp, err := client.Generate(context.Background(), ordinaryGenerateRequest()); resp != nil || !errors.Is(err, ErrMalformedResponse) {
t.Fatalf("invalid numeric usage response=%+v, error=%v", resp, err)
}
}
func checkOptionalResponseFieldOmission(t *testing.T) {
provider := newRecordingProvider(t)
client := newProviderClient(t, provider, OpenAICompatibleConfig{})