Skip to content

Commit 02b710c

Browse files
authored
feat: add Swarm tracing (strands-agents#461)
1 parent 1edd81a commit 02b710c

File tree

2 files changed

+50
-20
lines changed

2 files changed

+50
-20
lines changed

src/strands/multiagent/swarm.py

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,11 @@
2121
from dataclasses import dataclass, field
2222
from typing import Any, Callable, Tuple
2323

24+
from opentelemetry import trace as trace_api
25+
2426
from ..agent import Agent, AgentResult
2527
from ..agent.state import AgentState
28+
from ..telemetry import get_tracer
2629
from ..tools.decorator import tool
2730
from ..types.content import ContentBlock, Messages
2831
from ..types.event_loop import Metrics, Usage
@@ -229,6 +232,7 @@ def __init__(
229232
task="",
230233
completion_status=Status.PENDING,
231234
)
235+
self.tracer = get_tracer()
232236

233237
self._setup_swarm(nodes)
234238
self._inject_swarm_tools()
@@ -257,24 +261,26 @@ async def invoke_async(self, task: str | list[ContentBlock], **kwargs: Any) -> S
257261
)
258262

259263
start_time = time.time()
260-
try:
261-
logger.debug("current_node=<%s> | starting swarm execution with node", self.state.current_node.node_id)
262-
logger.debug(
263-
"max_handoffs=<%d>, max_iterations=<%d>, timeout=<%s>s | swarm execution config",
264-
self.max_handoffs,
265-
self.max_iterations,
266-
self.execution_timeout,
267-
)
264+
span = self.tracer.start_multiagent_span(task, "swarm")
265+
with trace_api.use_span(span, end_on_exit=True):
266+
try:
267+
logger.debug("current_node=<%s> | starting swarm execution with node", self.state.current_node.node_id)
268+
logger.debug(
269+
"max_handoffs=<%d>, max_iterations=<%d>, timeout=<%s>s | swarm execution config",
270+
self.max_handoffs,
271+
self.max_iterations,
272+
self.execution_timeout,
273+
)
268274

269-
await self._execute_swarm()
270-
except Exception:
271-
logger.exception("swarm execution failed")
272-
self.state.completion_status = Status.FAILED
273-
raise
274-
finally:
275-
self.state.execution_time = round((time.time() - start_time) * 1000)
275+
await self._execute_swarm()
276+
except Exception:
277+
logger.exception("swarm execution failed")
278+
self.state.completion_status = Status.FAILED
279+
raise
280+
finally:
281+
self.state.execution_time = round((time.time() - start_time) * 1000)
276282

277-
return self._build_result()
283+
return self._build_result()
278284

279285
def _setup_swarm(self, nodes: list[Agent]) -> None:
280286
"""Initialize swarm configuration."""

tests/strands/multiagent/test_swarm.py

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import math
22
import time
3-
from unittest.mock import MagicMock, Mock
3+
from unittest.mock import MagicMock, Mock, patch
44

55
import pytest
66

@@ -93,6 +93,22 @@ def mock_swarm(mock_agents):
9393
return swarm
9494

9595

96+
@pytest.fixture
97+
def mock_strands_tracer():
98+
with patch("strands.multiagent.swarm.get_tracer") as mock_get_tracer:
99+
mock_tracer_instance = MagicMock()
100+
mock_span = MagicMock()
101+
mock_tracer_instance.start_multiagent_span.return_value = mock_span
102+
mock_get_tracer.return_value = mock_tracer_instance
103+
yield mock_tracer_instance
104+
105+
106+
@pytest.fixture
107+
def mock_use_span():
108+
with patch("strands.multiagent.swarm.trace_api.use_span") as mock_use_span:
109+
yield mock_use_span
110+
111+
96112
def test_swarm_structure_and_nodes(mock_swarm, mock_agents):
97113
"""Test swarm structure and SwarmNode properties."""
98114
# Test swarm structure
@@ -214,7 +230,7 @@ def test_swarm_state_should_continue(mock_swarm):
214230

215231

216232
@pytest.mark.asyncio
217-
async def test_swarm_execution_async(mock_swarm, mock_agents):
233+
async def test_swarm_execution_async(mock_strands_tracer, mock_use_span, mock_swarm, mock_agents):
218234
"""Test asynchronous swarm execution."""
219235
# Execute swarm
220236
task = [ContentBlock(text="Analyze this task"), ContentBlock(text="Additional context")]
@@ -237,8 +253,11 @@ async def test_swarm_execution_async(mock_swarm, mock_agents):
237253
assert hasattr(result, "node_history")
238254
assert len(result.node_history) == 1
239255

256+
mock_strands_tracer.start_multiagent_span.assert_called()
257+
mock_use_span.assert_called_once()
240258

241-
def test_swarm_synchronous_execution(mock_agents):
259+
260+
def test_swarm_synchronous_execution(mock_strands_tracer, mock_use_span, mock_agents):
242261
"""Test synchronous swarm execution using __call__ method."""
243262
agents = list(mock_agents.values())
244263
swarm = Swarm(
@@ -279,6 +298,9 @@ def test_swarm_synchronous_execution(mock_agents):
279298
for node in swarm.nodes.values():
280299
node.executor.tool_registry.process_tools.assert_called()
281300

301+
mock_strands_tracer.start_multiagent_span.assert_called()
302+
mock_use_span.assert_called_once()
303+
282304

283305
def test_swarm_builder_validation(mock_agents):
284306
"""Test swarm builder validation and error handling."""
@@ -405,7 +427,7 @@ def test_swarm_tool_creation_and_execution():
405427
assert completion_result["status"] == "success"
406428

407429

408-
def test_swarm_failure_handling():
430+
def test_swarm_failure_handling(mock_strands_tracer, mock_use_span):
409431
"""Test swarm execution with agent failures."""
410432
# Test execution with agent failures
411433
failing_agent = create_mock_agent("failing_agent")
@@ -416,6 +438,8 @@ def test_swarm_failure_handling():
416438
# The swarm catches exceptions internally and sets status to FAILED
417439
result = failing_swarm("Test failure handling")
418440
assert result.status == Status.FAILED
441+
mock_strands_tracer.start_multiagent_span.assert_called()
442+
mock_use_span.assert_called_once()
419443

420444

421445
def test_swarm_metrics_handling():

0 commit comments

Comments
 (0)