vrr/open-code-review
2
0
Fork 0
mirror of https://github.com/alibaba/open-code-review.git synced 2026-07-10 01:39:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

feat: enhance token usage resolution for OpenAI and Anthropic compatibility

Browse source
This commit is contained in:
Pesto 2026-06-25 19:19:58 +02:00
parent 779c4a8290
commit 6dd3940567
2 changed files with 94 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

20
internal/llm/usage_resolver.go
View file

@ -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

80
internal/llm/usage_resolver_test.go Normal file
View file

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