chore: added swarm and graph spans

Author: poshinchen

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
@@ -113,7 +113,13 @@ def _start_span(
         if self.tracer is None:
             return None
 
-        context = trace_api.set_span_in_context(parent_span) if parent_span else 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)
+
         span = self.tracer.start_span(name=span_name, context=context, kind=span_kind)
 
         # Set start time as a common attribute
@@ -235,7 +241,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 +299,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 +503,63 @@ def end_agent_span(
 
         self._end_span(span, attributes, error)
 
+    def start_swarm_span(
+        self,
+        task: str | list[ContentBlock],
+    ) -> Optional[Span]:
+        """Start a new span for swarm invocation."""
+        attributes: Dict[str, AttributeValue] = {
+            "gen_ai.system": "strands-agents",
+            "gen_ai.agent.name": "swarm",
+            "gen_ai.operation.name": "invoke_swarm",
+        }
+
+        span = self._start_span("invoke_swarm", 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},
+            )
+
+    def start_graph_span(
+        self,
+        task: str,
+    ) -> Optional[Span]:
+        """Start a new span for graph invocation."""
+        attributes: Dict[str, AttributeValue] = {
+            "gen_ai.system": "strands-agents",
+            "gen_ai.agent.name": "graph",
+            "gen_ai.operation.name": "invoke_graph",
+        }
+
+        span = self._start_span("invoke_graph", attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
+
+        self._add_event(
+            span,
+            "gen_ai.user.message",
+            event_attributes={
+                "content": task,
+            },
+        )
+
+        return span
+
 
 # Singleton instance for global access
 _tracer_instance = None

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,116 @@ 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 start_swarm_span(
+        self,
+        task: str | list[ContentBlock],
+    ) -> Optional[Span]:
+        attributes: Dict[str, AttributeValue] = {
+            "gen_ai.system": "strands-agents",
+            "gen_ai.agent.name": "swarm",
+            "gen_ai.operation.name": "invoke_swarm",
+        }
+
+        span = self._start_span("invoke_swarm", attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
+
+        self._add_event(
+            span,
+            "gen_ai.user.message",
+            event_attributes={
+                "content": serialize(task),
+            },
+        )
+
+        return span
+"""
+
+
+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_swarm_span(task)
+
+        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_swarm_span(task)
+
+        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")