feat: added method for multiagent spans (#451)

Author: poshinchen

src/strands/multiagent/graph.py

@@ -21,7 +21,10 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Tuple, cast
 
+from opentelemetry import trace as trace_api
+
 from ..agent import Agent, AgentResult
+from ..telemetry import get_tracer
 from ..types.content import ContentBlock
 from ..types.event_loop import Metrics, Usage
 from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
@@ -249,6 +252,7 @@ def __init__(self, nodes: dict[str, GraphNode], edges: set[GraphEdge], entry_poi
         self.edges = edges
         self.entry_points = entry_points
         self.state = GraphState()
+        self.tracer = get_tracer()
 
     def execute(self, task: str | list[ContentBlock]) -> GraphResult:
         """Execute task synchronously."""
@@ -274,19 +278,20 @@ async def execute_async(self, task: str | list[ContentBlock]) -> GraphResult:
         )
 
         start_time = time.time()
-        try:
-            await self._execute_graph()
-            self.state.status = Status.COMPLETED
-            logger.debug("status=<%s> | graph execution completed", self.state.status)
-
-        except Exception:
-            logger.exception("graph execution failed")
-            self.state.status = Status.FAILED
-            raise
-        finally:
-            self.state.execution_time = round((time.time() - start_time) * 1000)
-
-        return self._build_result()
+        span = self.tracer.start_multiagent_span(task, "graph")
+        with trace_api.use_span(span, end_on_exit=True):
+            try:
+                await self._execute_graph()
+                self.state.status = Status.COMPLETED
+                logger.debug("status=<%s> | graph execution completed", self.state.status)
+
+            except Exception:
+                logger.exception("graph execution failed")
+                self.state.status = Status.FAILED
+                raise
+            finally:
+                self.state.execution_time = round((time.time() - start_time) * 1000)
+            return self._build_result()
 
     async def _execute_graph(self) -> None:
         """Unified execution flow with conditional routing."""

src/strands/telemetry/tracer.py

@@ -14,7 +14,7 @@
 from opentelemetry.trace import Span, StatusCode
 
 from ..agent.agent_result import AgentResult
-from ..types.content import Message, Messages
+from ..types.content import ContentBlock, Message, Messages
 from ..types.streaming import StopReason, Usage
 from ..types.tools import ToolResult, ToolUse
 from ..types.traces import AttributeValue
@@ -86,8 +86,6 @@ def __init__(
         """Initialize the tracer."""
         self.service_name = __name__
         self.tracer_provider: Optional[trace_api.TracerProvider] = None
-        self.tracer: Optional[trace_api.Tracer] = None
-
         self.tracer_provider = trace_api.get_tracer_provider()
         self.tracer = self.tracer_provider.get_tracer(self.service_name)
         ThreadingInstrumentor().instrument()
@@ -98,7 +96,7 @@ def _start_span(
         parent_span: Optional[Span] = None,
         attributes: Optional[Dict[str, AttributeValue]] = None,
         span_kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
-    ) -> Optional[Span]:
+    ) -> Span:
         """Generic helper method to start a span with common attributes.
 
         Args:
@@ -110,10 +108,13 @@ def _start_span(
         Returns:
             The created span, or None if tracing is not enabled
         """
-        if self.tracer is None:
-            return None
+        if not parent_span:
+            parent_span = trace_api.get_current_span()
+
+        context = None
+        if parent_span and parent_span.is_recording() and parent_span != trace_api.INVALID_SPAN:
+            context = trace_api.set_span_in_context(parent_span)
 
-        context = trace_api.set_span_in_context(parent_span) if parent_span else None
         span = self.tracer.start_span(name=span_name, context=context, kind=span_kind)
 
         # Set start time as a common attribute
@@ -235,7 +236,7 @@ def start_model_invoke_span(
         # Add additional kwargs as attributes
         attributes.update({k: v for k, v in kwargs.items() if isinstance(v, (str, int, float, bool))})
 
-        span = self._start_span("chat", parent_span, attributes, span_kind=trace_api.SpanKind.CLIENT)
+        span = self._start_span("chat", parent_span, attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
         for message in messages:
             self._add_event(
                 span,
@@ -293,8 +294,8 @@ def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None
         # Add additional kwargs as attributes
         attributes.update(kwargs)
 
-        span_name = f"Tool: {tool['name']}"
-        span = self._start_span(span_name, parent_span, attributes, span_kind=trace_api.SpanKind.INTERNAL)
+        span_name = f"execute_tool {tool['name']}"
+        span = self._start_span(span_name, parent_span, attributes=attributes, span_kind=trace_api.SpanKind.INTERNAL)
 
         self._add_event(
             span,
@@ -497,6 +498,41 @@ def end_agent_span(
 
         self._end_span(span, attributes, error)
 
+    def start_multiagent_span(
+        self,
+        task: str | list[ContentBlock],
+        instance: str,
+    ) -> Span:
+        """Start a new span for swarm invocation."""
+        attributes: Dict[str, AttributeValue] = {
+            "gen_ai.system": "strands-agents",
+            "gen_ai.agent.name": instance,
+            "gen_ai.operation.name": f"invoke_{instance}",
+        }
+
+        span = self._start_span(f"invoke_{instance}", attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
+        content = serialize(task) if isinstance(task, list) else task
+        self._add_event(
+            span,
+            "gen_ai.user.message",
+            event_attributes={"content": content},
+        )
+
+        return span
+
+    def end_swarm_span(
+        self,
+        span: Span,
+        result: Optional[str] = None,
+    ) -> None:
+        """End a swarm span with results."""
+        if result:
+            self._add_event(
+                span,
+                "gen_ai.choice",
+                event_attributes={"message": result},
+            )
+
 
 # Singleton instance for global access
 _tracer_instance = None

tests/strands/multiagent/test_graph.py

@@ -1,4 +1,4 @@
-from unittest.mock import AsyncMock, MagicMock, Mock
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 
@@ -101,6 +101,22 @@ def string_content_agent():
     return agent
 
 
+@pytest.fixture
+def mock_strands_tracer():
+    with patch("strands.multiagent.graph.get_tracer") as mock_get_tracer:
+        mock_tracer_instance = MagicMock()
+        mock_span = MagicMock()
+        mock_tracer_instance.start_multiagent_span.return_value = mock_span
+        mock_get_tracer.return_value = mock_tracer_instance
+        yield mock_tracer_instance
+
+
+@pytest.fixture
+def mock_use_span():
+    with patch("strands.multiagent.graph.trace_api.use_span") as mock_use_span:
+        yield mock_use_span
+
+
 @pytest.fixture
 def mock_graph(mock_agents, string_content_agent):
     """Create a graph for testing various scenarios."""
@@ -138,8 +154,9 @@ def always_false_condition(state: GraphState) -> bool:
 
 
 @pytest.mark.asyncio
-async def test_graph_execution(mock_graph, mock_agents, string_content_agent):
+async def test_graph_execution(mock_strands_tracer, mock_use_span, mock_graph, mock_agents, string_content_agent):
     """Test comprehensive graph execution with diverse nodes and conditional edges."""
+
     # Test graph structure
     assert len(mock_graph.nodes) == 8
     assert len(mock_graph.edges) == 8
@@ -214,9 +231,12 @@ async def test_graph_execution(mock_graph, mock_agents, string_content_agent):
     assert len(result.entry_points) == 1
     assert result.entry_points[0].node_id == "start_agent"
 
+    mock_strands_tracer.start_multiagent_span.assert_called()
+    mock_use_span.assert_called_once()
+
 
 @pytest.mark.asyncio
-async def test_graph_unsupported_node_type():
+async def test_graph_unsupported_node_type(mock_strands_tracer, mock_use_span):
     """Test unsupported executor type error handling."""
 
     class UnsupportedExecutor:
@@ -229,9 +249,12 @@ class UnsupportedExecutor:
     with pytest.raises(ValueError, match="Node 'unsupported_node' of type.*is not supported"):
         await graph.execute_async("test task")
 
+    mock_strands_tracer.start_multiagent_span.assert_called()
+    mock_use_span.assert_called_once()
+
 
 @pytest.mark.asyncio
-async def test_graph_execution_with_failures():
+async def test_graph_execution_with_failures(mock_strands_tracer, mock_use_span):
     """Test graph execution error handling and failure propagation."""
     failing_agent = Mock(spec=Agent)
     failing_agent.name = "failing_agent"
@@ -261,10 +284,12 @@ async def mock_stream_failure(*args, **kwargs):
     assert graph.state.status == Status.FAILED
     assert any(node.node_id == "fail_node" for node in graph.state.failed_nodes)
     assert len(graph.state.completed_nodes) == 0
+    mock_strands_tracer.start_multiagent_span.assert_called()
+    mock_use_span.assert_called_once()
 
 
 @pytest.mark.asyncio
-async def test_graph_edge_cases():
+async def test_graph_edge_cases(mock_strands_tracer, mock_use_span):
     """Test specific edge cases for coverage."""
     # Test entry node execution without dependencies
     entry_agent = create_mock_agent("entry_agent", "Entry response")
@@ -278,6 +303,8 @@ async def test_graph_edge_cases():
     # Verify entry node was called with original task
     entry_agent.stream_async.assert_called_once_with([{"text": "Original task"}])
     assert result.status == Status.COMPLETED
+    mock_strands_tracer.start_multiagent_span.assert_called()
+    mock_use_span.assert_called_once()
 
 
 def test_graph_builder_validation():
@@ -415,7 +442,7 @@ def test_condition(state):
     assert len(node.dependencies) == 0
 
 
-def test_graph_synchronous_execution(mock_agents):
+def test_graph_synchronous_execution(mock_strands_tracer, mock_use_span, mock_agents):
     """Test synchronous graph execution using execute method."""
     builder = GraphBuilder()
     builder.add_node(mock_agents["start_agent"], "start_agent")
@@ -444,3 +471,6 @@ def test_graph_synchronous_execution(mock_agents):
     # Verify return type is GraphResult
     assert isinstance(result, GraphResult)
     assert isinstance(result, MultiAgentResult)
+
+    mock_strands_tracer.start_multiagent_span.assert_called()
+    mock_use_span.assert_called_once()

tests/strands/telemetry/test_tracer.py

@@ -10,6 +10,7 @@
 )
 
 from strands.telemetry.tracer import JSONEncoder, Tracer, get_tracer, serialize
+from strands.types.content import ContentBlock
 from strands.types.streaming import StopReason, Usage
 
 
@@ -198,7 +199,91 @@ def test_start_tool_call_span(mock_tracer):
         span = tracer.start_tool_call_span(tool)
 
         mock_tracer.start_span.assert_called_once()
-        assert mock_tracer.start_span.call_args[1]["name"] == "Tool: test-tool"
+        assert mock_tracer.start_span.call_args[1]["name"] == "execute_tool test-tool"
+        mock_span.set_attribute.assert_any_call("gen_ai.tool.name", "test-tool")
+        mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
+        mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "execute_tool")
+        mock_span.set_attribute.assert_any_call("gen_ai.tool.call.id", "123")
+        mock_span.add_event.assert_any_call(
+            "gen_ai.tool.message", attributes={"role": "tool", "content": json.dumps({"param": "value"}), "id": "123"}
+        )
+        assert span is not None
+
+
+def test_start_swarm_call_span_with_string_task(mock_tracer):
+    """Test starting a swarm call span with task as string."""
+    with mock.patch("strands.telemetry.tracer.trace_api.get_tracer", return_value=mock_tracer):
+        tracer = Tracer()
+        tracer.tracer = mock_tracer
+
+        mock_span = mock.MagicMock()
+        mock_tracer.start_span.return_value = mock_span
+
+        task = "Design foo bar"
+
+        span = tracer.start_multiagent_span(task, "swarm")
+
+        mock_tracer.start_span.assert_called_once()
+        assert mock_tracer.start_span.call_args[1]["name"] == "invoke_swarm"
+        mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
+        mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "swarm")
+        mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "invoke_swarm")
+        mock_span.add_event.assert_any_call("gen_ai.user.message", attributes={"content": "Design foo bar"})
+        assert span is not None
+
+
+def test_start_swarm_span_with_contentblock_task(mock_tracer):
+    """Test starting a swarm call span with task as list of contentBlock."""
+    with mock.patch("strands.telemetry.tracer.trace_api.get_tracer", return_value=mock_tracer):
+        tracer = Tracer()
+        tracer.tracer = mock_tracer
+
+        mock_span = mock.MagicMock()
+        mock_tracer.start_span.return_value = mock_span
+
+        task = [ContentBlock(text="Original Task: foo bar")]
+
+        span = tracer.start_multiagent_span(task, "swarm")
+
+        mock_tracer.start_span.assert_called_once()
+        assert mock_tracer.start_span.call_args[1]["name"] == "invoke_swarm"
+        mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
+        mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "swarm")
+        mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "invoke_swarm")
+        mock_span.add_event.assert_any_call(
+            "gen_ai.user.message", attributes={"content": '[{"text": "Original Task: foo bar"}]'}
+        )
+        assert span is not None
+
+
+def test_end_swarm_span(mock_span):
+    """Test ending a tool call span."""
+    tracer = Tracer()
+    swarm_final_reuslt = "foo bar bar"
+
+    tracer.end_swarm_span(mock_span, swarm_final_reuslt)
+
+    mock_span.add_event.assert_called_with(
+        "gen_ai.choice",
+        attributes={"message": "foo bar bar"},
+    )
+
+
+def test_start_graph_call_span(mock_tracer):
+    """Test starting a graph call span."""
+    with mock.patch("strands.telemetry.tracer.trace_api.get_tracer", return_value=mock_tracer):
+        tracer = Tracer()
+        tracer.tracer = mock_tracer
+
+        mock_span = mock.MagicMock()
+        mock_tracer.start_span.return_value = mock_span
+
+        tool = {"name": "test-tool", "toolUseId": "123", "input": {"param": "value"}}
+
+        span = tracer.start_tool_call_span(tool)
+
+        mock_tracer.start_span.assert_called_once()
+        assert mock_tracer.start_span.call_args[1]["name"] == "execute_tool test-tool"
         mock_span.set_attribute.assert_any_call("gen_ai.tool.name", "test-tool")
         mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
         mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "execute_tool")