fix(ai): introduce message truncation for openai

Author: shellmayr

sentry_sdk/ai/utils.py

@@ -1,14 +1,24 @@
 import json
-
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Callable
+
     from sentry_sdk.tracing import Span
 
+from typing import TYPE_CHECKING
+
 import sentry_sdk
 from sentry_sdk.utils import logger
 
+if TYPE_CHECKING:
+    from typing import Any, Dict, List, Optional
+
+from sentry_sdk._types import AnnotatedValue
+from sentry_sdk.serializer import serialize
+
+MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
+
 
 class GEN_AI_ALLOWED_MESSAGE_ROLES:
     SYSTEM = "system"
@@ -95,3 +105,79 @@ def get_start_span_function():
         current_span is not None and current_span.containing_transaction is not None
     )
     return sentry_sdk.start_span if transaction_exists else sentry_sdk.start_transaction
+
+
+def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
+    # type: (List[Dict[str, Any]], int) -> List[Dict[str, Any]]
+    if not messages:
+        return messages
+
+    truncated_messages = list(messages)
+
+    while len(truncated_messages) > 1:
+        serialized_json = json.dumps(truncated_messages, separators=(",", ":"))
+        current_size = len(serialized_json.encode("utf-8"))
+
+        if current_size <= max_bytes:
+            break
+
+        truncated_messages.pop(0)
+
+    serialized_json = json.dumps(truncated_messages, separators=(",", ":"))
+    current_size = len(serialized_json.encode("utf-8"))
+
+    if current_size > max_bytes and len(truncated_messages) == 1:
+        message = truncated_messages[0].copy()
+        content = message.get("content", "")
+
+        if isinstance(content, str):
+            max_content_length = max_bytes // 2
+            while True:
+                message["content"] = content[:max_content_length]
+                test_json = json.dumps([message], separators=(",", ":"))
+                if len(test_json.encode("utf-8")) <= max_bytes:
+                    break
+                max_content_length = int(max_content_length * 0.9)
+                if max_content_length < 100:
+                    message["content"] = ""
+                    break
+
+            truncated_messages = [message]
+        elif isinstance(content, list):
+            content_copy = list(content)
+            while len(content_copy) > 0:
+                message["content"] = content_copy
+                test_json = json.dumps([message], separators=(",", ":"))
+                if len(test_json.encode("utf-8")) <= max_bytes:
+                    break
+                content_copy = content_copy[:-1]
+
+            if len(content_copy) == 0:
+                message["content"] = []
+
+            truncated_messages = [message]
+
+    return truncated_messages
+
+
+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
+
+    original_count = len(messages)
+    truncated_messages = truncate_messages_by_size(messages, max_bytes)
+
+    if not truncated_messages:
+        return None
+
+    truncated_count = len(truncated_messages)
+    n_removed = original_count - truncated_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)
+
+    return truncated_messages

sentry_sdk/client.py

@@ -606,6 +606,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:

sentry_sdk/integrations/openai.py

@@ -1,10 +1,13 @@
 from functools import wraps
-from collections.abc import Iterable
 
 import sentry_sdk
 from sentry_sdk import consts
 from sentry_sdk.ai.monitoring import record_token_usage
-from sentry_sdk.ai.utils import set_data_normalized, normalize_message_roles
+from sentry_sdk.ai.utils import (
+    set_data_normalized,
+    normalize_message_roles,
+    truncate_and_annotate_messages,
+)
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
@@ -18,19 +21,14 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, List, Optional, Callable, AsyncIterator, Iterator
+    from typing import Any, Iterable, List, Optional, Callable, AsyncIterator, Iterator
     from sentry_sdk.tracing import Span
 
 try:
     try:
-        from openai import NotGiven
+        from openai import NOT_GIVEN
     except ImportError:
-        NotGiven = None
-
-    try:
-        from openai import Omit
-    except ImportError:
-        Omit = None
+        NOT_GIVEN = None
 
     from openai.resources.chat.completions import Completions, AsyncCompletions
     from openai.resources import Embeddings, AsyncEmbeddings
@@ -189,9 +187,12 @@ def _set_input_data(span, kwargs, operation, integration):
         and integration.include_prompts
     ):
         normalized_messages = normalize_message_roles(messages)
-        set_data_normalized(
-            span, SPANDATA.GEN_AI_REQUEST_MESSAGES, normalized_messages, unpack=False
-        )
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
+        if messages_data is not None:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
+            )
 
     # Input attributes: Common
     set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "openai")
@@ -210,12 +211,12 @@ def _set_input_data(span, kwargs, operation, integration):
     for key, attribute in kwargs_keys_to_attributes.items():
         value = kwargs.get(key)
 
-        if value is not None and _is_given(value):
+        if value is not NOT_GIVEN and value is not None:
             set_data_normalized(span, attribute, value)
 
     # Input attributes: Tools
     tools = kwargs.get("tools")
-    if tools is not None and _is_given(tools) and len(tools) > 0:
+    if tools is not NOT_GIVEN and tools is not None and len(tools) > 0:
         set_data_normalized(
             span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
         )
