ref(openai): Revert input truncation

Author: alexander-alderman-webb

sentry_sdk/ai/utils.py

@@ -1,11 +1,10 @@
 import inspect
 import json
-from copy import deepcopy
 from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Dict, List, Optional, Tuple
+    from typing import Any, Callable, Dict, Tuple
 
     from sentry_sdk.tracing import Span
 
@@ -14,10 +13,6 @@
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
-MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
-# Maximum characters when only a single message is left after bytes truncation
-MAX_SINGLE_MESSAGE_CONTENT_CHARS = 10_000
-
 
 class GEN_AI_ALLOWED_MESSAGE_ROLES:
     SYSTEM = "system"
@@ -180,92 +175,6 @@ def _truncate_single_message_content_if_present(
     return message
 
 
-def _find_truncation_index(messages: "List[Dict[str, Any]]", max_bytes: int) -> int:
-    """
-    Find the index of the first message that would exceed the max bytes limit.
-    Compute the individual message sizes, and return the index of the first message from the back
-    of the list that would exceed the max bytes limit.
-    """
-    running_sum = 0
-    for idx in range(len(messages) - 1, -1, -1):
-        size = len(json.dumps(messages[idx], separators=(",", ":")).encode("utf-8"))
-        running_sum += size
-        if running_sum > max_bytes:
-            return idx + 1
-
-    return 0
-
-
-def truncate_messages_by_size(
-    messages: "List[Dict[str, Any]]",
-    max_bytes: int = MAX_GEN_AI_MESSAGE_BYTES,
-    max_single_message_chars: int = MAX_SINGLE_MESSAGE_CONTENT_CHARS,
-) -> "Tuple[List[Dict[str, Any]], int]":
-    """
-    Returns a truncated messages list, consisting of
-    - the last message, with its content truncated to `max_single_message_chars` characters,
-      if the last message's size exceeds `max_bytes` bytes; otherwise,
-    - the maximum number of messages, starting from the end of the `messages` list, whose total
-      serialized size does not exceed `max_bytes` bytes.
-
-    In the single message case, the serialized message size may exceed `max_bytes`, because
-    truncation is based only on character count in that case.
-    """
-    serialized_json = json.dumps(messages, separators=(",", ":"))
-    current_size = len(serialized_json.encode("utf-8"))
-
-    if current_size <= max_bytes:
-        return messages, 0
-
-    truncation_index = _find_truncation_index(messages, max_bytes)
-    if truncation_index < len(messages):
-        truncated_messages = messages[truncation_index:]
-    else:
-        truncation_index = len(messages) - 1
-        truncated_messages = messages[-1:]
-
-    if len(truncated_messages) == 1:
-        truncated_messages[0] = _truncate_single_message_content_if_present(
-            deepcopy(truncated_messages[0]), max_chars=max_single_message_chars
-        )
-
-    return truncated_messages, truncation_index
-
-
-def truncate_and_annotate_messages(
-    messages: "Optional[List[Dict[str, Any]]]",
-    span: "Any",
-    scope: "Any",
-    max_single_message_chars: int = MAX_SINGLE_MESSAGE_CONTENT_CHARS,
-) -> "Optional[List[Dict[str, Any]]]":
-    if not messages:
-        return None
-
-    truncated_message = _truncate_single_message_content_if_present(
-        deepcopy(messages[-1]), max_chars=max_single_message_chars
-    )
-    if len(messages) > 1:
-        scope._gen_ai_original_message_count[span.span_id] = len(messages)
-
-    return [truncated_message]
-
-
-def truncate_and_annotate_embedding_inputs(
-    messages: "Optional[List[Dict[str, Any]]]",
-    span: "Any",
-    scope: "Any",
-    max_bytes: int = MAX_GEN_AI_MESSAGE_BYTES,
-) -> "Optional[List[Dict[str, Any]]]":
-    if not messages:
-        return None
-
-    truncated_messages, removed_count = truncate_messages_by_size(messages, max_bytes)
-    if removed_count > 0:
-        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.

sentry_sdk/integrations/openai.py

@@ -10,8 +10,6 @@
 from sentry_sdk.ai.utils import (
     set_data_normalized,
     normalize_message_roles,
-    truncate_and_annotate_messages,
-    truncate_and_annotate_embedding_inputs,
 )
 from sentry_sdk.ai._openai_completions_api import (
     _is_system_instruction as _is_system_instruction_completions,
@@ -397,12 +395,9 @@ def _set_responses_api_input_data(
 
     if isinstance(messages, str):
         normalized_messages = normalize_message_roles([messages])  # type: ignore
-        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
-            )
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_REQUEST_MESSAGES, normalized_messages, unpack=False
+        )
 
         set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
         return
@@ -412,12 +407,9 @@ def _set_responses_api_input_data(
     ]
     if len(non_system_messages) > 0:
         normalized_messages = normalize_message_roles(non_system_messages)
-        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
-            )
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_REQUEST_MESSAGES, normalized_messages, unpack=False
+        )
 
     set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
 
@@ -471,12 +463,9 @@ def _set_completions_api_input_data(
 
     if isinstance(messages, str):
         normalized_messages = normalize_message_roles([messages])  # type: ignore
-        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
-            )
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_REQUEST_MESSAGES, normalized_messages, unpack=False
+        )
         set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
         return
 
@@ -502,12 +491,9 @@ def _set_completions_api_input_data(
     ]
     if len(non_system_messages) > 0:
         normalized_messages = normalize_message_roles(non_system_messages)
-        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
-            )
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_REQUEST_MESSAGES, normalized_messages, unpack=False
+        )
 
     set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
 
@@ -538,14 +524,9 @@ def _set_embeddings_input_data(
         set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
 
         normalized_messages = normalize_message_roles([messages])  # type: ignore
-        scope = sentry_sdk.get_current_scope()
-        messages_data = truncate_and_annotate_embedding_inputs(
-            normalized_messages, span, scope
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, normalized_messages, unpack=False
         )
-        if messages_data is not None:
-            set_data_normalized(
-                span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False
-            )
 
         return
 
@@ -559,14 +540,9 @@ def _set_embeddings_input_data(
 
     if len(messages) > 0:
         normalized_messages = normalize_message_roles(messages)
-        scope = sentry_sdk.get_current_scope()
-        messages_data = truncate_and_annotate_embedding_inputs(
-            normalized_messages, span, scope
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, normalized_messages, unpack=False
         )
-        if messages_data is not None:
-            set_data_normalized(
-                span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False
-            )
 
     set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
 

tests/integrations/openai/test_openai.py

@@ -3724,51 +3724,6 @@ def test_openai_message_role_mapping(
     assert stored_messages[0]["role"] == expected_role
 
 
-def test_openai_message_truncation(sentry_init, capture_items):
-    """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,
-    )
-    items = capture_items("transaction", "span")
-
-    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,
-        )
-
-    span = next(item.payload for item in items if item.type == "span")
-    assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["attributes"]
-
-    messages_data = span["attributes"][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)
-
-    (event,) = (item.payload for item in items if item.type == "transaction")
-    meta_path = event["_meta"]
-    span_meta = meta_path["spans"]["0"]["data"]
-    messages_meta = span_meta[SPANDATA.GEN_AI_REQUEST_MESSAGES]
-    assert "len" in messages_meta.get("", {})
-
-
 # noinspection PyTypeChecker
 def test_streaming_chat_completion_ttft(
     sentry_init, capture_items, get_model_response, server_side_event_chunks