Skip to content

Attach usage and metrics to MessageAddedEvent. #1575

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
esoler-sage wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Attach usage and metrics to MessageAddedEvent. #1575

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
4 changes: 3 additions & 1 deletion src/strands/event_loop/event_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -403,7 +403,9 @@ async def _handle_model_execution(

# Add the response message to the conversation
agent.messages.append(message)
await agent.hooks.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=message))
await agent.hooks.invoke_callbacks_async(
MessageAddedEvent(agent=agent, message=message, usage=usage, metrics=metrics)
)

# Update metrics
agent.event_loop_metrics.update_usage(usage)
Expand Down
3 changes: 3 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 @@ -105,6 +106,8 @@ class MessageAddedEvent(HookEvent):
"""

message: Message
usage: Usage | None = None
metrics: Metrics | None = None


@dataclass
Expand Down
28 changes: 24 additions & 4 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 @@ -178,7 +178,12 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, mock_model, tool_u
exception=None,
)

assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[1])
assert next(events) == MessageAddedEvent(
agent=agent,
message=agent.messages[1],
usage={"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
metrics={"latencyMs": 0, "timeToFirstByteMs": 0},
)
assert next(events) == BeforeToolCallEvent(
agent=agent,
selected_tool=agent_tool,
Expand All @@ -202,7 +207,12 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, mock_model, tool_u
),
exception=None,
)
assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[3])
assert next(events) == MessageAddedEvent(
agent=agent,
message=agent.messages[3],
usage={"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
metrics={"latencyMs": 0, "timeToFirstByteMs": 0},
)

assert next(events) == AfterInvocationEvent(agent=agent, result=result)

Expand Down Expand Up @@ -248,7 +258,12 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, mock_m
exception=None,
)

assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[1])
assert next(events) == MessageAddedEvent(
agent=agent,
message=agent.messages[1],
usage={"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
metrics={"latencyMs": 0, "timeToFirstByteMs": 0},
)
assert next(events) == BeforeToolCallEvent(
agent=agent,
selected_tool=agent_tool,
Expand All @@ -272,7 +287,12 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, mock_m
),
exception=None,
)
assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[3])
assert next(events) == MessageAddedEvent(
agent=agent,
message=agent.messages[3],
usage={"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
metrics={"latencyMs": 0, "timeToFirstByteMs": 0},
)

assert next(events) == AfterInvocationEvent(agent=agent, result=result)

Expand Down
5 changes: 4 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 @@ -884,7 +884,10 @@ async def test_event_loop_cycle_exception_model_hooks(mock_sleep, agent, model,

# Final message
assert next(events) == MessageAddedEvent(
agent=agent, message={"content": [{"text": "test text"}], "role": "assistant"}
agent=agent,
message={"content": [{"text": "test text"}], "role": "assistant"},
usage={"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
metrics={"latencyMs": 0, "timeToFirstByteMs": 0},
)


Expand Down
Loading