feat(gen_ai): add function set_conversation_id and managing functions on the Scope and apply it on the Span on .finish()

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

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)

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:

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

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

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