diff --git a/sentry_sdk/ai/__init__.py b/sentry_sdk/ai/__init__.py index fbcb9c061d..7f0d9f92f7 100644 --- a/sentry_sdk/ai/__init__.py +++ b/sentry_sdk/ai/__init__.py @@ -4,4 +4,5 @@ GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING, normalize_message_role, normalize_message_roles, + set_conversation_id, ) # noqa: F401 diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index a4ebe96d99..11d3a9be80 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -697,3 +697,11 @@ def truncate_and_annotate_messages( scope._gen_ai_original_message_count[span.span_id] = len(messages) return truncated_messages + + +def set_conversation_id(conversation_id: str) -> None: + """ + Set the conversation_id in the scope. + """ + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id(conversation_id) diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py index 6df26690c8..acd82a57ef 100644 --- a/sentry_sdk/scope.py +++ b/sentry_sdk/scope.py @@ -221,6 +221,7 @@ class Scope: "_breadcrumbs", "_n_breadcrumbs_truncated", "_gen_ai_original_message_count", + "_gen_ai_conversation_id", "_event_processors", "_error_processors", "_should_capture", @@ -303,6 +304,8 @@ def __copy__(self) -> "Scope": rv._attributes = self._attributes.copy() + rv._gen_ai_conversation_id = self._gen_ai_conversation_id + return rv @classmethod @@ -720,6 +723,8 @@ def clear(self) -> None: self._attributes: "Attributes" = {} + self._gen_ai_conversation_id: "Optional[str]" = None + @_attr_setter def level(self, value: "LogLevelStr") -> None: """ @@ -912,6 +917,26 @@ def remove_extra( """Removes a specific extra key.""" self._extras.pop(key, None) + def set_conversation_id(self, conversation_id: str) -> None: + """ + Sets the conversation ID for gen_ai spans. + + :param conversation_id: The conversation ID to set. + """ + self._gen_ai_conversation_id = conversation_id + + def get_conversation_id(self) -> "Optional[str]": + """ + Gets the conversation ID for gen_ai spans. + + :returns: The conversation ID, or None if not set. + """ + return self._gen_ai_conversation_id + + def remove_conversation_id(self) -> None: + """Removes the conversation ID.""" + self._gen_ai_conversation_id = None + def clear_breadcrumbs(self) -> None: """Clears breadcrumb buffer.""" self._breadcrumbs: "Deque[Breadcrumb]" = deque() @@ -1668,6 +1693,8 @@ def update_from_scope(self, scope: "Scope") -> None: self._gen_ai_original_message_count.update( scope._gen_ai_original_message_count ) + if scope._gen_ai_conversation_id: + self._gen_ai_conversation_id = scope._gen_ai_conversation_id if scope._span: self._span = scope._span if scope._attachments: diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index c4b38e4528..6dcc47d886 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -676,6 +676,17 @@ def finish( self.timestamp = datetime.now(timezone.utc) scope = scope or sentry_sdk.get_current_scope() + + # Copy conversation_id from scope to span data if this is an AI span + conversation_id = scope.get_conversation_id() + if conversation_id: + has_ai_op = "gen_ai.operation.name" in self._data + is_ai_span_op = self.op is not None and ( + self.op.startswith("ai.") or self.op.startswith("gen_ai.") + ) + if has_ai_op or is_ai_span_op: + self.set_data("gen_ai.conversation.id", conversation_id) + maybe_create_breadcrumbs_from_span(scope, self) return None diff --git a/tests/test_scope.py b/tests/test_scope.py index 86a0551a44..710ce33849 100644 --- a/tests/test_scope.py +++ b/tests/test_scope.py @@ -1021,3 +1021,54 @@ def test_trace_context_without_performance(sentry_init): assert trace_context["span_id"] == propagation_context.span_id assert trace_context["parent_span_id"] == propagation_context.parent_span_id assert "dynamic_sampling_context" in trace_context + + +def test_conversation_id_set_get(): + """Test that set_conversation_id and get_conversation_id work correctly.""" + scope = Scope() + assert scope.get_conversation_id() is None + + scope.set_conversation_id("test-conv-123") + assert scope.get_conversation_id() == "test-conv-123" + + +def test_conversation_id_remove(): + """Test that remove_conversation_id clears the conversation ID.""" + scope = Scope() + scope.set_conversation_id("test-conv-456") + assert scope.get_conversation_id() == "test-conv-456" + + scope.remove_conversation_id() + assert scope.get_conversation_id() is None + + +def test_conversation_id_overwrite(): + """Test that set_conversation_id overwrites existing value.""" + scope = Scope() + scope.set_conversation_id("first-conv") + scope.set_conversation_id("second-conv") + assert scope.get_conversation_id() == "second-conv" + + +def test_conversation_id_copy(): + """Test that conversation_id is preserved when scope is copied.""" + scope1 = Scope() + scope1.set_conversation_id("copy-test-conv") + + scope2 = copy.copy(scope1) + assert scope2.get_conversation_id() == "copy-test-conv" + + # Modifying copy should not affect original + scope2.set_conversation_id("modified-conv") + assert scope1.get_conversation_id() == "copy-test-conv" + assert scope2.get_conversation_id() == "modified-conv" + + +def test_conversation_id_clear(): + """Test that conversation_id is cleared when scope.clear() is called.""" + scope = Scope() + scope.set_conversation_id("clear-test-conv") + assert scope.get_conversation_id() == "clear-test-conv" + + scope.clear() + assert scope.get_conversation_id() is None diff --git a/tests/tracing/test_misc.py b/tests/tracing/test_misc.py index 619fff913b..8895c98dbc 100644 --- a/tests/tracing/test_misc.py +++ b/tests/tracing/test_misc.py @@ -605,3 +605,168 @@ def test_update_current_span(sentry_init, capture_events): "thread.id": mock.ANY, "thread.name": mock.ANY, } + + +class TestConversationIdPropagation: + """Tests for conversation_id propagation to AI spans.""" + + def test_conversation_id_propagates_to_span_with_gen_ai_operation_name( + self, sentry_init, capture_events + ): + """Span with gen_ai.operation.name data should get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-op-name-test") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(op="http.client") as span: + span.set_data("gen_ai.operation.name", "chat") + + (event,) = events + span_data = event["spans"][0]["data"] + assert span_data.get("gen_ai.conversation.id") == "conv-op-name-test" + + def test_conversation_id_propagates_to_span_with_ai_op( + self, sentry_init, capture_events + ): + """Span with ai.* op should get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-ai-op-test") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(op="ai.chat.completions"): + pass + + (event,) = events + span_data = event["spans"][0]["data"] + assert span_data.get("gen_ai.conversation.id") == "conv-ai-op-test" + + def test_conversation_id_propagates_to_span_with_gen_ai_op( + self, sentry_init, capture_events + ): + """Span with gen_ai.* op should get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-gen-ai-op-test") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(op="gen_ai.invoke_agent"): + pass + + (event,) = events + span_data = event["spans"][0]["data"] + assert span_data.get("gen_ai.conversation.id") == "conv-gen-ai-op-test" + + def test_conversation_id_not_propagated_to_non_ai_span( + self, sentry_init, capture_events + ): + """Non-AI span should NOT get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-should-not-appear") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(op="http.client") as span: + span.set_data("some.other.data", "value") + + (event,) = events + span_data = event["spans"][0]["data"] + assert "gen_ai.conversation.id" not in span_data + + def test_conversation_id_not_propagated_when_not_set( + self, sentry_init, capture_events + ): + """AI span should not have conversation_id if not set on scope.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + # Ensure no conversation_id is set + scope = sentry_sdk.get_current_scope() + scope.remove_conversation_id() + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(op="ai.chat.completions"): + pass + + (event,) = events + span_data = event["spans"][0]["data"] + assert "gen_ai.conversation.id" not in span_data + + def test_conversation_id_not_propagated_to_span_without_op( + self, sentry_init, capture_events + ): + """Span without op and without gen_ai.operation.name should NOT get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-no-op-test") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(name="unnamed-span") as span: + span.set_data("regular.data", "value") + + (event,) = events + span_data = event["spans"][0]["data"] + assert "gen_ai.conversation.id" not in span_data + + def test_conversation_id_propagates_with_gen_ai_operation_name_no_op( + self, sentry_init, capture_events + ): + """Span with gen_ai.operation.name but no op should still get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-no-op-but-data-test") + + with sentry_sdk.start_transaction(name="test-tx"): + with start_span(name="unnamed-span") as span: + span.set_data("gen_ai.operation.name", "embedding") + + (event,) = events + span_data = event["spans"][0]["data"] + assert span_data.get("gen_ai.conversation.id") == "conv-no-op-but-data-test" + + def test_conversation_id_propagates_to_transaction_with_ai_op( + self, sentry_init, capture_events + ): + """Transaction with ai.* op should get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-tx-ai-op-test") + + with sentry_sdk.start_transaction(op="ai.workflow", name="AI Workflow"): + pass + + (event,) = events + trace_data = event["contexts"]["trace"]["data"] + assert trace_data.get("gen_ai.conversation.id") == "conv-tx-ai-op-test" + + def test_conversation_id_not_propagated_to_non_ai_transaction( + self, sentry_init, capture_events + ): + """Non-AI transaction should NOT get conversation_id.""" + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + scope = sentry_sdk.get_current_scope() + scope.set_conversation_id("conv-tx-should-not-appear") + + with sentry_sdk.start_transaction(op="http.server", name="HTTP Request"): + pass + + (event,) = events + trace_data = event["contexts"]["trace"]["data"] + assert "gen_ai.conversation.id" not in trace_data