This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new a0b0c7a7 [integration][anthropic][python] Keep token usage on
non-tool_use responses (#959)
a0b0c7a7 is described below
commit a0b0c7a7e2ecf4d4179d153c84d7d55ed92c161e
Author: Edson <[email protected]>
AuthorDate: Sat Aug 8 10:36:47 2026 -0400
[integration][anthropic][python] Keep token usage on non-tool_use responses
(#959)
---
.../chat_models/anthropic/anthropic_chat_model.py | 1 +
.../tests/test_anthropic_response_parsing.py | 43 ++++++++++++++++++++++
2 files changed, 44 insertions(+)
diff --git
a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
index afcd3f12..f816a289 100644
---
a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
@@ -239,6 +239,7 @@ class AnthropicChatModelConnection(BaseChatModelConnection):
return ChatMessage(
role=MessageRole(message.role),
content=text,
+ extra_args=extra_args,
)
@override
diff --git
a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py
b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py
index 89a1d5bf..f11889e4 100644
---
a/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py
+++
b/python/flink_agents/integrations/chat_models/anthropic/tests/test_anthropic_response_parsing.py
@@ -94,3 +94,46 @@ def test_plain_text_response() -> None:
[ChatMessage(role=MessageRole.USER, content="hi")]
)
assert response.content == "Hello!"
+
+
+def test_plain_text_response_keeps_token_usage() -> None:
+ # Token usage must survive on non-tool_use responses too: the common
+ # end_turn path previously dropped extra_args entirely, so promptTokens /
+ # completionTokens never reached the token metrics recording.
+ message = Message(
+ id="m",
+ model="claude",
+ role="assistant",
+ type="message",
+ stop_reason="end_turn",
+ content=[TextBlock(type="text", text="Hello!")],
+ usage=Usage(input_tokens=7, output_tokens=3),
+ )
+ response = _connection_returning(message).chat(
+ [ChatMessage(role=MessageRole.USER, content="hi")],
+ model="claude-sonnet-4-5",
+ )
+ assert response.extra_args["model_name"] == "claude-sonnet-4-5"
+ assert response.extra_args["promptTokens"] == 7
+ assert response.extra_args["completionTokens"] == 3
+
+
+def test_tool_use_response_keeps_token_usage() -> None:
+ # Regression guard for the tool_use path, which already carried usage.
+ message = Message(
+ id="m",
+ model="claude",
+ role="assistant",
+ type="message",
+ stop_reason="tool_use",
+ content=[
+ ToolUseBlock(type="tool_use", id="t1", name="add", input={"a": 1,
"b": 2})
+ ],
+ usage=Usage(input_tokens=7, output_tokens=3),
+ )
+ response = _connection_returning(message).chat(
+ [ChatMessage(role=MessageRole.USER, content="add 1 and 2")],
+ model="claude-sonnet-4-5",
+ )
+ assert response.extra_args["promptTokens"] == 7
+ assert response.extra_args["completionTokens"] == 3
