add meta entries

Author: shellmayr

sentry_sdk/ai/message_utils.py

@@ -4,17 +4,9 @@
 if TYPE_CHECKING:
     from typing import Any, Dict, List, Optional
 
-try:
-    from sentry_sdk.serializer import serialize
-except ImportError:
-    # Fallback for cases where sentry_sdk isn't fully importable
-    def serialize(obj, **kwargs):
-        # type: (Any, **Any) -> Any
-        return obj
+from sentry_sdk.serializer import serialize
+from sentry_sdk._types import AnnotatedValue
 
-
-# Custom limit for gen_ai message serialization - 50% of MAX_EVENT_BYTES
-# to leave room for other event data while still being generous for messages
 MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
 
 
@@ -50,21 +42,26 @@ def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
 
 
 def serialize_gen_ai_messages(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
-    # type: (Optional[List[Dict[str, Any]]], int) -> Optional[str]
+    # type: (Optional[Any], int) -> Optional[str]
     """
     Serialize and truncate gen_ai messages for storage in spans.
 
     This function handles the complete workflow of:
-    1. Truncating messages to fit within size limits
-    2. Serializing them using Sentry's serializer
+    1. Truncating messages to fit within size limits (if not already done)
+    2. Serializing them using Sentry's serializer (which processes AnnotatedValue for _meta)
     3. Converting to JSON string for storage
 
-    :param messages: List of message objects or None
+    :param messages: List of message objects, AnnotatedValue, or None
     :param max_bytes: Maximum allowed size in bytes for the serialized messages
     :returns: JSON string of serialized messages or None if input was None/empty
     """
     if not messages:
         return None
+
+    if isinstance(messages, AnnotatedValue):
+        serialized_messages = serialize(messages, is_vars=False)
+        return json.dumps(serialized_messages, separators=(",", ":"))
+
     truncated_messages = truncate_messages_by_size(messages, max_bytes)
     if not truncated_messages:
         return None
@@ -96,44 +93,31 @@ def get_messages_metadata(original_messages, truncated_messages):
 
 
 def truncate_and_serialize_messages(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
-    # type: (Optional[List[Dict[str, Any]]], int) -> Dict[str, Any]
+    # type: (Optional[List[Dict[str, Any]]], int) -> Any
     """
-    One-stop function for gen_ai integrations to truncate and serialize messages.
+    Truncate messages and return AnnotatedValue for automatic _meta creation.
 
-    This is the main function that gen_ai integrations should use. It handles the
-    complete workflow and returns both the serialized data and metadata.
-
-    Example usage:
-        from sentry_sdk.ai.message_utils import truncate_and_serialize_messages
-
-        result = truncate_and_serialize_messages(messages)
-        if result['serialized_data']:
-            span.set_data('gen_ai.request.messages', result['serialized_data'])
-        if result['metadata']['was_truncated']:
-            # Log warning about truncation if desired
-            pass
+    This function handles truncation and returns the truncated messages wrapped in an
+    AnnotatedValue (when truncation occurs) so that Sentry's serializer can automatically
+    create the appropriate _meta structure.
 
     :param messages: List of message objects or None
     :param max_bytes: Maximum allowed size in bytes for the serialized messages
-    :returns: Dictionary containing 'serialized_data', 'metadata', and 'original_size'
+    :returns: List of messages, AnnotatedValue (if truncated), or None
     """
     if not messages:
-        return {
-            "serialized_data": None,
-            "metadata": get_messages_metadata([], []),
-            "original_size": 0,
-        }
-
-    original_serialized = serialize(messages, is_vars=False)
-    original_json = json.dumps(original_serialized, separators=(",", ":"))
-    original_size = len(original_json.encode("utf-8"))
+        return None
 
     truncated_messages = truncate_messages_by_size(messages, max_bytes)
-    serialized_data = serialize_gen_ai_messages(truncated_messages, max_bytes)
-    metadata = get_messages_metadata(messages, truncated_messages)
+    if not truncated_messages:
+        return None
 
-    return {
-        "serialized_data": serialized_data,
-        "metadata": metadata,
-        "original_size": original_size,
-    }
+    original_count = len(messages)
+    truncated_count = len(truncated_messages)
+
+    if original_count != truncated_count:
+        return AnnotatedValue(
+            value=serialize_gen_ai_messages(truncated_messages),
+            metadata={"len": original_count},
+        )
+    return truncated_messages

sentry_sdk/integrations/langchain.py

@@ -210,11 +210,9 @@ def on_llm_start(
             _set_tools_on_span(span, all_params.get("tools"))
 
             if should_send_default_pii() and self.include_prompts:
-                result = truncate_and_serialize_messages(prompts)
-                if result["serialized_data"]:
-                    span.set_data(
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                    )
+                messages_data = truncate_and_serialize_messages(prompts)
+                if messages_data is not None:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
     def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
@@ -265,11 +263,9 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                         normalized_messages.append(
                             self._normalize_langchain_message(message)
                         )
-                result = truncate_and_serialize_messages(normalized_messages)
-                if result["serialized_data"]:
-                    span.set_data(
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                    )
+                messages_data = truncate_and_serialize_messages(normalized_messages)
+                if messages_data is not None:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
     def on_chat_model_end(self, response, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
@@ -742,11 +738,9 @@ def new_invoke(self, *args, **kwargs):
                 and should_send_default_pii()
                 and integration.include_prompts
             ):
-                result = truncate_and_serialize_messages([input])
-                if result["serialized_data"]:
-                    span.set_data(
-                        SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                    )
+                messages_data = truncate_and_serialize_messages([input])
+                if messages_data is not None:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
             output = result.get("output")
             if (
@@ -795,11 +789,9 @@ def new_stream(self, *args, **kwargs):
             and should_send_default_pii()
             and integration.include_prompts
         ):
-            result = truncate_and_serialize_messages([input])
-            if result["serialized_data"]:
-                span.set_data(
-                    SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                )
+            messages_data = truncate_and_serialize_messages([input])
+            if messages_data is not None:
+                span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
         # Run the agent
         result = f(self, *args, **kwargs)

sentry_sdk/integrations/langgraph.py

@@ -181,11 +181,9 @@ def new_invoke(self, *args, **kwargs):
             ):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
-                    result = truncate_and_serialize_messages(input_messages)
-                    if result["serialized_data"]:
-                        span.set_data(
-                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                        )
+                    messages_data = truncate_and_serialize_messages(input_messages)
+                    if messages_data is not None:
+                        span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
             result = f(self, *args, **kwargs)
 
@@ -230,11 +228,9 @@ async def new_ainvoke(self, *args, **kwargs):
             ):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
-                    result = truncate_and_serialize_messages(input_messages)
-                    if result["serialized_data"]:
-                        span.set_data(
-                            SPANDATA.GEN_AI_REQUEST_MESSAGES, result["serialized_data"]
-                        )
+                    messages_data = truncate_and_serialize_messages(input_messages)
+                    if messages_data is not None:
+                        span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
             result = await f(self, *args, **kwargs)
 

tests/test_ai_message_utils.py

@@ -8,6 +8,8 @@
     get_messages_metadata,
     truncate_and_serialize_messages,
 )
