diff --git a/internal/llm/usage_resolver.go b/internal/llm/usage_resolver.go index 1345420..661d659 100644 --- a/internal/llm/usage_resolver.go +++ b/internal/llm/usage_resolver.go @@ -27,15 +27,23 @@ var completionTokensPaths = []string{ } var cacheReadTokensPaths = []string{ - "usage.cache_read_input_tokens", // Anthropic - "cache_read_input_tokens", // flat at root - "usage.prompt_tokens_details.cache_tokens_hit", // some providers - "usage.prompt_tokens_details.cache_tokens", // some providers + "usage.cache_read_input_tokens", // Anthropic + "cache_read_input_tokens", // flat at root + "data.usage.cache_read_input_tokens", // wrapped Anthropic-compatible proxy + "usage.prompt_tokens_details.cached_tokens", // OpenAI-compatible providers + "data.usage.prompt_tokens_details.cached_tokens", // wrapped OpenAI-compatible providers + "usage.prompt_tokens_details.cache_tokens_hit", // some providers + "data.usage.prompt_tokens_details.cache_tokens_hit", + "usage.prompt_tokens_details.cache_tokens", // some providers + "data.usage.prompt_tokens_details.cache_tokens", } var cacheWriteTokensPaths = []string{ - "usage.cache_creation_input_tokens", // Anthropic / proxy - "cache_creation_input_tokens", // flat at root + "usage.cache_creation_input_tokens", // Anthropic / proxy + "cache_creation_input_tokens", // flat at root + "data.usage.cache_creation_input_tokens", // wrapped Anthropic-compatible proxy + "usage.prompt_tokens_details.cache_creation_tokens", + "data.usage.prompt_tokens_details.cache_creation_tokens", } // totalTokensPaths is an ordered list of JSON paths to try when extracting diff --git a/internal/llm/usage_resolver_test.go b/internal/llm/usage_resolver_test.go new file mode 100644 index 0000000..5b67099 --- /dev/null +++ b/internal/llm/usage_resolver_test.go @@ -0,0 +1,80 @@ +package llm + +import "testing" + +func TestResolveUsageOpenAICompatibleCachedTokens(t *testing.T) { + usage := resolveUsage([]byte(`{ + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "prompt_tokens_details": { + "cached_tokens": 75 + } + } + }`)) + + if usage == nil { + t.Fatal("resolveUsage returned nil") + } + if usage.CacheReadTokens != 75 { + t.Errorf("CacheReadTokens = %d, want 75", usage.CacheReadTokens) + } + if usage.PromptTokens != 100 { + t.Errorf("PromptTokens = %d, want 100", usage.PromptTokens) + } + if usage.CompletionTokens != 20 { + t.Errorf("CompletionTokens = %d, want 20", usage.CompletionTokens) + } +} + +func TestResolveUsageWrappedCachedTokens(t *testing.T) { + usage := resolveUsage([]byte(`{ + "data": { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "prompt_tokens_details": { + "cached_tokens": 75, + "cache_creation_tokens": 10 + } + } + } + }`)) + + if usage == nil { + t.Fatal("resolveUsage returned nil") + } + if usage.CacheReadTokens != 75 { + t.Errorf("CacheReadTokens = %d, want 75", usage.CacheReadTokens) + } + if usage.CacheWriteTokens != 10 { + t.Errorf("CacheWriteTokens = %d, want 10", usage.CacheWriteTokens) + } + if usage.TotalTokens != 205 { + t.Errorf("TotalTokens = %d, want 205", usage.TotalTokens) + } +} + +func TestResolveUsageWrappedAnthropicCompatibleCacheTokens(t *testing.T) { + usage := resolveUsage([]byte(`{ + "data": { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 15 + } + } + }`)) + + if usage == nil { + t.Fatal("resolveUsage returned nil") + } + if usage.CacheReadTokens != 40 { + t.Errorf("CacheReadTokens = %d, want 40", usage.CacheReadTokens) + } + if usage.CacheWriteTokens != 15 { + t.Errorf("CacheWriteTokens = %d, want 15", usage.CacheWriteTokens) + } +}