Truncate single messages based on characters

Author: alexander-alderman-webb

sentry_sdk/ai/utils.py

@@ -1,4 +1,5 @@
 import json
+from copy import deepcopy
 from collections import deque
 from typing import TYPE_CHECKING
 from sys import getsizeof
@@ -12,6 +13,8 @@
 from sentry_sdk.utils import logger
 
 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:
@@ -101,27 +104,20 @@ def get_start_span_function():
     return sentry_sdk.start_span if transaction_exists else sentry_sdk.start_transaction
 
 
-def _truncate_single_message(message, max_bytes):
+def _truncate_single_message_content_if_present(message, max_chars):
     # type: (Dict[str, Any], int) -> Dict[str, Any]
     """
-    Truncate a single message to fit within max_bytes.
+    Truncate a single message to fit within max_chars.
     If the message is too large, truncate the content field.
     """
     if not isinstance(message, dict) or "content" not in message:
         return message
-    content = message.get("content", "")
+    content = message["content"]
 
-    if not isinstance(content, str) or len(content) <= max_bytes:
+    if not isinstance(content, str) or len(content) <= max_chars:
         return message
 
-    overhead_message = message.copy()
-    overhead_message["content"] = ""
-    overhead_size = len(
-        json.dumps(overhead_message, separators=(",", ":")).encode("utf-8")
-    )
-
-    available_content_bytes = max_bytes - overhead_size - 20
-    message["content"] = content[:available_content_bytes] + "..."
+    message["content"] = content[:max_chars] + "..."
     return message
 
 
@@ -142,22 +138,37 @@ def _find_truncation_index(messages, max_bytes):
     return 0
 
 
-def truncate_messages_by_size(messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES):
-    # type: (List[Dict[str, Any]], int) -> Tuple[List[Dict[str, Any]], int]
-    messages_with_truncated_content = [
-        _truncate_single_message(msg, max_bytes) for msg in messages
-    ]
-
-    serialized_json = json.dumps(messages_with_truncated_content, separators=(",", ":"))
+def truncate_messages_by_size(
+    messages,
+    max_bytes=MAX_GEN_AI_MESSAGE_BYTES,
+    max_single_message_chars=MAX_SINGLE_MESSAGE_CONTENT_CHARS,
+):
+    # type: (List[Dict[str, Any]], int, int) -> Tuple[List[Dict[str, Any]], int]
+    """
+    Returns a truncated messages array, consisting of
+    - the last message, with the messages's content truncated to `max_single_message_chars` characters,
+      if the last message's size exceeds `max_bytes`; otherwise,
+    - the maximum number of messages, starting from the end of the `messages` array, whose total
+      serialized size does not exceed `max_bytes` bytes.
+    """
+    serialized_json = json.dumps(messages, separators=(",", ":"))
     current_size = len(serialized_json.encode("utf-8"))
 
     if current_size <= max_bytes:
-        return messages_with_truncated_content, 0
+        return messages, 0
 
-    truncation_index = _find_truncation_index(
-        messages_with_truncated_content, max_bytes
+    truncation_index = _find_truncation_index(messages, max_bytes)
+    truncated_messages = (
+        messages[truncation_index:]
+        if truncation_index < len(messages)
+        else messages[-1:]
     )
-    return messages_with_truncated_content[truncation_index:], truncation_index
+    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(