From 9816296c25b378a4837977425e109c6978e9dd9e Mon Sep 17 00:00:00 2001 From: zhangzherui Date: Wed, 22 Jul 2026 23:55:54 +0800 Subject: [PATCH] feat(schema): track cache creation tokens --- components/model/callback_extra.go | 4 +++- schema/agentic_message.go | 3 +++ schema/agentic_message_test.go | 7 +++++-- schema/message.go | 7 ++++++- schema/message_test.go | 29 ++++++++++++++++++++++++++--- 5 files changed, 43 insertions(+), 7 deletions(-) diff --git a/components/model/callback_extra.go b/components/model/callback_extra.go index 8591c4373..af1392f34 100644 --- a/components/model/callback_extra.go +++ b/components/model/callback_extra.go @@ -44,8 +44,10 @@ type CompletionTokensDetails struct { // PromptTokenDetails provides a breakdown of prompt token usage. type PromptTokenDetails struct { - // Cached tokens present in the prompt. + // CachedTokens is the number of prompt tokens read from a provider cache. CachedTokens int + // CacheCreationTokens is the number of prompt tokens written to a provider cache. + CacheCreationTokens int } // Config is the config for the model. diff --git a/schema/agentic_message.go b/schema/agentic_message.go index 2474e5df8..a3e44fb85 100644 --- a/schema/agentic_message.go +++ b/schema/agentic_message.go @@ -1126,6 +1126,9 @@ func concatTokenUsage(usages []*TokenUsage) *TokenUsage { if usage.PromptTokenDetails.CachedTokens > ret.PromptTokenDetails.CachedTokens { ret.PromptTokenDetails.CachedTokens = usage.PromptTokenDetails.CachedTokens } + if usage.PromptTokenDetails.CacheCreationTokens > ret.PromptTokenDetails.CacheCreationTokens { + ret.PromptTokenDetails.CacheCreationTokens = usage.PromptTokenDetails.CacheCreationTokens + } if usage.CompletionTokensDetails.ReasoningTokens > ret.CompletionTokensDetails.ReasoningTokens { ret.CompletionTokensDetails.ReasoningTokens = usage.CompletionTokensDetails.ReasoningTokens } diff --git a/schema/agentic_message_test.go b/schema/agentic_message_test.go index 32dc96c2e..6f8e0f4bd 100644 --- a/schema/agentic_message_test.go +++ b/schema/agentic_message_test.go @@ -811,7 +811,8 @@ func TestConcatAgenticMessages(t *testing.T) { PromptTokens: 10, CompletionTokens: 5, PromptTokenDetails: PromptTokenDetails{ - CachedTokens: 3, + CachedTokens: 3, + CacheCreationTokens: 7, }, CompletionTokensDetails: CompletionTokensDetails{ ReasoningTokens: 2, @@ -827,7 +828,8 @@ func TestConcatAgenticMessages(t *testing.T) { PromptTokens: 8, CompletionTokens: 15, PromptTokenDetails: PromptTokenDetails{ - CachedTokens: 5, + CachedTokens: 5, + CacheCreationTokens: 4, }, CompletionTokensDetails: CompletionTokensDetails{ ReasoningTokens: 4, @@ -845,6 +847,7 @@ func TestConcatAgenticMessages(t *testing.T) { assert.Equal(t, 10, result.ResponseMeta.TokenUsage.PromptTokens) assert.Equal(t, 23, result.ResponseMeta.TokenUsage.TotalTokens) assert.Equal(t, 5, result.ResponseMeta.TokenUsage.PromptTokenDetails.CachedTokens) + assert.Equal(t, 7, result.ResponseMeta.TokenUsage.PromptTokenDetails.CacheCreationTokens) assert.Equal(t, 4, result.ResponseMeta.TokenUsage.CompletionTokensDetails.ReasoningTokens) }) diff --git a/schema/message.go b/schema/message.go index 892abe702..377dc97af 100644 --- a/schema/message.go +++ b/schema/message.go @@ -554,8 +554,10 @@ type CompletionTokensDetails struct { // PromptTokenDetails provides a breakdown of prompt token usage. type PromptTokenDetails struct { - // Cached tokens present in the prompt. + // CachedTokens is the number of prompt tokens read from a provider cache. CachedTokens int `json:"cached_tokens"` + // CacheCreationTokens is the number of prompt tokens written to a provider cache. + CacheCreationTokens int `json:"cache_creation_tokens"` } var _ MessagesTemplate = &Message{} @@ -1752,6 +1754,9 @@ func ConcatMessages(msgs []*Message) (*Message, error) { if msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens > ret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens { ret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens = msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens } + if msg.ResponseMeta.Usage.PromptTokenDetails.CacheCreationTokens > ret.ResponseMeta.Usage.PromptTokenDetails.CacheCreationTokens { + ret.ResponseMeta.Usage.PromptTokenDetails.CacheCreationTokens = msg.ResponseMeta.Usage.PromptTokenDetails.CacheCreationTokens + } if msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens > ret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens { ret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens = msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens diff --git a/schema/message_test.go b/schema/message_test.go index bb086b8c6..6e99632f7 100644 --- a/schema/message_test.go +++ b/schema/message_test.go @@ -18,6 +18,7 @@ package schema import ( "context" + "encoding/json" "reflect" "sync" "testing" @@ -27,6 +28,25 @@ import ( "github.com/cloudwego/eino/internal/generic" ) +func TestPromptTokenDetailsJSON(t *testing.T) { + details := PromptTokenDetails{ + CachedTokens: 3, + CacheCreationTokens: 5, + } + + data, err := json.Marshal(details) + assert.NoError(t, err) + assert.JSONEq(t, `{"cached_tokens":3,"cache_creation_tokens":5}`, string(data)) + + var decoded PromptTokenDetails + assert.NoError(t, json.Unmarshal(data, &decoded)) + assert.Equal(t, details, decoded) + + var legacy PromptTokenDetails + assert.NoError(t, json.Unmarshal([]byte(`{"cached_tokens":3}`), &legacy)) + assert.Equal(t, PromptTokenDetails{CachedTokens: 3}, legacy) +} + func TestMessageTemplate(t *testing.T) { pyFmtMessage := UserMessage("input: {question}") jinja2Message := UserMessage("input: {{question}}") @@ -157,7 +177,8 @@ func TestConcatMessage(t *testing.T) { CompletionTokens: 15, PromptTokens: 30, PromptTokenDetails: PromptTokenDetails{ - CachedTokens: 15, + CachedTokens: 15, + CacheCreationTokens: 20, }, CompletionTokensDetails: CompletionTokensDetails{ ReasoningTokens: 8, @@ -179,7 +200,8 @@ func TestConcatMessage(t *testing.T) { CompletionTokens: 10, PromptTokens: 20, PromptTokenDetails: PromptTokenDetails{ - CachedTokens: 10, + CachedTokens: 10, + CacheCreationTokens: 20, }, CompletionTokensDetails: CompletionTokensDetails{ ReasoningTokens: 5, @@ -201,7 +223,8 @@ func TestConcatMessage(t *testing.T) { CompletionTokens: 15, PromptTokens: 30, PromptTokenDetails: PromptTokenDetails{ - CachedTokens: 15, + CachedTokens: 15, + CacheCreationTokens: 12, }, CompletionTokensDetails: CompletionTokensDetails{ ReasoningTokens: 8,