use annotatedvalue instead of previous approach

shellmayr · shellmayr · commit 8f0e0492add2 · 2025-10-14T14:15:26.000+02:00
diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py
@@ -109,26 +109,11 @@ def get_start_span_function():
 
 def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
     # type: (List[Dict[str, Any]], int) -> List[Dict[str, Any]]
-    """
-    Truncate messages by removing the oldest ones until the serialized size is within limits.
-    If the last message is still too large, truncate its content instead of removing it entirely.
-
-    This function prioritizes keeping the most recent messages while ensuring the total
-    serialized size stays under the specified byte limit. It uses the Sentry serializer
-    to get accurate size estimates that match what will actually be sent.
-
-    Always preserves at least one message, even if content needs to be truncated.
-
-    :param messages: List of message objects (typically with 'role' and 'content' keys)
-    :param max_bytes: Maximum allowed size in bytes for the serialized messages
-    :returns: Truncated list of messages that fits within the size limit
-    """
     if not messages:
         return messages
 
     truncated_messages = list(messages)
 
-    # First, remove older messages until we're under the limit or have only one message left
     while len(truncated_messages) > 1:
         serialized = serialize(
             truncated_messages, is_vars=False, max_value_length=round(max_bytes * 0.8)
@@ -139,10 +124,8 @@ def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
         if current_size <= max_bytes:
             break
 
-        truncated_messages.pop(0)  # Remove oldest message
+        truncated_messages.pop(0)
 
-    # If we still have one message but it's too large, truncate its content
-    # This ensures we always preserve at least one message
     if len(truncated_messages) == 1:
         serialized = serialize(
             truncated_messages, is_vars=False, max_value_length=round(max_bytes * 0.8)
@@ -151,7 +134,6 @@ def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
         current_size = len(serialized_json.encode("utf-8"))
 
         if current_size > max_bytes:
-            # Truncate the content of the last message
             last_message = truncated_messages[0].copy()
             content = last_message.get("content", "")
 
@@ -162,71 +144,24 @@ def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
     return truncated_messages
 
 
-def serialize_gen_ai_messages(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
-    # 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 (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, 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
-    """
+def truncate_and_annotate_messages(
+    messages, span, scope, max_bytes=MAX_GEN_AI_MESSAGE_BYTES
+):
+    # type: (Optional[List[Dict[str, Any]]], Any, Any, int) -> Optional[List[Dict[str, Any]]]
     if not messages:
         return None
 
-    if isinstance(messages, AnnotatedValue):
-        serialized_messages = serialize(
-            messages, is_vars=False, max_value_length=round(max_bytes * 0.8)
-        )
-        return json.dumps(serialized_messages, separators=(",", ":"))
-
+    original_count = len(messages)
     truncated_messages = truncate_messages_by_size(messages, max_bytes)
-    serialized_messages = serialize(
-        truncated_messages, is_vars=False, max_value_length=round(max_bytes * 0.8)
-    )
-
-    return json.dumps(serialized_messages)
 
-
-def truncate_and_serialize_messages(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
-    # type: (Optional[List[Dict[str, Any]]], int) -> Any
-    """
-    Truncate messages and return serialized string or AnnotatedValue for automatic _meta creation.
-
-    This function handles truncation and always returns serialized JSON strings. When truncation
-    occurs, it wraps the serialized string in an AnnotatedValue 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: JSON string, AnnotatedValue containing JSON string (if truncated), or None
-    """
-    if not messages:
-        return None
-
-    truncated_messages = truncate_messages_by_size(messages, max_bytes)
     if not truncated_messages:
         return None
 
-    # Always serialize to JSON string
-    serialized_json = serialize_gen_ai_messages(truncated_messages, max_bytes)
-    if not serialized_json:
-        return None
-
-    original_count = len(messages)
     truncated_count = len(truncated_messages)
+    n_removed = original_count - truncated_count
 
-    # If truncation occurred, wrap the serialized string in AnnotatedValue for _meta
-    if original_count != truncated_count:
-        return AnnotatedValue(
-            value=serialized_json,
-            metadata={"len": original_count},
-        )
+    if n_removed > 0:
+        scope._gen_ai_messages_truncated[span.span_id] = n_removed
+        span.set_data("_gen_ai_messages_original_count", original_count)
 
-    # No truncation, return plain serialized string
-    return serialized_json
+    return truncated_messages
diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py
@@ -614,6 +614,32 @@ def _prepare_event(
             event["breadcrumbs"] = AnnotatedValue(
                 event.get("breadcrumbs", []), {"len": previous_total_breadcrumbs}
             )
+
+        # Annotate truncated gen_ai messages in spans
+        if scope is not None and scope._gen_ai_messages_truncated:
+            spans = event.get("spans", [])
+            if isinstance(spans, AnnotatedValue):
+                spans = spans.value
+
+            for span in spans:
+                if isinstance(span, dict):
+                    span_id = span.get("span_id")
+                    if span_id and span_id in scope._gen_ai_messages_truncated:
+                        span_data = span.get("data", {})
+                        original_count = span_data.pop(
+                            "_gen_ai_messages_original_count", None
+                        )
+                        if (
+                            original_count is not None
+                            and SPANDATA.GEN_AI_REQUEST_MESSAGES in span_data
+                        ):
+                            span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES] = (
+                                AnnotatedValue(
+                                    span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES],
+                                    {"len": original_count},
+                                )
+                            )
+
         # Postprocess the event here so that annotated types do
         # generally not surface in before_send
         if event is not None:
diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py
@@ -6,7 +6,7 @@
 from sentry_sdk.ai.utils import (
     set_data_normalized,
     normalize_message_roles,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
     get_start_span_function,
 )
 from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
@@ -146,9 +146,12 @@ def _set_input_data(span, kwargs, integration):
                 normalized_messages.append(message)
 
         role_normalized_messages = normalize_message_roles(normalized_messages)
-        serialized_messages = truncate_and_serialize_messages(role_normalized_messages)
-        if serialized_messages is not None:
-            span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, serialized_messages)
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(
+            role_normalized_messages, span, scope
+        )
+        if messages_data is not None:
+            span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
     set_data_normalized(
         span, SPANDATA.GEN_AI_RESPONSE_STREAMING, kwargs.get("stream", False)
diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py
@@ -9,7 +9,7 @@
     normalize_message_roles,
     set_data_normalized,
     get_start_span_function,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
@@ -222,7 +222,10 @@ def on_llm_start(
                     }
                     for prompt in prompts
                 ]
-                messages_data = truncate_and_serialize_messages(normalized_messages)
+                scope = sentry_sdk.get_current_scope()
+                messages_data = truncate_and_annotate_messages(
+                    normalized_messages, span, scope
+                )
                 if messages_data is not None:
                     span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
@@ -276,7 +279,10 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                             self._normalize_langchain_message(message)
                         )
                 normalized_messages = normalize_message_roles(normalized_messages)
-                messages_data = truncate_and_serialize_messages(normalized_messages)
+                scope = sentry_sdk.get_current_scope()
+                messages_data = truncate_and_annotate_messages(
+                    normalized_messages, span, scope
+                )
                 if messages_data is not None:
                     span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
@@ -752,7 +758,10 @@ def new_invoke(self, *args, **kwargs):
                 and integration.include_prompts
             ):
                 normalized_messages = normalize_message_roles([input])
-                messages_data = truncate_and_serialize_messages(normalized_messages)
+                scope = sentry_sdk.get_current_scope()
+                messages_data = truncate_and_annotate_messages(
+                    normalized_messages, span, scope
+                )
                 if messages_data is not None:
                     span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
@@ -804,7 +813,9 @@ def new_stream(self, *args, **kwargs):
             and integration.include_prompts
         ):
             normalized_messages = normalize_message_roles([input])
-            messages_data = truncate_and_serialize_messages(normalized_messages)
+            messages_data = truncate_and_annotate_messages(
+                normalized_messages, span, sentry_sdk.get_current_scope()
+            )
             if messages_data is not None:
                 span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py
@@ -5,7 +5,7 @@
 from sentry_sdk.ai.utils import (
     set_data_normalized,
     normalize_message_roles,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
@@ -185,8 +185,9 @@ def new_invoke(self, *args, **kwargs):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
                     normalized_input_messages = normalize_message_roles(input_messages)
-                    messages_data = truncate_and_serialize_messages(
-                        normalized_input_messages
+                    scope = sentry_sdk.get_current_scope()
+                    messages_data = truncate_and_annotate_messages(
+                        normalized_input_messages, span, scope
                     )
                     if messages_data is not None:
                         span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
@@ -235,8 +236,9 @@ async def new_ainvoke(self, *args, **kwargs):
                 input_messages = _parse_langgraph_messages(args[0])
                 if input_messages:
                     normalized_input_messages = normalize_message_roles(input_messages)
-                    messages_data = truncate_and_serialize_messages(
-                        normalized_input_messages
+                    scope = sentry_sdk.get_current_scope()
+                    messages_data = truncate_and_annotate_messages(
+                        normalized_input_messages, span, scope
                     )
                     if messages_data is not None:
                         span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py
@@ -6,7 +6,8 @@
 from sentry_sdk.ai.utils import (
     get_start_span_function,
     set_data_normalized,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
+    normalize_message_roles,
 )
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
@@ -76,7 +77,9 @@ def _input_callback(kwargs):
 
     # Record messages if allowed
     if messages and should_send_default_pii() and integration.include_prompts:
-        messages_data = truncate_and_serialize_messages(messages)
+        normalized_messages = normalize_message_roles(messages)
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
         if messages_data is not None:
             span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py
@@ -6,7 +6,7 @@
 from sentry_sdk.ai.utils import (
     set_data_normalized,
     normalize_message_roles,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
@@ -187,7 +187,8 @@ def _set_input_data(span, kwargs, operation, integration):
         and integration.include_prompts
     ):
         normalized_messages = normalize_message_roles(messages)
-        messages_data = truncate_and_serialize_messages(normalized_messages)
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
         if messages_data is not None:
             span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py
@@ -3,7 +3,7 @@
     get_start_span_function,
     set_data_normalized,
     normalize_message_roles,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.scope import should_send_default_pii
@@ -62,7 +62,10 @@ def invoke_agent_span(context, agent, kwargs):
 
         if len(messages) > 0:
             normalized_messages = normalize_message_roles(messages)
-            messages_data = truncate_and_serialize_messages(normalized_messages)
+            scope = sentry_sdk.get_current_scope()
+            messages_data = truncate_and_annotate_messages(
+                normalized_messages, span, scope
+            )
             if messages_data is not None:
                 span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
diff --git a/sentry_sdk/integrations/openai_agents/utils.py b/sentry_sdk/integrations/openai_agents/utils.py
@@ -4,7 +4,7 @@
     normalize_message_roles,
     set_data_normalized,
     normalize_message_role,
-    truncate_and_serialize_messages,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import SPANDATA, SPANSTATUS, OP
 from sentry_sdk.integrations import DidNotEnable
@@ -137,7 +137,10 @@ def _set_input_data(span, get_response_kwargs):
                 )
 
     role_normalized_messages = normalize_message_roles(request_messages)
-    messages_data = truncate_and_serialize_messages(role_normalized_messages)
+    scope = sentry_sdk.get_current_scope()
+    messages_data = truncate_and_annotate_messages(
+        role_normalized_messages, span, scope
+    )
     if messages_data is not None:
         span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data)
 
diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
@@ -188,6 +188,7 @@ class Scope:
         "_extras",
         "_breadcrumbs",
         "_n_breadcrumbs_truncated",
+        "_gen_ai_messages_truncated",
         "_event_processors",
         "_error_processors",
         "_should_capture",
@@ -213,6 +214,7 @@ def __init__(self, ty=None, client=None):
         self._name = None  # type: Optional[str]
         self._propagation_context = None  # type: Optional[PropagationContext]
         self._n_breadcrumbs_truncated = 0  # type: int
+        self._gen_ai_messages_truncated = {}  # type: Dict[str, int]
 
         self.client = NonRecordingClient()  # type: sentry_sdk.client.BaseClient
 
@@ -247,6 +249,7 @@ def __copy__(self):
 
         rv._breadcrumbs = copy(self._breadcrumbs)
         rv._n_breadcrumbs_truncated = self._n_breadcrumbs_truncated
+        rv._gen_ai_messages_truncated = self._gen_ai_messages_truncated.copy()
         rv._event_processors = self._event_processors.copy()
         rv._error_processors = self._error_processors.copy()
         rv._propagation_context = self._propagation_context
@@ -1583,6 +1586,8 @@ def update_from_scope(self, scope):
             self._n_breadcrumbs_truncated = (
                 self._n_breadcrumbs_truncated + scope._n_breadcrumbs_truncated
             )
+        if scope._gen_ai_messages_truncated:
+            self._gen_ai_messages_truncated.update(scope._gen_ai_messages_truncated)
         if scope._span:
             self._span = scope._span
         if scope._attachments:
diff --git a/tests/test_ai_message_utils.py b/tests/test_ai_message_utils.py
