chore: update spanKind and attributes for tokens (#296)

* chore: update spanKind

* chore: added additional token semantic conventions

Author: poshinchen

src/strands/telemetry/tracer.py

@@ -77,9 +77,6 @@ class Tracer:
 
     When the OTEL_EXPORTER_OTLP_ENDPOINT environment variable is set, traces
     are sent to the OTLP endpoint.
-
-    When the STRANDS_OTEL_ENABLE_CONSOLE_EXPORT environment variable is set,
-    traces are printed to the console.
     """
 
     def __init__(
@@ -103,13 +100,15 @@ def _start_span(
         span_name: str,
         parent_span: Optional[Span] = None,
         attributes: Optional[Dict[str, AttributeValue]] = None,
+        span_kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
     ) -> Optional[Span]:
         """Generic helper method to start a span with common attributes.
 
         Args:
             span_name: Name of the span to create
             parent_span: Optional parent span to link this span to
             attributes: Dictionary of attributes to set on the span
+            span_kind: enum of OptenTelemetry SpanKind
 
         Returns:
             The created span, or None if tracing is not enabled
@@ -118,7 +117,7 @@ def _start_span(
             return None
 
         context = trace_api.set_span_in_context(parent_span) if parent_span else None
-        span = self.tracer.start_span(name=span_name, context=context)
+        span = self.tracer.start_span(name=span_name, context=context, kind=span_kind)
 
         # Set start time as a common attribute
         span.set_attribute("gen_ai.event.start_time", datetime.now(timezone.utc).isoformat())
@@ -219,7 +218,7 @@ def start_model_invoke_span(
         """
         attributes: Dict[str, AttributeValue] = {
             "gen_ai.system": "strands-agents",
-            "agent.name": agent_name,
+            "gen_ai.operation.name": "chat",
             "gen_ai.agent.name": agent_name,
             "gen_ai.prompt": serialize(messages),
         }
@@ -230,7 +229,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))})
 
-        return self._start_span("Model invoke", parent_span, attributes)
+        return self._start_span("Model invoke", parent_span, attributes, span_kind=trace_api.SpanKind.CLIENT)
 
     def end_model_invoke_span(
         self, span: Span, message: Message, usage: Usage, error: Optional[Exception] = None
@@ -246,7 +245,9 @@ def end_model_invoke_span(
         attributes: Dict[str, AttributeValue] = {
             "gen_ai.completion": serialize(message["content"]),
             "gen_ai.usage.prompt_tokens": usage["inputTokens"],
+            "gen_ai.usage.input_tokens": usage["inputTokens"],
             "gen_ai.usage.completion_tokens": usage["outputTokens"],
+            "gen_ai.usage.output_tokens": usage["outputTokens"],
             "gen_ai.usage.total_tokens": usage["totalTokens"],
         }
 
@@ -265,6 +266,7 @@ def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None
         """
         attributes: Dict[str, AttributeValue] = {
             "gen_ai.prompt": serialize(tool),
+            "gen_ai.system": "strands-agents",
             "tool.name": tool["name"],
             "tool.id": tool["toolUseId"],
             "tool.parameters": serialize(tool["input"]),
@@ -274,7 +276,7 @@ def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None
         attributes.update(kwargs)
 
         span_name = f"Tool: {tool['name']}"
-        return self._start_span(span_name, parent_span, attributes)
+        return self._start_span(span_name, parent_span, attributes, span_kind=trace_api.SpanKind.INTERNAL)
 
     def end_tool_call_span(
         self, span: Span, tool_result: Optional[ToolResult], error: Optional[Exception] = None
@@ -335,7 +337,7 @@ def start_event_loop_cycle_span(
         attributes.update({k: v for k, v in kwargs.items() if isinstance(v, (str, int, float, bool))})
 
         span_name = f"Cycle {event_loop_cycle_id}"
-        return self._start_span(span_name, parent_span, attributes)
+        return self._start_span(span_name, parent_span, attributes, span_kind=trace_api.SpanKind.INTERNAL)
 
     def end_event_loop_cycle_span(
         self,
@@ -405,7 +407,7 @@ def start_agent_span(
         # Add additional kwargs as attributes
         attributes.update({k: v for k, v in kwargs.items() if isinstance(v, (str, int, float, bool))})
 
-        return self._start_span(agent_name, attributes=attributes)
+        return self._start_span(agent_name, attributes=attributes, span_kind=trace_api.SpanKind.CLIENT)
 
     def end_agent_span(
         self,
@@ -436,6 +438,8 @@ def end_agent_span(
                     {
                         "gen_ai.usage.prompt_tokens": accumulated_usage["inputTokens"],
                         "gen_ai.usage.completion_tokens": accumulated_usage["outputTokens"],
+                        "gen_ai.usage.input_tokens": accumulated_usage["inputTokens"],
+                        "gen_ai.usage.output_tokens": accumulated_usage["outputTokens"],
                         "gen_ai.usage.total_tokens": accumulated_usage["totalTokens"],
                     }
                 )

tests/strands/telemetry/test_tracer.py

@@ -5,6 +5,7 @@
 
 import pytest
 from opentelemetry.trace import (
+    SpanKind,
     StatusCode,  # type: ignore
 )
 
@@ -75,7 +76,7 @@ def test_start_span(mock_tracer):
 
         span = tracer._start_span("test_span", attributes={"key": "value"})
 
-        mock_tracer.start_span.assert_called_once_with(name="test_span", context=None)
+        mock_tracer.start_span.assert_called_once_with(name="test_span", context=None, kind=SpanKind.INTERNAL)
         mock_span.set_attribute.assert_any_call("key", "value")
         assert span is not None
 
@@ -151,8 +152,10 @@ def test_start_model_invoke_span(mock_tracer):
 
         mock_tracer.start_span.assert_called_once()
         assert mock_tracer.start_span.call_args[1]["name"] == "Model invoke"
+        assert mock_tracer.start_span.call_args[1]["kind"] == SpanKind.CLIENT
         mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
-        mock_span.set_attribute.assert_any_call("agent.name", "TestAgent")
+        mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "chat")
+        mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "TestAgent")
         mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id)
         assert span is not None
 
@@ -167,7 +170,9 @@ def test_end_model_invoke_span(mock_span):
 
     mock_span.set_attribute.assert_any_call("gen_ai.completion", json.dumps(message["content"]))
     mock_span.set_attribute.assert_any_call("gen_ai.usage.prompt_tokens", 10)
+    mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 10)
     mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 20)
+    mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 20)
     mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 30)
     mock_span.set_status.assert_called_once_with(StatusCode.OK)
     mock_span.end.assert_called_once()
@@ -294,7 +299,9 @@ def test_end_agent_span(mock_span):
 
     mock_span.set_attribute.assert_any_call("gen_ai.completion", "Agent response")
     mock_span.set_attribute.assert_any_call("gen_ai.usage.prompt_tokens", 50)
+    mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 50)
     mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 100)
+    mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 100)
     mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150)
     mock_span.set_status.assert_called_once_with(StatusCode.OK)
     mock_span.end.assert_called_once()