diff --git a/src/strands/telemetry/tracer.py b/src/strands/telemetry/tracer.py index 6ab33301a..80fb86c40 100644 --- a/src/strands/telemetry/tracer.py +++ b/src/strands/telemetry/tracer.py @@ -110,6 +110,17 @@ def _parse_semconv_opt_in(self) -> set[str]: opt_in_env = os.getenv("OTEL_SEMCONV_STABILITY_OPT_IN", "") return {value.strip() for value in opt_in_env.split(",")} + @property + def is_langfuse(self) -> bool: + """Check if Langfuse is configured as the OTLP endpoint. + + Returns: + True if Langfuse is the OTLP endpoint, False otherwise. + """ + return "langfuse" in os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "") or "langfuse" in os.getenv( + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "" + ) + def _start_span( self, span_name: str, @@ -142,23 +153,10 @@ def _start_span( # Add all provided attributes if attributes: - self._set_attributes(span, attributes) + span.set_attributes(attributes) return span - def _set_attributes(self, span: Span, attributes: dict[str, AttributeValue]) -> None: - """Set attributes on a span, handling different value types appropriately. - - Args: - span: The span to set attributes on - attributes: Dictionary of attributes to set - """ - if not span: - return - - for key, value in attributes.items(): - span.set_attribute(key, value) - def _add_optional_usage_and_metrics_attributes( self, attributes: dict[str, AttributeValue], usage: Usage, metrics: Metrics ) -> None: @@ -203,7 +201,7 @@ def _end_span( # Add any additional attributes if attributes: - self._set_attributes(span, attributes) + span.set_attributes(attributes) # Handle error if present if error: @@ -236,17 +234,24 @@ def end_span_with_error(self, span: Span, error_message: str, exception: Excepti error = exception or Exception(error_message) self._end_span(span, error=error) - def _add_event(self, span: Span | None, event_name: str, event_attributes: Attributes) -> None: + def _add_event( + self, span: Span | None, event_name: str, event_attributes: Attributes, to_span_attributes: bool = False + ) -> 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 + to_span_attributes: Add the attributes to span attributes """ if not span: return + # Add to span attribute since some backend can't read the events + if to_span_attributes and event_attributes: + span.set_attributes(event_attributes) + span.add_event(event_name, attributes=event_attributes) def _get_event_name_for_message(self, message: Message) -> str: @@ -358,6 +363,7 @@ def end_model_invoke_span( ] ), }, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -366,7 +372,7 @@ def end_model_invoke_span( event_attributes={"finish_reason": str(stop_reason), "message": serialize(message["content"])}, ) - self._set_attributes(span, attributes) + span.set_attributes(attributes) def start_tool_call_span( self, @@ -423,6 +429,7 @@ def start_tool_call_span( ] ) }, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -476,6 +483,7 @@ def end_tool_call_span(self, span: Span, tool_result: ToolResult | None, error: ] ) }, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -572,6 +580,7 @@ def end_event_loop_cycle_span( ] ) }, + to_span_attributes=self.is_langfuse, ) else: self._add_event(span, "gen_ai.choice", event_attributes=event_attributes) @@ -666,6 +675,7 @@ def end_agent_span( ] ) }, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -675,9 +685,7 @@ def end_agent_span( ) if hasattr(response, "metrics") and hasattr(response.metrics, "accumulated_usage"): - if "langfuse" in os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "") or "langfuse" in os.getenv( - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "" - ): + if self.is_langfuse: attributes.update({"langfuse.observation.type": "span"}) accumulated_usage = response.metrics.accumulated_usage attributes.update( @@ -736,6 +744,7 @@ def start_multiagent_span( span, "gen_ai.client.inference.operation.details", {"gen_ai.input.messages": serialize([{"role": "user", "parts": parts}])}, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -767,6 +776,7 @@ def end_swarm_span( ] ) }, + to_span_attributes=self.is_langfuse, ) else: self._add_event( @@ -816,7 +826,10 @@ def _add_event_messages(self, span: Span, messages: Messages) -> None: {"role": message["role"], "parts": self._map_content_blocks_to_otel_parts(message["content"])} ) self._add_event( - span, "gen_ai.client.inference.operation.details", {"gen_ai.input.messages": serialize(input_messages)} + span, + "gen_ai.client.inference.operation.details", + {"gen_ai.input.messages": serialize(input_messages)}, + to_span_attributes=self.is_langfuse, ) else: for message in messages: diff --git a/tests/strands/telemetry/test_tracer.py b/tests/strands/telemetry/test_tracer.py index 6ea605083..da7f010e2 100644 --- a/tests/strands/telemetry/test_tracer.py +++ b/tests/strands/telemetry/test_tracer.py @@ -79,22 +79,11 @@ 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, kind=SpanKind.INTERNAL) - mock_span.set_attribute.assert_any_call("key", "value") + # Check that set_attributes was called with the provided attributes + mock_span.set_attributes.assert_called_once_with({"key": "value"}) assert span is not None -def test_set_attributes(mock_span): - """Test setting attributes on a span.""" - tracer = Tracer() - attributes = {"str_attr": "value", "int_attr": 123, "bool_attr": True} - - tracer._set_attributes(mock_span, attributes) - - # Check that set_attribute was called for each attribute - calls = [mock.call(k, v) for k, v in attributes.items()] - mock_span.set_attribute.assert_has_calls(calls, any_order=True) - - def test_end_span_no_span(): """Test ending a span when span is None.""" tracer = Tracer() @@ -109,7 +98,8 @@ def test_end_span(mock_span): tracer._end_span(mock_span, attributes) - mock_span.set_attribute.assert_any_call("key", "value") + # Check that set_attributes was called with the provided attributes + mock_span.set_attributes.assert_called_once_with({"key": "value"}) mock_span.set_status.assert_called_once_with(StatusCode.OK) mock_span.end.assert_called_once() @@ -158,11 +148,14 @@ def test_start_model_invoke_span(mock_tracer): mock_tracer.start_span.assert_called_once() assert mock_tracer.start_span.call_args[1]["name"] == "chat" assert mock_tracer.start_span.call_args[1]["kind"] == SpanKind.INTERNAL - mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents") - mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "chat") - mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id) - mock_span.set_attribute.assert_any_call("custom_key", "custom_value") - mock_span.set_attribute.assert_any_call("user_id", "12345") + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "chat", + "gen_ai.system": "strands-agents", + "custom_key": "custom_value", + "user_id": "12345", + "gen_ai.request.model": model_id, + "agent_name": "TestAgent", + }) mock_span.add_event.assert_called_with( "gen_ai.user.message", attributes={"content": json.dumps(messages[0]["content"])} ) @@ -195,9 +188,13 @@ def test_start_model_invoke_span_latest_conventions(mock_tracer, monkeypatch): mock_tracer.start_span.assert_called_once() assert mock_tracer.start_span.call_args[1]["name"] == "chat" assert mock_tracer.start_span.call_args[1]["kind"] == SpanKind.INTERNAL - mock_span.set_attribute.assert_any_call("gen_ai.provider.name", "strands-agents") - mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "chat") - mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id) + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "chat", + "gen_ai.provider.name": "strands-agents", + "gen_ai.request.model": model_id, + "agent_name": "TestAgent", + }) mock_span.add_event.assert_called_with( "gen_ai.client.inference.operation.details", attributes={ @@ -235,13 +232,15 @@ def test_end_model_invoke_span(mock_span): tracer.end_model_invoke_span(mock_span, message, usage, metrics, stop_reason) - 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_attribute.assert_any_call("gen_ai.server.request.duration", 20) - mock_span.set_attribute.assert_any_call("gen_ai.server.time_to_first_token", 10) + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.completion_tokens": 20, + "gen_ai.usage.output_tokens": 20, + "gen_ai.usage.total_tokens": 30, + "gen_ai.server.time_to_first_token": 10, + "gen_ai.server.request.duration": 20, + }) mock_span.add_event.assert_called_with( "gen_ai.choice", attributes={"message": json.dumps(message["content"]), "finish_reason": "end_turn"}, @@ -260,13 +259,15 @@ def test_end_model_invoke_span_latest_conventions(mock_span, monkeypatch): tracer.end_model_invoke_span(mock_span, message, usage, metrics, stop_reason) - 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_attribute.assert_any_call("gen_ai.server.time_to_first_token", 10) - mock_span.set_attribute.assert_any_call("gen_ai.server.request.duration", 20) + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.completion_tokens": 20, + "gen_ai.usage.output_tokens": 20, + "gen_ai.usage.total_tokens": 30, + "gen_ai.server.time_to_first_token": 10, + "gen_ai.server.request.duration": 20, + }) mock_span.add_event.assert_called_with( "gen_ai.client.inference.operation.details", attributes={ @@ -299,12 +300,15 @@ def test_start_tool_call_span(mock_tracer): 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") - mock_span.set_attribute.assert_any_call("gen_ai.tool.call.id", "123") - mock_span.set_attribute.assert_any_call("session_id", "abc123") - mock_span.set_attribute.assert_any_call("environment", "production") + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.tool.name": "test-tool", + "gen_ai.system": "strands-agents", + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.call.id": "123", + "session_id": "abc123", + "environment": "production", + }) mock_span.add_event.assert_any_call( "gen_ai.tool.message", attributes={"role": "tool", "content": json.dumps({"param": "value"}), "id": "123"} ) @@ -327,10 +331,13 @@ def test_start_tool_call_span_latest_conventions(mock_tracer, monkeypatch): 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.provider.name", "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.set_attributes.assert_called_once_with({ + "gen_ai.tool.name": "test-tool", + "gen_ai.provider.name": "strands-agents", + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.call.id": "123", + }) mock_span.add_event.assert_called_with( "gen_ai.client.inference.operation.details", attributes={ @@ -370,11 +377,14 @@ def test_start_swarm_call_span_with_string_task(mock_tracer): 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.set_attribute.assert_any_call("workflow_id", "wf-789") - mock_span.set_attribute.assert_any_call("priority", "high") + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "invoke_swarm", + "gen_ai.system": "strands-agents", + "gen_ai.agent.name": "swarm", + "workflow_id": "wf-789", + "priority": "high", + }) mock_span.add_event.assert_any_call("gen_ai.user.message", attributes={"content": "Design foo bar"}) assert span is not None @@ -394,9 +404,12 @@ def test_start_swarm_span_with_contentblock_task(mock_tracer): 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.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "invoke_swarm", + "gen_ai.system": "strands-agents", + "gen_ai.agent.name": "swarm", + }) mock_span.add_event.assert_any_call( "gen_ai.user.message", attributes={"content": '[{"text": "Original Task: foo bar"}]'} ) @@ -447,9 +460,12 @@ def test_start_swarm_span_with_contentblock_task_latest_conventions(mock_tracer, 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.provider.name", "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.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "invoke_swarm", + "gen_ai.provider.name": "strands-agents", + "gen_ai.agent.name": "swarm", + }) mock_span.add_event.assert_any_call( "gen_ai.client.inference.operation.details", attributes={ @@ -512,10 +528,13 @@ def test_start_graph_call_span(mock_tracer): 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") - mock_span.set_attribute.assert_any_call("gen_ai.tool.call.id", "123") + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "execute_tool", + "gen_ai.system": "strands-agents", + "gen_ai.tool.name": "test-tool", + "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"} ) @@ -529,7 +548,7 @@ def test_end_tool_call_span(mock_span): tracer.end_tool_call_span(mock_span, tool_result) - mock_span.set_attribute.assert_any_call("gen_ai.tool.status", "success") + mock_span.set_attributes.assert_called_once_with({"gen_ai.tool.status": "success"}) mock_span.add_event.assert_called_with( "gen_ai.choice", attributes={"message": json.dumps(tool_result.get("content")), "id": ""}, @@ -546,7 +565,7 @@ def test_end_tool_call_span_latest_conventions(mock_span, monkeypatch): tracer.end_tool_call_span(mock_span, tool_result) - mock_span.set_attribute.assert_any_call("gen_ai.tool.status", "success") + mock_span.set_attributes.assert_called_once_with({"gen_ai.tool.status": "success"}) mock_span.add_event.assert_called_with( "gen_ai.client.inference.operation.details", attributes={ @@ -589,9 +608,12 @@ 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"] == "execute_event_loop_cycle" - mock_span.set_attribute.assert_any_call("event_loop.cycle_id", "cycle-123") - mock_span.set_attribute.assert_any_call("request_id", "req-456") - mock_span.set_attribute.assert_any_call("trace_level", "debug") + + mock_span.set_attributes.assert_called_once_with({ + "event_loop.cycle_id": "cycle-123", + "request_id": "req-456", + "trace_level": "debug", + }) mock_span.add_event.assert_any_call( "gen_ai.user.message", attributes={"content": json.dumps([{"text": "Hello"}])} ) @@ -615,7 +637,8 @@ def test_start_event_loop_cycle_span_latest_conventions(mock_tracer, monkeypatch mock_tracer.start_span.assert_called_once() assert mock_tracer.start_span.call_args[1]["name"] == "execute_event_loop_cycle" - mock_span.set_attribute.assert_any_call("event_loop.cycle_id", "cycle-123") + + mock_span.set_attributes.assert_called_once_with({"event_loop.cycle_id": "cycle-123"}) mock_span.add_event.assert_any_call( "gen_ai.client.inference.operation.details", attributes={ @@ -707,10 +730,15 @@ def test_start_agent_span(mock_tracer): mock_tracer.start_span.assert_called_once() assert mock_tracer.start_span.call_args[1]["name"] == "invoke_agent WeatherAgent" assert mock_tracer.start_span.call_args[1]["kind"] == SpanKind.INTERNAL - mock_span.set_attribute.assert_any_call("gen_ai.system", "strands-agents") - mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "WeatherAgent") - mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id) - mock_span.set_attribute.assert_any_call("custom_attr", "value") + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "invoke_agent", + "gen_ai.system": "strands-agents", + "gen_ai.agent.name": "WeatherAgent", + "gen_ai.request.model": model_id, + "gen_ai.agent.tools": json.dumps(tools), + "custom_attr": "value", + }) mock_span.add_event.assert_any_call("gen_ai.user.message", attributes={"content": json.dumps(content)}) assert span is not None @@ -740,10 +768,15 @@ def test_start_agent_span_latest_conventions(mock_tracer, monkeypatch): mock_tracer.start_span.assert_called_once() assert mock_tracer.start_span.call_args[1]["name"] == "invoke_agent WeatherAgent" - mock_span.set_attribute.assert_any_call("gen_ai.provider.name", "strands-agents") - mock_span.set_attribute.assert_any_call("gen_ai.agent.name", "WeatherAgent") - mock_span.set_attribute.assert_any_call("gen_ai.request.model", model_id) - mock_span.set_attribute.assert_any_call("custom_attr", "value") + + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.operation.name": "invoke_agent", + "gen_ai.provider.name": "strands-agents", + "gen_ai.agent.name": "WeatherAgent", + "gen_ai.request.model": model_id, + "gen_ai.agent.tools": json.dumps(tools), + "custom_attr": "value", + }) mock_span.add_event.assert_any_call( "gen_ai.client.inference.operation.details", attributes={ @@ -770,13 +803,17 @@ def test_end_agent_span(mock_span): tracer.end_agent_span(mock_span, mock_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_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 0) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 0) + mock_span.set_attributes.assert_called_once_with( + { + "gen_ai.usage.prompt_tokens": 50, + "gen_ai.usage.input_tokens": 50, + "gen_ai.usage.completion_tokens": 100, + "gen_ai.usage.output_tokens": 100, + "gen_ai.usage.total_tokens": 150, + "gen_ai.usage.cache_read_input_tokens": 0, + "gen_ai.usage.cache_write_input_tokens": 0, + } + ) mock_span.add_event.assert_any_call( "gen_ai.choice", attributes={"message": "Agent response", "finish_reason": "end_turn"}, @@ -800,13 +837,19 @@ def test_end_agent_span_with_langfuse_observation_type(mock_span, monkeypatch): mock_response.__str__ = mock.MagicMock(return_value="Agent response") tracer.end_agent_span(mock_span, mock_response) - mock_span.set_attribute.assert_any_call("langfuse.observation.type", "span") - 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.output_tokens", 100) - mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 0) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 0) + + mock_span.set_attributes.assert_called_once_with( + { + "langfuse.observation.type": "span", + "gen_ai.usage.prompt_tokens": 50, + "gen_ai.usage.input_tokens": 50, + "gen_ai.usage.completion_tokens": 100, + "gen_ai.usage.output_tokens": 100, + "gen_ai.usage.total_tokens": 150, + "gen_ai.usage.cache_read_input_tokens": 0, + "gen_ai.usage.cache_write_input_tokens": 0, + } + ) mock_span.add_event.assert_any_call( "gen_ai.choice", attributes={"message": "Agent response", "finish_reason": "end_turn"}, @@ -831,13 +874,17 @@ def test_end_agent_span_latest_conventions(mock_span, monkeypatch): tracer.end_agent_span(mock_span, mock_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_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 0) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 0) + mock_span.set_attributes.assert_called_once_with( + { + "gen_ai.usage.prompt_tokens": 50, + "gen_ai.usage.input_tokens": 50, + "gen_ai.usage.completion_tokens": 100, + "gen_ai.usage.output_tokens": 100, + "gen_ai.usage.total_tokens": 150, + "gen_ai.usage.cache_read_input_tokens": 0, + "gen_ai.usage.cache_write_input_tokens": 0, + } + ) mock_span.add_event.assert_called_with( "gen_ai.client.inference.operation.details", attributes={ @@ -872,15 +919,17 @@ def test_end_model_invoke_span_with_cache_metrics(mock_span): tracer.end_model_invoke_span(mock_span, message, usage, metrics, stop_reason) - 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_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 5) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 3) - mock_span.set_attribute.assert_any_call("gen_ai.server.request.duration", 10) - mock_span.set_attribute.assert_any_call("gen_ai.server.time_to_first_token", 5) + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.completion_tokens": 20, + "gen_ai.usage.output_tokens": 20, + "gen_ai.usage.total_tokens": 30, + "gen_ai.usage.cache_read_input_tokens": 5, + "gen_ai.usage.cache_write_input_tokens": 3, + "gen_ai.server.request.duration": 10, + "gen_ai.server.time_to_first_token": 5, + }) def test_end_agent_span_with_cache_metrics(mock_span): @@ -904,13 +953,15 @@ def test_end_agent_span_with_cache_metrics(mock_span): tracer.end_agent_span(mock_span, mock_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_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 25) - mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 10) + mock_span.set_attributes.assert_called_once_with({ + "gen_ai.usage.prompt_tokens": 50, + "gen_ai.usage.input_tokens": 50, + "gen_ai.usage.completion_tokens": 100, + "gen_ai.usage.output_tokens": 100, + "gen_ai.usage.total_tokens": 150, + "gen_ai.usage.cache_read_input_tokens": 25, + "gen_ai.usage.cache_write_input_tokens": 10, + }) mock_span.set_status.assert_called_once_with(StatusCode.OK) mock_span.end.assert_called_once() @@ -1444,3 +1495,80 @@ def test_start_agent_span_includes_tool_definitions_when_enabled(monkeypatch): ] expected_json = serialize(expected_tool_details) assert attributes["gen_ai.tool.definitions"] == expected_json + + +def test_end_model_invoke_span_langfuse_adds_attributes(mock_span, monkeypatch): + """Test that end_model_invoke_span adds attributes via set_attributes for Langfuse.""" + monkeypatch.setenv("OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://us.cloud.langfuse.com") + + tracer = Tracer() + message = {"role": "assistant", "content": [{"text": "Response"}]} + usage = Usage(inputTokens=10, outputTokens=20, totalTokens=30) + metrics = Metrics(latencyMs=20, timeToFirstByteMs=10) + stop_reason: StopReason = "end_turn" + + tracer.end_model_invoke_span(mock_span, message, usage, metrics, stop_reason) + + expected_output = serialize( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "Response"}], + "finish_reason": "end_turn", + } + ] + ) + + assert mock_span.set_attributes.call_count == 2 + mock_span.set_attributes.assert_any_call({"gen_ai.output.messages": expected_output}) + mock_span.set_attributes.assert_any_call({ + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.completion_tokens": 20, + "gen_ai.usage.output_tokens": 20, + "gen_ai.usage.total_tokens": 30, + "gen_ai.server.time_to_first_token": 10, + "gen_ai.server.request.duration": 20, + }) + + mock_span.add_event.assert_called_with( + "gen_ai.client.inference.operation.details", + attributes={"gen_ai.output.messages": expected_output}, + ) + + +def test_end_model_invoke_span_non_langfuse_no_extra_attributes(mock_span, monkeypatch): + """Test that end_model_invoke_span doesn't add extra attributes for non-Langfuse endpoints.""" + monkeypatch.setenv("OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://api.honeycomb.io") + + tracer = Tracer() + message = {"role": "assistant", "content": [{"text": "Response"}]} + usage = Usage(inputTokens=10, outputTokens=20, totalTokens=30) + metrics = Metrics(latencyMs=20, timeToFirstByteMs=10) + stop_reason: StopReason = "end_turn" + + tracer.end_model_invoke_span(mock_span, message, usage, metrics, stop_reason) + + # Verify that set_attribute was NOT called with gen_ai.output.messages + # (it should only be in the event, not as an attribute) + expected_output = serialize( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "Response"}], + "finish_reason": "end_turn", + } + ] + ) + + # Check that gen_ai.output.messages was not set as an attribute + set_attribute_calls = [call[0][0] for call in mock_span.set_attribute.call_args_list] + assert "gen_ai.output.messages" not in set_attribute_calls + + # But verify that add_event was still called + mock_span.add_event.assert_called_with( + "gen_ai.client.inference.operation.details", + attributes={"gen_ai.output.messages": expected_output}, + )