@@ -695,15 +696,3 @@ async def _sentry_patched_responses_async(*args, **kwargs):
         return await _execute_async(f, *args, **kwargs)
 
     return _sentry_patched_responses_async
-
-
-def _is_given(obj):
-    # type: (Any) -> bool
-    """
-    Check for givenness safely across different openai versions.
-    """
-    if NotGiven is not None and isinstance(obj, NotGiven):
-        return False
-    if Omit is not None and isinstance(obj, Omit):
-        return False
-    return True

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:

tests/integrations/openai/test_openai.py

@@ -1,3 +1,4 @@
+import json
 import pytest
 
 from sentry_sdk.utils import package_version
@@ -7,11 +8,6 @@
 except ImportError:
     NOT_GIVEN = None
 
-try:
-    from openai import omit
-except ImportError:
-    omit = None
-
 from openai import AsyncOpenAI, OpenAI, AsyncStream, Stream, OpenAIError
 from openai.types import CompletionUsage, CreateEmbeddingResponse, Embedding
 from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionChunk
@@ -44,6 +40,9 @@
     OpenAIIntegration,
     _calculate_token_usage,
 )
+from sentry_sdk.ai.utils import MAX_GEN_AI_MESSAGE_BYTES
+from sentry_sdk._types import AnnotatedValue
+from sentry_sdk.serializer import serialize
 
 from unittest import mock  # python 3.3 and above
 
@@ -1429,7 +1428,7 @@ async def test_streaming_responses_api_async(
 )
 @pytest.mark.parametrize(
     "tools",
-    [[], None, NOT_GIVEN, omit],
+    [[], None, NOT_GIVEN],
 )
 def test_empty_tools_in_chat_completion(sentry_init, capture_events, tools):
     sentry_init(
@@ -1456,6 +1455,7 @@ def test_empty_tools_in_chat_completion(sentry_init, capture_events, tools):
 
 def test_openai_message_role_mapping(sentry_init, capture_events):
     """Test that OpenAI integration properly maps message roles like 'ai' to 'assistant'"""
+
     sentry_init(
         integrations=[OpenAIIntegration(include_prompts=True)],
         traces_sample_rate=1.0,
@@ -1465,7 +1465,6 @@ def test_openai_message_role_mapping(sentry_init, capture_events):
 
     client = OpenAI(api_key="z")
     client.chat.completions._post = mock.Mock(return_value=EXAMPLE_CHAT_COMPLETION)
-
     # Test messages with mixed roles including "ai" that should be mapped to "assistant"
     test_messages = [
         {"role": "system", "content": "You are helpful."},
@@ -1476,11 +1475,9 @@ def test_openai_message_role_mapping(sentry_init, capture_events):
 
     with start_transaction(name="openai tx"):
         client.chat.completions.create(model="test-model", messages=test_messages)
-
+    # Verify that the span was created correctly
     (event,) = events
     span = event["spans"][0]
-
-    # Verify that the span was created correctly
     assert span["op"] == "gen_ai.chat"
     assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"]
 
@@ -1505,3 +1502,55 @@ def test_openai_message_role_mapping(sentry_init, capture_events):
     # Verify no "ai" roles remain
     roles = [msg["role"] for msg in stored_messages]
     assert "ai" not in roles
+
+
+def test_openai_message_truncation(sentry_init, capture_events):
+    """Test that large messages are truncated properly in OpenAI integration."""
+    sentry_init(
+        integrations=[OpenAIIntegration(include_prompts=True)],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+    events = capture_events()
+
+    client = OpenAI(api_key="z")
+    client.chat.completions._post = mock.Mock(return_value=EXAMPLE_CHAT_COMPLETION)
+
+    large_content = (
+        "This is a very long message that will exceed our size limits. " * 1000
+    )
+    large_messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": large_content},
+        {"role": "assistant", "content": large_content},
+        {"role": "user", "content": large_content},
+    ]
+
+    with start_transaction(name="openai tx"):
+        client.chat.completions.create(
+            model="some-model",
+            messages=large_messages,
+        )
+
+    (event,) = events
+    span = event["spans"][0]
+    assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"]
+
+    messages_data = span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]
+    assert isinstance(messages_data, str)
+
+    parsed_messages = json.loads(messages_data)
+    assert isinstance(parsed_messages, list)
+    assert len(parsed_messages) <= len(large_messages)
+
+    if "_meta" in event and len(parsed_messages) < len(large_messages):
+        meta_path = event["_meta"]
+        if (
+            "spans" in meta_path
+            and "0" in meta_path["spans"]
+            and "data" in meta_path["spans"]["0"]
+        ):
+            span_meta = meta_path["spans"]["0"]["data"]
+            if SPANDATA.GEN_AI_REQUEST_MESSAGES in span_meta:
+                messages_meta = span_meta[SPANDATA.GEN_AI_REQUEST_MESSAGES]
+                assert "len" in messages_meta.get("", {})