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

Author: constantinius

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()

sentry_sdk/tracing.py

@@ -676,6 +676,12 @@ 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 a gen_ai span
+        conversation_id = scope.get_conversation_id()
+        if conversation_id and any(key.startswith("gen_ai.") for key in self._data):
+            self.set_data("gen_ai.conversation.id", conversation_id)
+
         maybe_create_breadcrumbs_from_span(scope, self)
 
         return None