Skip to content

Commit 2131d2d

Browse files
committed
feat: added agent_invocations
1 parent 0c640e8 commit 2131d2d

File tree

3 files changed

+195
-23
lines changed

3 files changed

+195
-23
lines changed

src/strands/agent/agent.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -562,6 +562,8 @@ async def stream_async(
562562
"""
563563
self._interrupt_state.resume(prompt)
564564

565+
self.event_loop_metrics.reset_usage_metrics()
566+
565567
merged_state = {}
566568
if kwargs:
567569
warnings.warn("`**kwargs` parameter is deprecating, use `invocation_state` instead.", stacklevel=2)

src/strands/telemetry/metrics.py

Lines changed: 107 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,34 @@ def add_call(
151151
metrics_client.tool_error_count.add(1, attributes=attributes)
152152

153153

154+
@dataclass
155+
class EventLoopCycleMetric:
156+
"""Aggregated metrics for a single event loop cycle.
157+
158+
Attributes:
159+
event_loop_cycle_id: Current eventLoop cycle id.
160+
usage: Total token usage for the entire cycle (succeeded model invocation, excluding tool invocations).
161+
"""
162+
163+
event_loop_cycle_id: str
164+
usage: Usage
165+
166+
167+
@dataclass
168+
class AgentInvocation:
169+
"""Metrics for a single agent invocation.
170+
171+
AgentInvocation contains all the event loop cycles and accumulated token usage for that invocation.
172+
173+
Attributes:
174+
cycles: List of event loop cycles that occurred during this invocation.
175+
usage: Accumulated token usage for this invocation across all cycles.
176+
"""
177+
178+
cycles: list[EventLoopCycleMetric] = field(default_factory=list)
179+
usage: Usage = field(default_factory=lambda: Usage(inputTokens=0, outputTokens=0, totalTokens=0))
180+
181+
154182
@dataclass
155183
class EventLoopMetrics:
156184
"""Aggregated metrics for an event loop's execution.
@@ -159,15 +187,17 @@ class EventLoopMetrics:
159187
cycle_count: Number of event loop cycles executed.
160188
tool_metrics: Metrics for each tool used, keyed by tool name.
161189
cycle_durations: List of durations for each cycle in seconds.
190+
agent_invocations: Agent invocation metrics containing cycles and usage data.
162191
traces: List of execution traces.
163-
accumulated_usage: Accumulated token usage across all model invocations.
192+
accumulated_usage: Accumulated token usage across all model invocations (across all requests).
164193
accumulated_metrics: Accumulated performance metrics across all model invocations.
165194
"""
166195

167196
cycle_count: int = 0
168-
tool_metrics: Dict[str, ToolMetrics] = field(default_factory=dict)
169-
cycle_durations: List[float] = field(default_factory=list)
170-
traces: List[Trace] = field(default_factory=list)
197+
tool_metrics: dict[str, ToolMetrics] = field(default_factory=dict)
198+
cycle_durations: list[float] = field(default_factory=list)
199+
agent_invocations: list[AgentInvocation] = field(default_factory=list)
200+
traces: list[Trace] = field(default_factory=list)
171201
accumulated_usage: Usage = field(default_factory=lambda: Usage(inputTokens=0, outputTokens=0, totalTokens=0))
172202
accumulated_metrics: Metrics = field(default_factory=lambda: Metrics(latencyMs=0))
173203

@@ -176,14 +206,23 @@ def _metrics_client(self) -> "MetricsClient":
176206
"""Get the singleton MetricsClient instance."""
177207
return MetricsClient()
178208

209+
@property
210+
def current_agent_invocation(self) -> Optional[AgentInvocation]:
211+
"""Get the current (most recent) agent invocation.
212+
213+
Returns:
214+
The most recent AgentInvocation, or None if no invocations exist.
215+
"""
216+
return self.agent_invocations[-1] if self.agent_invocations else None
217+
179218
def start_cycle(
180219
self,
181-
attributes: Optional[Dict[str, Any]] = None,
220+
attributes: Dict[str, Any],
182221
) -> Tuple[float, Trace]:
183222
"""Start a new event loop cycle and create a trace for it.
184223
185224
Args:
186-
attributes: attributes of the metrics.
225+
attributes: attributes of the metrics, including event_loop_cycle_id.
187226
188227
Returns:
189228
A tuple containing the start time and the cycle trace object.
@@ -194,6 +233,18 @@ def start_cycle(
194233
start_time = time.time()
195234
cycle_trace = Trace(f"Cycle {self.cycle_count}", start_time=start_time)
196235
self.traces.append(cycle_trace)
236+
237+
# Ensure there's at least one agent invocation
238+
if not self.agent_invocations:
239+
self.agent_invocations.append(AgentInvocation())
240+
241+
self.agent_invocations[-1].cycles.append(
242+
EventLoopCycleMetric(
243+
event_loop_cycle_id=attributes["event_loop_cycle_id"],
244+
usage=Usage(inputTokens=0, outputTokens=0, totalTokens=0),
245+
)
246+
)
247+
197248
return start_time, cycle_trace
198249

199250
def end_cycle(self, start_time: float, cycle_trace: Trace, attributes: Optional[Dict[str, Any]] = None) -> None:
@@ -252,32 +303,57 @@ def add_tool_usage(
252303
)
253304
tool_trace.end()
254305

306+
def _accumulate_usage(self, target: Usage, source: Usage) -> None:
307+
"""Helper method to accumulate usage from source to target.
308+
309+
Args:
310+
target: The Usage object to accumulate into.
311+
source: The Usage object to accumulate from.
312+
"""
313+
target["inputTokens"] += source["inputTokens"]
314+
target["outputTokens"] += source["outputTokens"]
315+
target["totalTokens"] += source["totalTokens"]
316+
317+
if "cacheReadInputTokens" in source:
318+
target["cacheReadInputTokens"] = target.get("cacheReadInputTokens", 0) + source["cacheReadInputTokens"]
319+
320+
if "cacheWriteInputTokens" in source:
321+
target["cacheWriteInputTokens"] = target.get("cacheWriteInputTokens", 0) + source["cacheWriteInputTokens"]
322+
255323
def update_usage(self, usage: Usage) -> None:
256324
"""Update the accumulated token usage with new usage data.
257325
258326
Args:
259327
usage: The usage data to add to the accumulated totals.
260328
"""
329+
# Record metrics to OpenTelemetry
261330
self._metrics_client.event_loop_input_tokens.record(usage["inputTokens"])
262331
self._metrics_client.event_loop_output_tokens.record(usage["outputTokens"])
263-
self.accumulated_usage["inputTokens"] += usage["inputTokens"]
264-
self.accumulated_usage["outputTokens"] += usage["outputTokens"]
265-
self.accumulated_usage["totalTokens"] += usage["totalTokens"]
266332

267-
# Handle optional cached token metrics
333+
# Handle optional cached token metrics for OpenTelemetry
268334
if "cacheReadInputTokens" in usage:
269-
cache_read_tokens = usage["cacheReadInputTokens"]
270-
self._metrics_client.event_loop_cache_read_input_tokens.record(cache_read_tokens)
271-
self.accumulated_usage["cacheReadInputTokens"] = (
272-
self.accumulated_usage.get("cacheReadInputTokens", 0) + cache_read_tokens
273-
)
274-
335+
self._metrics_client.event_loop_cache_read_input_tokens.record(usage["cacheReadInputTokens"])
275336
if "cacheWriteInputTokens" in usage:
276-
cache_write_tokens = usage["cacheWriteInputTokens"]
277-
self._metrics_client.event_loop_cache_write_input_tokens.record(cache_write_tokens)
278-
self.accumulated_usage["cacheWriteInputTokens"] = (
279-
self.accumulated_usage.get("cacheWriteInputTokens", 0) + cache_write_tokens
280-
)
337+
self._metrics_client.event_loop_cache_write_input_tokens.record(usage["cacheWriteInputTokens"])
338+
339+
self._accumulate_usage(self.accumulated_usage, usage)
340+
341+
if not self.agent_invocations:
342+
self.agent_invocations.append(AgentInvocation())
343+
344+
self._accumulate_usage(self.agent_invocations[-1].usage, usage)
345+
346+
if self.agent_invocations[-1].cycles:
347+
current_cycle = self.agent_invocations[-1].cycles[-1]
348+
self._accumulate_usage(current_cycle.usage, usage)
349+
350+
def reset_usage_metrics(self) -> None:
351+
"""Start a new agent invocation by creating a new AgentInvocation.
352+
353+
This should be called at the start of a new request to begin tracking
354+
a new agent invocation with fresh usage and cycle data.
355+
"""
356+
self.agent_invocations.append(AgentInvocation())
281357

282358
def update_metrics(self, metrics: Metrics) -> None:
283359
"""Update the accumulated performance metrics with new metrics data.
@@ -322,6 +398,16 @@ def get_summary(self) -> Dict[str, Any]:
322398
"traces": [trace.to_dict() for trace in self.traces],
323399
"accumulated_usage": self.accumulated_usage,
324400
"accumulated_metrics": self.accumulated_metrics,
401+
"agent_invocations": [
402+
{
403+
"usage": invocation.usage,
404+
"cycles": [
405+
{"event_loop_cycle_id": cycle.event_loop_cycle_id, "usage": cycle.usage}
406+
for cycle in invocation.cycles
407+
],
408+
}
409+
for invocation in self.agent_invocations
410+
],
325411
}
326412
return summary
327413

tests/strands/telemetry/test_metrics.py

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,12 @@ def test_tool_metrics_add_call(success, tool, tool_metrics, mock_get_meter_provi
240240
@unittest.mock.patch.object(strands.telemetry.metrics.uuid, "uuid4")
241241
def test_event_loop_metrics_start_cycle(mock_uuid4, mock_time, event_loop_metrics, mock_get_meter_provider):
242242
mock_time.return_value = 1
243-
mock_uuid4.return_value = "i1"
243+
mock_event_loop_cycle_id = "i1"
244+
mock_uuid4.return_value = mock_event_loop_cycle_id
244245

245-
tru_start_time, tru_cycle_trace = event_loop_metrics.start_cycle()
246+
tru_start_time, tru_cycle_trace = event_loop_metrics.start_cycle(
247+
attributes={"event_loop_cycle_id": mock_event_loop_cycle_id}
248+
)
246249
exp_start_time, exp_cycle_trace = 1, strands.telemetry.metrics.Trace("Cycle 1")
247250

248251
tru_attrs = {"cycle_count": event_loop_metrics.cycle_count, "traces": event_loop_metrics.traces}
@@ -256,6 +259,13 @@ def test_event_loop_metrics_start_cycle(mock_uuid4, mock_time, event_loop_metric
256259
and tru_attrs == exp_attrs
257260
)
258261

262+
assert len(event_loop_metrics.agent_invocations) == 1
263+
assert len(event_loop_metrics.agent_invocations[0].cycles) == 1
264+
assert event_loop_metrics.agent_invocations[0].cycles[0].event_loop_cycle_id == "i1"
265+
assert event_loop_metrics.agent_invocations[0].cycles[0].usage["inputTokens"] == 0
266+
assert event_loop_metrics.agent_invocations[0].cycles[0].usage["outputTokens"] == 0
267+
assert event_loop_metrics.agent_invocations[0].cycles[0].usage["totalTokens"] == 0
268+
259269

260270
@unittest.mock.patch.object(strands.telemetry.metrics.time, "time")
261271
def test_event_loop_metrics_end_cycle(mock_time, trace, event_loop_metrics, mock_get_meter_provider):
@@ -324,13 +334,23 @@ def test_event_loop_metrics_add_tool_usage(mock_time, trace, tool, event_loop_me
324334

325335

326336
def test_event_loop_metrics_update_usage(usage, event_loop_metrics, mock_get_meter_provider):
337+
event_loop_metrics.start_cycle(attributes={"event_loop_cycle_id": "test-cycle"})
338+
327339
for _ in range(3):
328340
event_loop_metrics.update_usage(usage)
329341

330342
tru_usage = event_loop_metrics.accumulated_usage
331343
exp_usage = Usage(inputTokens=3, outputTokens=6, totalTokens=9, cacheWriteInputTokens=6)
332344

333345
assert tru_usage == exp_usage
346+
347+
assert event_loop_metrics.current_agent_invocation.usage == exp_usage
348+
349+
assert len(event_loop_metrics.agent_invocations) == 1
350+
assert len(event_loop_metrics.agent_invocations[0].cycles) == 1
351+
assert event_loop_metrics.agent_invocations[0].cycles[0].event_loop_cycle_id == "test-cycle"
352+
assert event_loop_metrics.agent_invocations[0].cycles[0].usage == exp_usage
353+
334354
mock_get_meter_provider.return_value.get_meter.assert_called()
335355
metrics_client = event_loop_metrics._metrics_client
336356
metrics_client.event_loop_input_tokens.record.assert_called()
@@ -370,6 +390,7 @@ def test_event_loop_metrics_get_summary(trace, tool, event_loop_metrics, mock_ge
370390
"outputTokens": 0,
371391
"totalTokens": 0,
372392
},
393+
"agent_invocations": [],
373394
"average_cycle_time": 0,
374395
"tool_usage": {
375396
"tool1": {
@@ -476,3 +497,66 @@ def test_use_ProxyMeter_if_no_global_meter_provider():
476497

477498
# Verify it's using a _ProxyMeter
478499
assert isinstance(metrics_client.meter, _ProxyMeter)
500+
501+
502+
def test_current_agent_invocation_property(usage, event_loop_metrics, mock_get_meter_provider):
503+
"""Test the current_agent_invocation property getter"""
504+
# Initially, no invocations exist
505+
assert event_loop_metrics.current_agent_invocation is None
506+
507+
event_loop_metrics.start_cycle(attributes={"event_loop_cycle_id": "cycle-1"})
508+
event_loop_metrics.update_usage(usage)
509+
510+
# Current_agent_invocation should return the first invocation
511+
current = event_loop_metrics.current_agent_invocation
512+
assert current is not None
513+
assert current.usage["inputTokens"] == 1
514+
assert len(current.cycles) == 1
515+
516+
event_loop_metrics.reset_usage_metrics()
517+
event_loop_metrics.start_cycle(attributes={"event_loop_cycle_id": "cycle-2"})
518+
usage2 = Usage(inputTokens=10, outputTokens=20, totalTokens=30)
519+
event_loop_metrics.update_usage(usage2)
520+
521+
# Should return the second invocation
522+
current = event_loop_metrics.current_agent_invocation
523+
assert current is not None
524+
assert current.usage["inputTokens"] == 10
525+
assert len(current.cycles) == 1
526+
527+
assert len(event_loop_metrics.agent_invocations) == 2
528+
529+
assert current is event_loop_metrics.agent_invocations[-1]
530+
531+
532+
def test_reset_usage_metrics(usage, event_loop_metrics, mock_get_meter_provider):
533+
"""Test that reset_usage_metrics creates a new agent invocation but preserves accumulated_usage"""
534+
# Add some usage across multiple cycles in first invocation
535+
event_loop_metrics.start_cycle(attributes={"event_loop_cycle_id": "cycle-1"})
536+
event_loop_metrics.update_usage(usage)
537+
538+
event_loop_metrics.start_cycle(attributes={"event_loop_cycle_id": "cycle-2"})
539+
usage2 = Usage(inputTokens=10, outputTokens=20, totalTokens=30)
540+
event_loop_metrics.update_usage(usage2)
541+
542+
assert len(event_loop_metrics.agent_invocations) == 1
543+
assert event_loop_metrics.current_agent_invocation.usage["inputTokens"] == 11
544+
assert len(event_loop_metrics.current_agent_invocation.cycles) == 2
545+
assert event_loop_metrics.accumulated_usage["inputTokens"] == 11
546+
547+
# Reset - creates a new invocation
548+
event_loop_metrics.reset_usage_metrics()
549+
550+
assert len(event_loop_metrics.agent_invocations) == 2
551+
552+
assert event_loop_metrics.current_agent_invocation.usage["inputTokens"] == 0
553+
assert event_loop_metrics.current_agent_invocation.usage["outputTokens"] == 0
554+
assert event_loop_metrics.current_agent_invocation.usage["totalTokens"] == 0
555+
assert len(event_loop_metrics.current_agent_invocation.cycles) == 0
556+
557+
# Verify first invocation data is preserved
558+
assert event_loop_metrics.agent_invocations[0].usage["inputTokens"] == 11
559+
assert len(event_loop_metrics.agent_invocations[0].cycles) == 2
560+
561+
# Verify accumulated_usage is NOT cleared
562+
assert event_loop_metrics.accumulated_usage["inputTokens"] == 11

0 commit comments

Comments
 (0)