update tests

Author: alexander-alderman-webb

tests/test_ai_monitoring.py

@@ -8,6 +8,7 @@
 from sentry_sdk.ai.monitoring import ai_track
 from sentry_sdk.ai.utils import (
     MAX_GEN_AI_MESSAGE_BYTES,
+    MAX_SINGLE_MESSAGE_CONTENT_CHARS,
     set_data_normalized,
     truncate_and_annotate_messages,
     truncate_messages_by_size,
@@ -177,7 +178,6 @@ async def async_tool(**kwargs):
 
 @pytest.fixture
 def sample_messages():
-    """Sample messages similar to what gen_ai integrations would use"""
     return [
         {"role": "system", "content": "You are a helpful assistant."},
         {
@@ -226,8 +226,7 @@ def test_truncation_removes_oldest_first(self, large_messages):
         )
         assert len(result) < len(large_messages)
 
-        if result:
-            assert result[-1] == large_messages[-1]
+        assert result[-1] == large_messages[-1]
         assert truncation_index == len(large_messages) - len(result)
 
     def test_empty_messages_list(self):
@@ -278,26 +277,22 @@ def test_progressive_truncation(self, large_messages):
             assert current_count >= 1
             prev_count = current_count
 
-    def test_individual_message_truncation(self):
-        large_content = "This is a very long message. " * 1000
+    def test_single_message_truncation(self):
+        large_content = "This is a very long message. " * 10_000
 
         messages = [
             {"role": "system", "content": "You are a helpful assistant."},
             {"role": "user", "content": large_content},
         ]
 
         result, truncation_index = truncate_messages_by_size(
-            messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES
+            messages, max_single_message_chars=MAX_SINGLE_MESSAGE_CONTENT_CHARS
         )
 
-        assert len(result) > 0
-
-        total_size = len(json.dumps(result, separators=(",", ":")).encode("utf-8"))
-        assert total_size <= MAX_GEN_AI_MESSAGE_BYTES
-
-        for msg in result:
-            msg_size = len(json.dumps(msg, separators=(",", ":")).encode("utf-8"))
-            assert msg_size <= MAX_GEN_AI_MESSAGE_BYTES
+        assert len(result) == 1
+        assert (
+            len(result[0]["content"].rstrip("...")) <= MAX_SINGLE_MESSAGE_CONTENT_CHARS
+        )
 
         # If the last message is too large, the system message is not present
         system_msgs = [m for m in result if m.get("role") == "system"]
@@ -309,53 +304,6 @@ def test_individual_message_truncation(self):
         assert user_msgs[0]["content"].endswith("...")
         assert len(user_msgs[0]["content"]) < len(large_content)
 
-    def test_combined_individual_and_array_truncation(self):
-        huge_content = "X" * 25000
-        medium_content = "Y" * 5000
-
-        messages = [
-            {"role": "system", "content": medium_content},
-            {"role": "user", "content": huge_content},
-            {"role": "assistant", "content": medium_content},
-            {"role": "user", "content": "small"},
-        ]
-
-        result, truncation_index = truncate_messages_by_size(
-            messages, max_bytes=MAX_GEN_AI_MESSAGE_BYTES
-        )
-
-        assert len(result) > 0
-
-        total_size = len(json.dumps(result, separators=(",", ":")).encode("utf-8"))
-        assert total_size <= MAX_GEN_AI_MESSAGE_BYTES
-
-        for msg in result:
-            msg_size = len(json.dumps(msg, separators=(",", ":")).encode("utf-8"))
-            assert msg_size <= MAX_GEN_AI_MESSAGE_BYTES
-
-        # The last user "small" message should always be present and untruncated
-        last_user_msgs = [
-            m for m in result if m.get("role") == "user" and m["content"] == "small"
-        ]
-        assert len(last_user_msgs) == 1
-
-        # If the huge message is present, it must be truncated
-        for user_msg in [
-            m for m in result if m.get("role") == "user" and "X" in m["content"]
-        ]:
-            assert user_msg["content"].endswith("...")
-            assert len(user_msg["content"]) < len(huge_content)
-
-        # The medium messages, if present, should not be truncated
-        for expected_role in ["system", "assistant"]:
-            role_msgs = [m for m in result if m.get("role") == expected_role]
-            if role_msgs:
-                assert role_msgs[0]["content"].startswith("Y")
-                assert len(role_msgs[0]["content"]) <= len(medium_content)
-                assert not role_msgs[0]["content"].endswith("...") or len(
-                    role_msgs[0]["content"]
-                ) == len(medium_content)
-
 
 class TestTruncateAndAnnotateMessages:
     def test_no_truncation_returns_list(self, sample_messages):