Skip to content

chore: update attributes for semantic conventions #182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 56 additions & 28 deletions src/strands/telemetry/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.resources import Resource # type: ignore[attr-defined]
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor
from opentelemetry.trace import StatusCode
Expand Down Expand Up @@ -237,6 +237,21 @@ def _set_attributes(self, span: trace.Span, attributes: Dict[str, AttributeValue
for key, value in attributes.items():
span.set_attribute(key, value)

def _add_event(
self, span: Optional[trace.Span], event_name: str, event_attributes: Dict[str, AttributeValue]
) -> None:
"""Add an event with attributes to a span.

Args:
span: The span to add the event to
event_name: Name of the event
event_attributes: Dictionary of attributes to set on the event
"""
if not span:
return

span.add_event(event_name, event_attributes)

def _end_span(
self,
span: trace.Span,
Expand Down Expand Up @@ -314,18 +329,20 @@ def start_model_invoke_span(
"""
attributes: Dict[str, AttributeValue] = {
"gen_ai.system": "strands-agents",
"agent.name": agent_name,
"gen_ai.operation.name": "invoke_model",
"gen_ai.agent.name": agent_name,
"gen_ai.prompt": serialize(messages),
}

if model_id:
attributes["gen_ai.request.model"] = model_id

# 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)
span = self._start_span("Model invoke", parent_span, attributes)
self._add_event(
span, "gen_ai.user.message", {"gen_ai.system": "strands-agents", "content": serialize(messages)}
)
return span

def end_model_invoke_span(
self, span: trace.Span, message: Message, usage: Usage, error: Optional[Exception] = None
Expand All @@ -339,13 +356,17 @@ def end_model_invoke_span(
error: Optional exception if the model call failed.
"""
attributes: Dict[str, AttributeValue] = {
"gen_ai.completion": serialize(message["content"]),
"gen_ai.usage.prompt_tokens": usage["inputTokens"],
"gen_ai.usage.completion_tokens": usage["outputTokens"],
"gen_ai.usage.usage.output_tokens": usage["outputTokens"],

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this double usage redundant?

"gen_ai.usage.input_tokens": usage["inputTokens"],
"gen_ai.usage.total_tokens": usage["totalTokens"],
}

self._end_span(span, attributes, error)
self._add_event(
span,
"gen_ai.assistant.message",
event_attributes={"gen_ai.system": "strands-agents", "content": serialize(message["content"])},
)

def start_tool_call_span(
self, tool: ToolUse, parent_span: Optional[trace.Span] = None, **kwargs: Any
Expand All @@ -361,10 +382,9 @@ def start_tool_call_span(
The created span, or None if tracing is not enabled.
"""
attributes: Dict[str, AttributeValue] = {
"gen_ai.prompt": serialize(tool),
"tool.name": tool["name"],
"tool.id": tool["toolUseId"],
"tool.parameters": serialize(tool["input"]),
"gen_ai.tool.name": tool["name"],
"gen_ai.tool.id": tool["toolUseId"],
"gen_ai.tool.parameters": serialize(tool),
}

# Add additional kwargs as attributes
Expand All @@ -391,12 +411,19 @@ def end_tool_call_span(
tool_result_content_json = serialize(tool_result.get("content"))
attributes.update(
{
"tool.result": tool_result_content_json,
"gen_ai.completion": tool_result_content_json,
"gen_ai.tool.id": tool_result.get("toolUseId", ""),
"tool.status": status_str,
}
)

self._add_event(
span,
"gen_ai.tool.message",
event_attributes={
"gen_ai.system": "strands-agents",
"tool.status": status_str,
"content": tool_result_content_json,
},
)
self._end_span(span, attributes, error)

def start_event_loop_cycle_span(
Expand All @@ -421,7 +448,6 @@ def start_event_loop_cycle_span(
parent_span = parent_span if parent_span else event_loop_kwargs.get("event_loop_parent_span")

attributes: Dict[str, AttributeValue] = {
"gen_ai.prompt": serialize(messages),
"event_loop.cycle_id": event_loop_cycle_id,
}

Expand All @@ -432,7 +458,9 @@ 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)
span = self._start_span(span_name, parent_span, attributes)
self._add_event(span, "gen_ai.eventloop.cycle", event_attributes={"content": serialize(messages)})
return span

def end_event_loop_cycle_span(
self,
Expand All @@ -450,13 +478,15 @@ def end_event_loop_cycle_span(
error: Optional exception if the cycle failed.
"""
attributes: Dict[str, AttributeValue] = {
"gen_ai.completion": serialize(message["content"]),
"content": serialize(message["content"]),
}

if tool_result_message:
attributes["tool.result"] = serialize(tool_result_message["content"])
attributes["tool.content"] = serialize(tool_result_message["content"])

self._end_span(span, attributes, error)
self._add_event(span, "gen_ai.eventloop.message", attributes)

self._end_span(span, error=error)

def start_agent_span(
self,
Expand All @@ -482,7 +512,6 @@ def start_agent_span(
"""
attributes: Dict[str, AttributeValue] = {
"gen_ai.system": "strands-agents",
"agent.name": agent_name,
"gen_ai.agent.name": agent_name,
"gen_ai.prompt": prompt,
}
Expand All @@ -492,7 +521,6 @@ def start_agent_span(

if tools:
tools_json = serialize(tools)
attributes["agent.tools"] = tools_json
attributes["gen_ai.agent.tools"] = tools_json

# Add custom trace attributes if provided
Expand Down Expand Up @@ -521,18 +549,18 @@ def end_agent_span(
attributes: Dict[str, AttributeValue] = {}

if response:
attributes.update(
{
"gen_ai.completion": str(response),
}
self._add_event(
span,
"gen_ai.assistant.message",
event_attributes={"gen_ai.system": "strands-agents", "content": str(response)},
)

if hasattr(response, "metrics") and hasattr(response.metrics, "accumulated_usage"):
accumulated_usage = response.metrics.accumulated_usage
attributes.update(
{
"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.usage.output_tokens": accumulated_usage["outputTokens"],
"gen_ai.usage.total_tokens": accumulated_usage["totalTokens"],
}
)
Expand Down
66 changes: 48 additions & 18 deletions tests/strands/telemetry/test_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,26 @@ def test_start_span(mock_tracer):
assert span is not None


def test_add_event(mock_span):
"""Test adding an event to a span."""
tracer = Tracer()

event_name = "test_event"
event_attributes = {"key1": "value1", "key2": "value2"}

tracer._add_event(mock_span, event_name, event_attributes)

mock_span.add_event.assert_called_once_with(event_name, event_attributes)


def test_add_event_with_none_span():
"""Test adding an event to a None span."""
tracer = Tracer()

# This should not raise an exception
tracer._add_event(None, "test_event", {"key": "value"})


def test_set_attributes(mock_span):
"""Test setting attributes on a span."""
tracer = Tracer()
Expand Down Expand Up @@ -274,10 +294,12 @@ def test_start_model_invoke_span(mock_tracer):

span = tracer.start_model_invoke_span(agent_name="TestAgent", messages=messages, model_id=model_id)

mock_tracer.start_span.assert_called_once()
mock_span.add_event.assert_called_once_with(
"gen_ai.user.message", {"gen_ai.system": "strands-agents", "content": serialize(messages)}
)
assert mock_tracer.start_span.call_args[1]["name"] == "Model invoke"
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.agent.name", "TestAgent")
mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id)
assert span is not None

Expand All @@ -290,9 +312,11 @@ def test_end_model_invoke_span(mock_span):

tracer.end_model_invoke_span(mock_span, message, usage)

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.completion_tokens", 20)
mock_span.add_event.assert_called_once_with(
"gen_ai.assistant.message", {"gen_ai.system": "strands-agents", "content": serialize(message["content"])}
)
mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 10)
mock_span.set_attribute.assert_any_call("gen_ai.usage.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()
Expand All @@ -313,12 +337,12 @@ def test_start_tool_call_span(mock_tracer):

mock_tracer.start_span.assert_called_once()
assert mock_tracer.start_span.call_args[1]["name"] == "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.tool.id", "123")
mock_span.set_attribute.assert_any_call(
"gen_ai.prompt", json.dumps({"name": "test-tool", "toolUseId": "123", "input": {"param": "value"}})
"gen_ai.tool.parameters", json.dumps({"name": "test-tool", "toolUseId": "123", "input": {"param": "value"}})
)
mock_span.set_attribute.assert_any_call("tool.name", "test-tool")
mock_span.set_attribute.assert_any_call("tool.id", "123")
mock_span.set_attribute.assert_any_call("tool.parameters", json.dumps({"param": "value"}))
assert span is not None


Expand All @@ -329,10 +353,13 @@ def test_end_tool_call_span(mock_span):

tracer.end_tool_call_span(mock_span, tool_result)

mock_span.set_attribute.assert_any_call("tool.result", json.dumps(tool_result.get("content")))
mock_span.set_attribute.assert_any_call("gen_ai.completion", json.dumps(tool_result.get("content")))
mock_span.set_attribute.assert_any_call("gen_ai.tool.id", "")
mock_span.set_attribute.assert_any_call("tool.status", "success")
mock_span.set_status.assert_called_once_with(StatusCode.OK)
mock_span.add_event.assert_called_once_with(
"gen_ai.tool.message",
{"gen_ai.system": "strands-agents", "tool.status": "success", "content": json.dumps(tool_result["content"])},
)
mock_span.end.assert_called_once()


Expand All @@ -352,8 +379,8 @@ def test_start_event_loop_cycle_span(mock_tracer):

mock_tracer.start_span.assert_called_once()
assert mock_tracer.start_span.call_args[1]["name"] == "Cycle cycle-123"
mock_span.set_attribute.assert_any_call("gen_ai.prompt", json.dumps(messages))
mock_span.set_attribute.assert_any_call("event_loop.cycle_id", "cycle-123")
mock_span.add_event.assert_called_once_with("gen_ai.eventloop.cycle", {"content": json.dumps(messages)})
assert span is not None


Expand All @@ -365,8 +392,10 @@ def test_end_event_loop_cycle_span(mock_span):

tracer.end_event_loop_cycle_span(mock_span, message, tool_result_message)

mock_span.set_attribute.assert_any_call("gen_ai.completion", json.dumps(message["content"]))
mock_span.set_attribute.assert_any_call("tool.result", json.dumps(tool_result_message["content"]))
mock_span.add_event.assert_called_once_with(
"gen_ai.eventloop.message",
{"content": json.dumps(message["content"]), "tool.content": json.dumps(tool_result_message["content"])},
)
mock_span.set_status.assert_called_once_with(StatusCode.OK)
mock_span.end.assert_called_once()

Expand Down Expand Up @@ -396,7 +425,6 @@ def test_start_agent_span(mock_tracer):
mock_tracer.start_span.assert_called_once()
assert mock_tracer.start_span.call_args[1]["name"] == "WeatherAgent"
mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents")
mock_span.set_attribute.assert_any_call("agent.name", "WeatherAgent")
mock_span.set_attribute.assert_any_call("gen_ai.prompt", prompt)
mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id)
mock_span.set_attribute.assert_any_call("custom_attr", "value")
Expand All @@ -417,10 +445,12 @@ def test_end_agent_span(mock_span):

tracer.end_agent_span(mock_span, mock_response)

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.completion_tokens", 100)
mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 50)
mock_span.set_attribute.assert_any_call("gen_ai.usage.usage.output_tokens", 100)
mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150)
mock_span.add_event.assert_any_call(
"gen_ai.assistant.message", {"gen_ai.system": "strands-agents", "content": "Agent response"}
)
mock_span.set_status.assert_called_once_with(StatusCode.OK)
mock_span.end.assert_called_once()

Expand Down