From 953debae406689a97103fad03819e06e4233fe88 Mon Sep 17 00:00:00 2001 From: sara tadayon Date: Mon, 23 Feb 2026 11:09:18 -0700 Subject: [PATCH] feat: expose usage and metrics in AfterModelCallEvent --- src/strands/event_loop/event_loop.py | 3 +- src/strands/hooks/events.py | 5 ++ tests/fixtures/mocked_model_provider.py | 6 +++ .../strands/agent/hooks/test_agent_events.py | 48 +++++++++++++++++++ tests/strands/agent/test_agent_hooks.py | 8 ++++ tests/strands/agent/test_retry.py | 6 +++ tests/strands/event_loop/test_event_loop.py | 11 ++++- 7 files changed, 84 insertions(+), 3 deletions(-) diff --git a/src/strands/event_loop/event_loop.py b/src/strands/event_loop/event_loop.py index 3113ddb79..eb15bff99 100644 --- a/src/strands/event_loop/event_loop.py +++ b/src/strands/event_loop/event_loop.py @@ -346,8 +346,7 @@ async def _handle_model_execution( agent=agent, invocation_state=invocation_state, stop_response=AfterModelCallEvent.ModelStopResponse( - stop_reason=stop_reason, - message=message, + stop_reason=stop_reason, message=message, usage=usage, metrics=metrics ), ) diff --git a/src/strands/hooks/events.py b/src/strands/hooks/events.py index 8d3e5d280..0fd811c7b 100644 --- a/src/strands/hooks/events.py +++ b/src/strands/hooks/events.py @@ -13,6 +13,7 @@ from ..agent.agent_result import AgentResult from ..types.content import Message, Messages +from ..types.event_loop import Metrics, Usage from ..types.interrupt import _Interruptible from ..types.streaming import StopReason from ..types.tools import AgentTool, ToolResult, ToolUse @@ -269,10 +270,14 @@ class ModelStopResponse: Attributes: stop_reason: The reason the model stopped generating. message: The generated message from the model. + usage: Token usage information for model interactions. + metrics: Performance metrics for model interactions. """ message: Message stop_reason: StopReason + usage: Usage | None = None + metrics: Metrics | None = None invocation_state: dict[str, Any] = field(default_factory=dict) stop_response: ModelStopResponse | None = None diff --git a/tests/fixtures/mocked_model_provider.py b/tests/fixtures/mocked_model_provider.py index f1c5cae77..a0a15a865 100644 --- a/tests/fixtures/mocked_model_provider.py +++ b/tests/fixtures/mocked_model_provider.py @@ -106,3 +106,9 @@ def map_agent_message_to_events(self, agent_message: Message | RedactionMessage) yield {"contentBlockStop": {}} yield {"messageStop": {"stopReason": stop_reason}} + yield { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } diff --git a/tests/strands/agent/hooks/test_agent_events.py b/tests/strands/agent/hooks/test_agent_events.py index 02c367ccc..7755f7b10 100644 --- a/tests/strands/agent/hooks/test_agent_events.py +++ b/tests/strands/agent/hooks/test_agent_events.py @@ -140,6 +140,14 @@ async def test_stream_e2e_success(alist): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "tool_use"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, { "message": { "content": [ @@ -198,6 +206,14 @@ async def test_stream_e2e_success(alist): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "tool_use"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, { "message": { "content": [ @@ -256,6 +272,14 @@ async def test_stream_e2e_success(alist): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "tool_use"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, { "message": { "content": [ @@ -307,6 +331,14 @@ async def test_stream_e2e_success(alist): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "end_turn"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, {"message": {"content": [{"text": "I invoked the tools!"}], "role": "assistant"}}, { "result": AgentResult( @@ -371,6 +403,14 @@ async def test_stream_e2e_throttle_and_redact(alist, mock_sleep): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "guardrail_intervened"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, {"message": {"content": [{"text": "INPUT BLOCKED!"}], "role": "assistant"}}, { "result": AgentResult( @@ -435,6 +475,14 @@ async def test_stream_e2e_reasoning_redacted_content(alist): }, {"event": {"contentBlockStop": {}}}, {"event": {"messageStop": {"stopReason": "end_turn"}}}, + { + "event": { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + } + }, { "message": { "content": [ diff --git a/tests/strands/agent/test_agent_hooks.py b/tests/strands/agent/test_agent_hooks.py index 4397b9628..6c3712a03 100644 --- a/tests/strands/agent/test_agent_hooks.py +++ b/tests/strands/agent/test_agent_hooks.py @@ -175,6 +175,8 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, mock_model, tool_u "role": "assistant", }, stop_reason="tool_use", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), exception=None, ) @@ -201,6 +203,8 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, mock_model, tool_u stop_response=AfterModelCallEvent.ModelStopResponse( message=mock_model.agent_responses[1], stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), exception=None, ) @@ -248,6 +252,8 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, mock_m "role": "assistant", }, stop_reason="tool_use", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), exception=None, ) @@ -274,6 +280,8 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, mock_m stop_response=AfterModelCallEvent.ModelStopResponse( message=mock_model.agent_responses[1], stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), exception=None, ) diff --git a/tests/strands/agent/test_retry.py b/tests/strands/agent/test_retry.py index 830c1b5b8..45d6c7b22 100644 --- a/tests/strands/agent/test_retry.py +++ b/tests/strands/agent/test_retry.py @@ -170,6 +170,8 @@ async def test_model_retry_strategy_no_retry_on_success(): stop_response=AfterModelCallEvent.ModelStopResponse( message={"role": "assistant", "content": [{"text": "Success"}]}, stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), ) @@ -203,6 +205,8 @@ async def test_model_retry_strategy_reset_on_success(mock_sleep): stop_response=AfterModelCallEvent.ModelStopResponse( message={"role": "assistant", "content": [{"text": "Success"}]}, stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), ) await strategy._handle_after_model_call(event2) @@ -281,6 +285,8 @@ async def test_model_retry_strategy_backwards_compatible_event_cleared_on_succes stop_response=AfterModelCallEvent.ModelStopResponse( message={"role": "assistant", "content": [{"text": "Success"}]}, stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), ) diff --git a/tests/strands/event_loop/test_event_loop.py b/tests/strands/event_loop/test_event_loop.py index 8c6155e20..e2cc751aa 100644 --- a/tests/strands/event_loop/test_event_loop.py +++ b/tests/strands/event_loop/test_event_loop.py @@ -840,6 +840,12 @@ async def test_event_loop_cycle_exception_model_hooks(mock_sleep, agent, model, [ {"contentBlockDelta": {"delta": {"text": "test text"}}}, {"contentBlockStop": {}}, + { + "metadata": { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 100}, + } + }, ] ), ] @@ -878,7 +884,10 @@ async def test_event_loop_cycle_exception_model_hooks(mock_sleep, agent, model, agent=agent, invocation_state=ANY, stop_response=AfterModelCallEvent.ModelStopResponse( - message={"content": [{"text": "test text"}], "role": "assistant"}, stop_reason="end_turn" + message={"content": [{"text": "test text"}], "role": "assistant"}, + stop_reason="end_turn", + usage={"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + metrics={"latencyMs": 100}, ), exception=None, )