+from sentry_sdk._types import AnnotatedValue
+from sentry_sdk.serializer import serialize
 
 
 @pytest.fixture
@@ -407,3 +409,86 @@ def test_truncation_keeps_most_recent(self):
                 assert (
                     "Message 9" in last_kept_content or "Message 8" in last_kept_content
                 )
+
+
+class TestMetaSupport:
+    """Test that _meta entries are created correctly when truncation occurs"""
+
+    def test_annotated_value_returned_on_truncation(self, large_messages):
+        """Test that truncate_and_serialize_messages returns AnnotatedValue when truncation occurs"""
+        # Force truncation with a limit that will keep at least one message
+        # Each large message is ~30KB, so 50KB should keep 1-2 messages but force truncation
+        small_limit = 50_000  # 50KB to force truncation but keep some messages
+        result = truncate_and_serialize_messages(large_messages, max_bytes=small_limit)
+
+        # Should return an AnnotatedValue when truncation occurs
+        assert isinstance(result, AnnotatedValue)
+        assert result.metadata == {"len": len(large_messages)}
+
+        # The value should be the truncated messages
+        assert isinstance(result.value, list)
+        assert len(result.value) < len(large_messages)
+
+    def test_no_annotated_value_when_no_truncation(self, sample_messages):
+        """Test that truncate_and_serialize_messages returns plain list when no truncation occurs"""
+        result = truncate_and_serialize_messages(sample_messages)
+
+        # Should return plain list when no truncation occurs
+        assert not isinstance(result, AnnotatedValue)
+        assert isinstance(result, list)
+        assert len(result) == len(sample_messages)
+        assert result == sample_messages
+
+    def test_meta_structure_in_serialized_output(self, large_messages):
+        """Test that _meta structure is created correctly in serialized output"""
+        # Force truncation with a limit that will keep at least one message
+        small_limit = 50_000  # 50KB to force truncation but keep some messages
+        annotated_messages = truncate_and_serialize_messages(
+            large_messages, max_bytes=small_limit
+        )
+
+        # Simulate how the serializer would process this (like it does in actual span data)
+        test_data = {"gen_ai": {"request": {"messages": annotated_messages}}}
+
+        # Serialize using Sentry's serializer (which processes AnnotatedValue)
+        serialized = serialize(test_data, is_vars=False)
+
+        # Check that _meta structure was created
+        assert "_meta" in serialized
+        assert "gen_ai" in serialized["_meta"]
+        assert "request" in serialized["_meta"]["gen_ai"]
+        assert "messages" in serialized["_meta"]["gen_ai"]["request"]
+        assert serialized["_meta"]["gen_ai"]["request"]["messages"][""] == {
+            "len": len(large_messages)
+        }
+
+        # Check that the actual data is still there
+        assert "gen_ai" in serialized
+        assert "request" in serialized["gen_ai"]
+        assert "messages" in serialized["gen_ai"]["request"]
+        assert isinstance(serialized["gen_ai"]["request"]["messages"], list)
+        assert len(serialized["gen_ai"]["request"]["messages"]) < len(large_messages)
+
+    def test_serialize_gen_ai_messages_handles_annotated_value(self, large_messages):
+        """Test that serialize_gen_ai_messages handles AnnotatedValue input correctly"""
+        # Create an AnnotatedValue manually
+        truncated = large_messages[:2]  # Keep only first 2 messages
+        annotated = AnnotatedValue(
+            value=truncated, metadata={"len": len(large_messages)}
+        )
+
+        # serialize_gen_ai_messages should handle it
+        result = serialize_gen_ai_messages(annotated)
+
+        assert result is not None
+        parsed = json.loads(result)
+        assert isinstance(parsed, list)
+        assert len(parsed) == 2  # Only 2 messages kept
+
+    def test_empty_messages_no_annotated_value(self):
+        """Test that empty messages don't create AnnotatedValue"""
+        result = truncate_and_serialize_messages([])
+        assert result is None
+
+        result = truncate_and_serialize_messages(None)
+        assert result is None