From 4b876ea9ca16797eff0211596de2cbf34663b2a5 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Sat, 20 Sep 2025 12:33:37 -0700
Subject: [PATCH] test_sanitize_request_body_for_spend_logs_payload_mixed_types

---
 .../test_spend_tracking_utils.py              | 78 ++++++++++++++-----
 1 file changed, 58 insertions(+), 20 deletions(-)

diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
index 10d13ca3ee..4159f05cca 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
@@ -58,26 +58,52 @@ def test_sanitize_request_body_for_spend_logs_payload_long_string():
 
 
 def test_sanitize_request_body_for_spend_logs_payload_nested_dict():
-    request_body = {"outer": {"inner": {"text": "a" * 2000, "normal": "short"}}}
+    from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB
+
+    # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB
+    long_string = "a" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500)
+    request_body = {"outer": {"inner": {"text": long_string, "normal": "short"}}}
     sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
-    assert len(sanitized["outer"]["inner"]["text"]) == 1000 + len(
-        f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
-    )
+    
+    # Calculate expected lengths based on actual MAX_STRING_LENGTH_PROMPT_IN_DB
+    start_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.35)
+    end_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.65)
+    total_keep = start_chars + end_chars
+    if total_keep > MAX_STRING_LENGTH_PROMPT_IN_DB:
+        end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars
+    
+    skipped_chars = len(long_string) - total_keep
+    expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..."
+    expected_length = start_chars + len(expected_truncation_message) + end_chars
+    
+    assert len(sanitized["outer"]["inner"]["text"]) == expected_length
     assert sanitized["outer"]["inner"]["normal"] == "short"
 
 
 def test_sanitize_request_body_for_spend_logs_payload_nested_list():
+    from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB
+
+    # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB
+    long_string = "a" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500)
     request_body = {
-        "items": [{"text": "a" * 2000}, {"text": "short"}, [{"text": "a" * 2000}]]
+        "items": [{"text": long_string}, {"text": "short"}, [{"text": long_string}]]
     }
     sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
-    assert len(sanitized["items"][0]["text"]) == 1000 + len(
-        f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
-    )
+    
+    # Calculate expected lengths based on actual MAX_STRING_LENGTH_PROMPT_IN_DB
+    start_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.35)
+    end_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.65)
+    total_keep = start_chars + end_chars
+    if total_keep > MAX_STRING_LENGTH_PROMPT_IN_DB:
+        end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars
+    
+    skipped_chars = len(long_string) - total_keep
+    expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..."
+    expected_length = start_chars + len(expected_truncation_message) + end_chars
+    
+    assert len(sanitized["items"][0]["text"]) == expected_length
     assert sanitized["items"][1]["text"] == "short"
-    assert len(sanitized["items"][2][0]["text"]) == 1000 + len(
-        f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
-    )
+    assert len(sanitized["items"][2][0]["text"]) == expected_length
 
 
 def test_sanitize_request_body_for_spend_logs_payload_non_string_values():
@@ -93,21 +119,33 @@ def test_sanitize_request_body_for_spend_logs_payload_empty():
 
 
 def test_sanitize_request_body_for_spend_logs_payload_mixed_types():
+    from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB
+
+    # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB
+    long_string = "a" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500)
     request_body = {
-        "text": "a" * 2000,
+        "text": long_string,
         "number": 42,
-        "nested": {"list": ["short", "a" * 2000], "dict": {"key": "a" * 2000}},
+        "nested": {"list": ["short", long_string], "dict": {"key": long_string}},
     }
     sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
-    assert len(sanitized["text"]) == 1000 + len(f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)")
+    
+    # Calculate expected lengths based on actual MAX_STRING_LENGTH_PROMPT_IN_DB
+    start_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.35)
+    end_chars = int(MAX_STRING_LENGTH_PROMPT_IN_DB * 0.65)
+    total_keep = start_chars + end_chars
+    if total_keep > MAX_STRING_LENGTH_PROMPT_IN_DB:
+        end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars
+    
+    skipped_chars = len(long_string) - total_keep
+    expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..."
+    expected_length = start_chars + len(expected_truncation_message) + end_chars
+    
+    assert len(sanitized["text"]) == expected_length
     assert sanitized["number"] == 42
     assert sanitized["nested"]["list"][0] == "short"
-    assert len(sanitized["nested"]["list"][1]) == 1000 + len(
-        f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
-    )
-    assert len(sanitized["nested"]["dict"]["key"]) == 1000 + len(
-        f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
-    )
+    assert len(sanitized["nested"]["list"][1]) == expected_length
+    assert len(sanitized["nested"]["dict"]["key"]) == expected_length
 
 
 def test_sanitize_request_body_for_spend_logs_payload_circular_reference():