diff --git a/libs/acl/openai/cache_usage_test.go b/libs/acl/openai/cache_usage_test.go new file mode 100644 index 000000000..987ba7398 --- /dev/null +++ b/libs/acl/openai/cache_usage_test.go @@ -0,0 +1,85 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package openai + +import ( + "encoding/json" + "testing" + + "github.com/cloudwego/eino/schema" + openai "github.com/meguminnnnnnnnn/go-openai" + "github.com/stretchr/testify/assert" +) + +func TestToEinoTokenUsageCachedTokens(t *testing.T) { + tests := []struct { + name string + usage *openai.Usage + want int + }{ + { + name: "OpenAI prompt token details", + usage: &openai.Usage{ + PromptTokensDetails: &openai.PromptTokensDetails{CachedTokens: 64}, + ExtraFields: map[string]json.RawMessage{ + "prompt_cache_hit_tokens": json.RawMessage(`32`), + }, + }, + want: 64, + }, + { + name: "DeepSeek top-level cache hit extension", + usage: &openai.Usage{ + ExtraFields: map[string]json.RawMessage{ + "prompt_cache_hit_tokens": json.RawMessage(`48`), + }, + }, + want: 48, + }, + { + name: "zero standard detail falls back to DeepSeek extension", + usage: &openai.Usage{ + PromptTokensDetails: &openai.PromptTokensDetails{}, + ExtraFields: map[string]json.RawMessage{ + "prompt_cache_hit_tokens": json.RawMessage(`32`), + }, + }, + want: 32, + }, + { + name: "malformed extension", + usage: &openai.Usage{ + ExtraFields: map[string]json.RawMessage{ + "prompt_cache_hit_tokens": json.RawMessage(`"invalid"`), + }, + }, + want: 0, + }, + { + name: "missing cache details", + usage: &openai.Usage{}, + want: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := toEinoTokenUsage(tt.usage) + assert.Equal(t, schema.PromptTokenDetails{CachedTokens: tt.want}, got.PromptTokenDetails) + }) + } +} diff --git a/libs/acl/openai/chat_model.go b/libs/acl/openai/chat_model.go index 4b3f98f25..4219cfc7e 100644 --- a/libs/acl/openai/chat_model.go +++ b/libs/acl/openai/chat_model.go @@ -1335,6 +1335,19 @@ func toEinoTokenUsage(usage *openai.Usage) *schema.TokenUsage { if usage.PromptTokensDetails != nil { promptTokenDetails.CachedTokens = usage.PromptTokensDetails.CachedTokens } + if promptTokenDetails.CachedTokens == 0 { + // DeepSeek's OpenAI-compatible API reports cache hits as a top-level + // usage extension instead of prompt_tokens_details.cached_tokens. + // go-openai preserves unknown usage fields in ExtraFields. Falling back + // on a zero standard value also covers gateways that synthesize an empty + // prompt_tokens_details object while retaining DeepSeek's extension. + if raw, ok := usage.ExtraFields["prompt_cache_hit_tokens"]; ok { + var cachedTokens int + if err := json.Unmarshal(raw, &cachedTokens); err == nil { + promptTokenDetails.CachedTokens = cachedTokens + } + } + } completionTokensDetails := schema.CompletionTokensDetails{} if usage.CompletionTokensDetails != nil { completionTokensDetails.ReasoningTokens = usage.CompletionTokensDetails.ReasoningTokens