From 8d11868cb2af5d220740cbd1b5965d9024e5ff01 Mon Sep 17 00:00:00 2001 From: pangju Date: Wed, 2 Sep 2026 20:18:24 +0800 Subject: [PATCH 1/2] feat(agenticopenai): expose cache write tokens --- components/model/agenticopenai/consts.go | 2 + .../model/agenticopenai/message_extra.go | 58 +++++++ .../model/agenticopenai/message_extra_test.go | 163 ++++++++++++++++++ .../agenticopenai/responses_convertor.go | 1 + .../responses_event_convertor.go | 30 ++-- 5 files changed, 241 insertions(+), 13 deletions(-) create mode 100644 components/model/agenticopenai/message_extra.go create mode 100644 components/model/agenticopenai/message_extra_test.go diff --git a/components/model/agenticopenai/consts.go b/components/model/agenticopenai/consts.go index 0d8415eee..c92c254f3 100644 --- a/components/model/agenticopenai/consts.go +++ b/components/model/agenticopenai/consts.go @@ -22,6 +22,8 @@ const responsesImplType = "AgenticOpenAI/Responses" const defaultBaseURL = "https://api.openai.com/v1" +const keyOfCacheWriteTokens = "_eino_openai_cache_write_tokens" + type ServerToolName string const ( diff --git a/components/model/agenticopenai/message_extra.go b/components/model/agenticopenai/message_extra.go new file mode 100644 index 000000000..10f54d55e --- /dev/null +++ b/components/model/agenticopenai/message_extra.go @@ -0,0 +1,58 @@ +/* + * Copyright 2026 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 + * + * https://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 agenticopenai + +import ( + "strconv" + + "github.com/cloudwego/eino/schema" + "github.com/openai/openai-go/v3/responses" +) + +// GetCacheWriteTokens returns the OpenAI cache_write_tokens count from the +// message. This is the number of input tokens written to the prompt cache. +// Pricing for cache writes depends on the model. +// +// The cache-read side is available through the standard token usage path: +// +// msg.ResponseMeta.TokenUsage.PromptTokenDetails.CachedTokens +// +// When streaming, OpenAI reports cache-write usage on a response lifecycle +// event. schema.ConcatAgenticMessages merges Extra maps with a last-value-wins +// policy for int, so the final concatenated message preserves the count. +func GetCacheWriteTokens(msg *schema.AgenticMessage) (int, bool) { + if msg == nil || msg.Extra == nil { + return 0, false + } + tokens, ok := msg.Extra[keyOfCacheWriteTokens].(int) + return tokens, ok +} + +func cacheWriteTokensExtra(resp *responses.Response) map[string]any { + if resp == nil { + return nil + } + field, ok := resp.Usage.InputTokensDetails.JSON.ExtraFields["cache_write_tokens"] + if !ok { + return nil + } + tokens, err := strconv.Atoi(field.Raw()) + if err != nil || tokens <= 0 { + return nil + } + return map[string]any{keyOfCacheWriteTokens: tokens} +} diff --git a/components/model/agenticopenai/message_extra_test.go b/components/model/agenticopenai/message_extra_test.go new file mode 100644 index 000000000..41b623d66 --- /dev/null +++ b/components/model/agenticopenai/message_extra_test.go @@ -0,0 +1,163 @@ +/* + * Copyright 2026 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 + * + * https://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 agenticopenai + +import ( + "encoding/json" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" + "github.com/openai/openai-go/v3/responses" +) + +func TestGetCacheWriteTokens(t *testing.T) { + tests := []struct { + name string + msg *schema.AgenticMessage + wantTokens int + wantOK bool + }{ + {name: "nil message"}, + {name: "nil extra", msg: &schema.AgenticMessage{}}, + {name: "missing key", msg: &schema.AgenticMessage{Extra: map[string]any{"other": 42}}}, + { + name: "cache write tokens present", + msg: &schema.AgenticMessage{Extra: map[string]any{ + keyOfCacheWriteTokens: 1234, + }}, + wantTokens: 1234, + wantOK: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotTokens, gotOK := GetCacheWriteTokens(tt.msg) + if gotTokens != tt.wantTokens || gotOK != tt.wantOK { + t.Fatalf("GetCacheWriteTokens() = (%d, %v), want (%d, %v)", + gotTokens, gotOK, tt.wantTokens, tt.wantOK) + } + }) + } +} + +func TestCacheWriteTokensSetOnGeneratedMessage(t *testing.T) { + t.Run("zero cache write tokens are omitted", func(t *testing.T) { + resp := mustUnmarshalResponse(t, `{ + "id": "resp_1", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 100, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 20, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 120 + } + }`) + + msg, err := toOutputMessage(resp, &model.Options{}) + if err != nil { + t.Fatal(err) + } + if tokens, ok := GetCacheWriteTokens(msg); ok || tokens != 0 { + t.Fatalf("expected (0, false), got (%d, %v)", tokens, ok) + } + }) + + t.Run("nonzero cache write tokens are exposed", func(t *testing.T) { + resp := mustUnmarshalResponse(t, `{ + "id": "resp_1", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 600, + "input_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 500}, + "output_tokens": 20, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 620 + } + }`) + + msg, err := toOutputMessage(resp, &model.Options{}) + if err != nil { + t.Fatal(err) + } + if tokens, ok := GetCacheWriteTokens(msg); !ok || tokens != 500 { + t.Fatalf("expected (500, true), got (%d, %v)", tokens, ok) + } + }) +} + +func TestCacheWriteTokensSetOnStreamingResponse(t *testing.T) { + resp := mustUnmarshalResponse(t, `{ + "id": "resp_1", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 400, + "input_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 300}, + "output_tokens": 20, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 420 + } + }`) + + sr, sw := schema.Pipe[*model.AgenticCallbackOutput](1) + reader := sr.Copy(1)[0] + sender := newCallbackSender(sw, &model.AgenticConfig{}) + sender.sendResponse(resp, nil) + + out, err := reader.Recv() + if err != nil { + t.Fatal(err) + } + if tokens, ok := GetCacheWriteTokens(out.Message); !ok || tokens != 300 { + t.Fatalf("expected (300, true), got (%d, %v)", tokens, ok) + } +} + +func TestCacheWriteTokensPreservedAfterConcat(t *testing.T) { + usageChunk := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + Extra: map[string]any{keyOfCacheWriteTokens: 300}, + } + textChunk := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlockChunk(&schema.AssistantGenText{Text: "hello"}, &schema.StreamingMeta{Index: 0}), + }, + } + + msg, err := schema.ConcatAgenticMessages([]*schema.AgenticMessage{usageChunk, textChunk}) + if err != nil { + t.Fatal(err) + } + if tokens, ok := GetCacheWriteTokens(msg); !ok || tokens != 300 { + t.Fatalf("expected (300, true), got (%d, %v)", tokens, ok) + } +} + +func mustUnmarshalResponse(t *testing.T, raw string) *responses.Response { + t.Helper() + var resp responses.Response + if err := json.Unmarshal([]byte(raw), &resp); err != nil { + t.Fatalf("json.Unmarshal(response) error = %v", err) + } + return &resp +} diff --git a/components/model/agenticopenai/responses_convertor.go b/components/model/agenticopenai/responses_convertor.go index 52cc2297e..447071caa 100644 --- a/components/model/agenticopenai/responses_convertor.go +++ b/components/model/agenticopenai/responses_convertor.go @@ -1598,6 +1598,7 @@ func toOutputMessage(resp *responses.Response, options *model.Options) (msg *sch Role: schema.AgenticRoleTypeAssistant, ContentBlocks: blocks, ResponseMeta: responseObjectToResponseMeta(resp), + Extra: cacheWriteTokensExtra(resp), } return msg, nil diff --git a/components/model/agenticopenai/responses_event_convertor.go b/components/model/agenticopenai/responses_event_convertor.go index 8d9692e99..d4ea3f29c 100644 --- a/components/model/agenticopenai/responses_event_convertor.go +++ b/components/model/agenticopenai/responses_event_convertor.go @@ -48,24 +48,19 @@ func receivedStreamingResponse(sr *ssestream.Stream[responses.ResponseStreamEven _ = sw.Send(nil, fmt.Errorf("received error event: code=%s message=%s", variant.Code, variant.Message)) case responses.ResponseCreatedEvent: - meta := responseObjectToResponseMeta(&variant.Response) - sender.sendMeta(meta, nil) + sender.sendResponse(&variant.Response, nil) case responses.ResponseInProgressEvent: - meta := responseObjectToResponseMeta(&variant.Response) - sender.sendMeta(meta, nil) + sender.sendResponse(&variant.Response, nil) case responses.ResponseCompletedEvent: - meta := responseObjectToResponseMeta(&variant.Response) - sender.sendMeta(meta, nil) + sender.sendResponse(&variant.Response, nil) case responses.ResponseIncompleteEvent: - meta := responseObjectToResponseMeta(&variant.Response) - sender.sendMeta(meta, nil) + sender.sendResponse(&variant.Response, nil) case responses.ResponseFailedEvent: - meta := responseObjectToResponseMeta(&variant.Response) - sender.sendMeta(meta, nil) + sender.sendResponse(&variant.Response, nil) case responses.ResponseOutputItemAddedEvent: blocks, err := receiver.itemAddedEventToContentBlock(variant) @@ -225,16 +220,24 @@ func newCallbackSender(sw *schema.StreamWriter[*model.AgenticCallbackOutput], co } func (s *callbackSender) sendMeta(meta *schema.AgenticResponseMeta, err error) { - s.send(meta, nil, err) + s.send(meta, nil, nil, err) +} + +func (s *callbackSender) sendResponse(resp *responses.Response, err error) { + if resp == nil { + s.send(nil, nil, nil, err) + return + } + s.send(responseObjectToResponseMeta(resp), nil, cacheWriteTokensExtra(resp), err) } func (s *callbackSender) sendBlock(block *schema.ContentBlock, err error) { if block != nil || err != nil { - s.send(nil, block, err) + s.send(nil, block, nil, err) } } -func (s *callbackSender) send(meta *schema.AgenticResponseMeta, block *schema.ContentBlock, err error) { +func (s *callbackSender) send(meta *schema.AgenticResponseMeta, block *schema.ContentBlock, extra map[string]any, err error) { if err != nil { _ = s.sw.Send(nil, fmt.Errorf("%s: %w", s.errHeader, err)) return @@ -243,6 +246,7 @@ func (s *callbackSender) send(meta *schema.AgenticResponseMeta, block *schema.Co msg := &schema.AgenticMessage{ Role: schema.AgenticRoleTypeAssistant, ResponseMeta: meta, + Extra: extra, } if block != nil { From b7fdff07eb85a0a93506c75943742e366654a02d Mon Sep 17 00:00:00 2001 From: pangju Date: Wed, 2 Sep 2026 20:25:41 +0800 Subject: [PATCH 2/2] test(agenticopenai): cover cache write edge cases --- .../model/agenticopenai/message_extra_test.go | 58 +++++++++++-------- .../responses_event_convertor.go | 4 -- .../responses_event_convertor_test.go | 16 +---- .../agenticopenai/responses_model_test.go | 30 +++++++--- 4 files changed, 58 insertions(+), 50 deletions(-) diff --git a/components/model/agenticopenai/message_extra_test.go b/components/model/agenticopenai/message_extra_test.go index 41b623d66..d0c55c09b 100644 --- a/components/model/agenticopenai/message_extra_test.go +++ b/components/model/agenticopenai/message_extra_test.go @@ -18,6 +18,7 @@ package agenticopenai import ( "encoding/json" + "fmt" "testing" "github.com/cloudwego/eino/components/model" @@ -104,31 +105,42 @@ func TestCacheWriteTokensSetOnGeneratedMessage(t *testing.T) { }) } -func TestCacheWriteTokensSetOnStreamingResponse(t *testing.T) { - resp := mustUnmarshalResponse(t, `{ - "id": "resp_1", - "status": "completed", - "output": [], - "usage": { - "input_tokens": 400, - "input_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 300}, - "output_tokens": 20, - "output_tokens_details": {"reasoning_tokens": 0}, - "total_tokens": 420 - } - }`) +func TestCacheWriteTokensInvalidValuesAreOmitted(t *testing.T) { + tests := []struct { + name string + cacheWriteDetails string + }{ + {name: "absent", cacheWriteDetails: `"cached_tokens": 0`}, + {name: "null", cacheWriteDetails: `"cached_tokens": 0, "cache_write_tokens": null`}, + {name: "string", cacheWriteDetails: `"cached_tokens": 0, "cache_write_tokens": "300"`}, + {name: "fractional", cacheWriteDetails: `"cached_tokens": 0, "cache_write_tokens": 1.5`}, + {name: "negative", cacheWriteDetails: `"cached_tokens": 0, "cache_write_tokens": -1`}, + {name: "overflow", cacheWriteDetails: `"cached_tokens": 0, "cache_write_tokens": 9223372036854775808`}, + } - sr, sw := schema.Pipe[*model.AgenticCallbackOutput](1) - reader := sr.Copy(1)[0] - sender := newCallbackSender(sw, &model.AgenticConfig{}) - sender.sendResponse(resp, nil) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp := mustUnmarshalResponse(t, fmt.Sprintf(`{ + "id": "resp_1", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 100, + "input_tokens_details": {%s}, + "output_tokens": 20, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 120 + } + }`, tt.cacheWriteDetails)) - out, err := reader.Recv() - if err != nil { - t.Fatal(err) - } - if tokens, ok := GetCacheWriteTokens(out.Message); !ok || tokens != 300 { - t.Fatalf("expected (300, true), got (%d, %v)", tokens, ok) + msg, err := toOutputMessage(resp, &model.Options{}) + if err != nil { + t.Fatal(err) + } + if tokens, ok := GetCacheWriteTokens(msg); ok || tokens != 0 { + t.Fatalf("expected (0, false), got (%d, %v)", tokens, ok) + } + }) } } diff --git a/components/model/agenticopenai/responses_event_convertor.go b/components/model/agenticopenai/responses_event_convertor.go index d4ea3f29c..b2b8849d7 100644 --- a/components/model/agenticopenai/responses_event_convertor.go +++ b/components/model/agenticopenai/responses_event_convertor.go @@ -219,10 +219,6 @@ func newCallbackSender(sw *schema.StreamWriter[*model.AgenticCallbackOutput], co } } -func (s *callbackSender) sendMeta(meta *schema.AgenticResponseMeta, err error) { - s.send(meta, nil, nil, err) -} - func (s *callbackSender) sendResponse(resp *responses.Response, err error) { if resp == nil { s.send(nil, nil, nil, err) diff --git a/components/model/agenticopenai/responses_event_convertor_test.go b/components/model/agenticopenai/responses_event_convertor_test.go index c494ffe1d..d7f49d40f 100644 --- a/components/model/agenticopenai/responses_event_convertor_test.go +++ b/components/model/agenticopenai/responses_event_convertor_test.go @@ -764,20 +764,6 @@ func TestNewCallbackSender(t *testing.T) { assert.Equal(t, config, s.config) } -func TestCallbackSenderSendMeta(t *testing.T) { - sr, sw := schema.Pipe[*model.AgenticCallbackOutput](8) - r := sr.Copy(1)[0] - s := newCallbackSender(sw, &model.AgenticConfig{}) - - meta := &schema.AgenticResponseMeta{} - s.sendMeta(meta, nil) - - out, err := r.Recv() - assert.NoError(t, err) - assert.NotNil(t, out) - assert.NotNil(t, out.Message.ResponseMeta) -} - func TestCallbackSenderSendBlock(t *testing.T) { sr, sw := schema.Pipe[*model.AgenticCallbackOutput](8) r := sr.Copy(1)[0] @@ -798,7 +784,7 @@ func TestCallbackSenderSendError(t *testing.T) { s := newCallbackSender(sw, &model.AgenticConfig{}) s.errHeader = "test error" - s.sendMeta(nil, errors.New("error")) + s.sendResponse(nil, errors.New("error")) _, err := r.Recv() assert.Error(t, err) diff --git a/components/model/agenticopenai/responses_model_test.go b/components/model/agenticopenai/responses_model_test.go index 06903ae2e..7c05b4c5f 100644 --- a/components/model/agenticopenai/responses_model_test.go +++ b/components/model/agenticopenai/responses_model_test.go @@ -124,13 +124,22 @@ func TestModelStream(t *testing.T) { mockey.Mock((*responses.ResponseService).NewStreaming).Return(mockStream).Build() - // Mock AsAny to return a completed event + completedResponse := mustUnmarshalResponse(t, `{ + "id": "resp_1", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 400, + "input_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 300}, + "output_tokens": 20, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 420 + } + }`) + + // Mock AsAny to return a completed event with cache-write usage. mockey.Mock(responses.ResponseStreamEventUnion.AsAny).Return(responses.ResponseCompletedEvent{ - Response: responses.Response{ - Output: []responses.ResponseOutputItemUnion{ - {Type: "message", ID: "m1", Status: "completed"}, - }, - }, + Response: *completedResponse, }).Build() s, err := m.Stream(ctx, input) @@ -138,8 +147,13 @@ func TestModelStream(t *testing.T) { assert.NotNil(t, s) defer s.Close() - // The stream should eventually close without errors - // We just verify it was created successfully + chunk, err := s.Recv() + assert.NoError(t, err) + if assert.NotNil(t, chunk) { + tokens, ok := GetCacheWriteTokens(chunk) + assert.True(t, ok) + assert.Equal(t, 300, tokens) + } }) mockey.PatchConvey("genRequest error", func() {