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

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"`
}