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

14
internal/domain/domain.go
View File

@@ -306,8 +306,18 @@ type StructuredOutputJSONSpec struct {
// GenerateResponse is the response received from the LLM client.
type GenerateResponse struct {
Content string
Usage TokenUsage
Content string
Usage TokenUsage
UsagePresence *TokenUsagePresence
}
// TokenUsagePresence records which token fields a provider reported.
type TokenUsagePresence struct {
PromptTokens bool
CompletionTokens bool
TotalTokens bool
CachedTokens bool
CacheWriteTokens bool
}
// TokenUsage tracks token consumption.

34
internal/llm/openai_compatible_client.go
View File

@@ -187,15 +187,29 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
return &domain.GenerateResponse{
Content: *content,
Usage: domain.TokenUsage{
PromptTokens: wireResp.Usage.PromptTokens,
CompletionTokens: wireResp.Usage.CompletionTokens,
TotalTokens: wireResp.Usage.TotalTokens,
CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens,
CacheWriteTokens: wireResp.Usage.CacheWriteTokens,
PromptTokens: tokenCount(wireResp.Usage.PromptTokens),
CompletionTokens: tokenCount(wireResp.Usage.CompletionTokens),
TotalTokens: tokenCount(wireResp.Usage.TotalTokens),
CachedTokens: tokenCount(wireResp.Usage.PromptTokensDetails.CachedTokens),
CacheWriteTokens: tokenCount(wireResp.Usage.CacheWriteTokens),
},
UsagePresence: &domain.TokenUsagePresence{
PromptTokens: wireResp.Usage.PromptTokens != nil,
CompletionTokens: wireResp.Usage.CompletionTokens != nil,
TotalTokens: wireResp.Usage.TotalTokens != nil,
CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens != nil,
CacheWriteTokens: wireResp.Usage.CacheWriteTokens != nil,
},
}, nil
}
func tokenCount(value *int) int {
if value == nil {
return 0
}
return *value
}
func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {
limited := &io.LimitedReader{
R: body,
@@ -389,13 +403,13 @@ type openAIChatResponse struct {
Message openAIChatResponseMessage `json:"message"`
} `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokens *int `json:"prompt_tokens"`
CompletionTokens *int `json:"completion_tokens"`
TotalTokens *int `json:"total_tokens"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
CachedTokens *int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CacheWriteTokens int `json:"cache_write_tokens"`
CacheWriteTokens *int `json:"cache_write_tokens"`
} `json:"usage"`
}

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{})