Skip to content

feat: expose usage and metrics in AfterModelCallEvent #1753

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
svtadayon wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: expose usage and metrics in AfterModelCallEvent #1753

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions src/strands/event_loop/event_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
),
)

Expand Down
5 changes: 5 additions & 0 deletions src/strands/hooks/events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions tests/fixtures/mocked_model_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
}
}
48 changes: 48 additions & 0 deletions tests/strands/agent/hooks/test_agent_events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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": [
Expand Down
8 changes: 8 additions & 0 deletions tests/strands/agent/test_agent_hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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,
)
Expand Down Expand Up @@ -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,
)
Expand All @@ -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,
)
Expand Down
6 changes: 6 additions & 0 deletions tests/strands/agent/test_retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
),
)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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},
),
)

Expand Down
11 changes: 10 additions & 1 deletion tests/strands/event_loop/test_event_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
}
},
]
),
]
Expand Down Expand Up @@ -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,
)
Expand Down