diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 358b2e6d72c..681f1080022 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -450,8 +450,9 @@ _STRUCTURAL_KEYS: Final = frozenset({"role", "type", "id", "tool_call_id", "name _MIN_UNQUOTED_TEXT_LENGTH: Final = 4 # pydantic shows a long input value as its repr's first 25 and last 24 characters around "...", quotes and # container keys included, so each side of the cut shows at most 24 characters of a quoted text -_CUT_ECHO_HEAD: Final = re.compile(r"(?<=['\"])[^'\"]{1,24}(?=\.\.\.)") -_CUT_ECHO_TAIL: Final = re.compile(r"(?<=\.\.\.)[^'\"]{1,24}(?=['\"])") +_CUT_SIDE_LENGTH: Final = 24 +_CUT_ECHO_HEAD: Final = re.compile(rf"(?<=['\"])[^'\"]{{1,{_CUT_SIDE_LENGTH}}}(?=\.\.\.)") +_CUT_ECHO_TAIL: Final = re.compile(rf"(?<=\.\.\.)[^'\"]{{1,{_CUT_SIDE_LENGTH}}}(?=['\"])") _MAX_SCRUB_COMPARISONS: Final = 10_000_000 @@ -493,6 +494,15 @@ def _echo_forms(text: str) -> frozenset[str]: return frozenset({text, json.dumps(text)[1:-1], json.dumps(text, ensure_ascii=False)[1:-1], repr(text)[1:-1]}) +def _cut_sides(text: str) -> str: + """ + The start and end of ``text`` that a pydantic cut can show, for a text too long to be quoted whole, with + ``text``'s quote characters between them so repr escapes the sides as it escapes the whole ``text``. + """ + quotes: Final = "".join(quote for quote in "'\"" if quote in text) + return f"{text[:_CUT_SIDE_LENGTH]}{quotes}{text[-_CUT_SIDE_LENGTH:]}" + + def _cut_echoes(text: str, forms: Collection[str]) -> tuple[tuple[str, str], ...]: """ Each side of a quoted value ``text`` cuts short with ``...`` that starts or ends one of ``forms``, as @@ -524,9 +534,14 @@ def scrub_request_text(text: str, texts: Collection[str]) -> str: return text if len(text) * len(texts) > _MAX_SCRUB_COMPARISONS: return REDACTED_BY_LITELLM - forms: Final = frozenset(itertools.chain.from_iterable(map(_echo_forms, texts))) + forms: Final = frozenset( + itertools.chain.from_iterable(_echo_forms(each) for each in texts if len(each) <= len(text)) + ) + cut_forms: Final = frozenset( + itertools.chain.from_iterable(_echo_forms(_cut_sides(each)) for each in texts if len(each) > len(text)) + ) echoes: Final = sorted( - (*((form, _echo_pattern(form)) for form in forms if form in text), *_cut_echoes(text, forms)), + (*((form, _echo_pattern(form)) for form in forms if form in text), *_cut_echoes(text, forms | cut_forms)), key=lambda echo: len(echo[0]), reverse=True, ) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 632957d08ee..78739ab6f25 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -11,6 +11,7 @@ import importlib import json import os import threading +import tracemalloc from collections.abc import AsyncIterator, Mapping from datetime import datetime from types import MappingProxyType, SimpleNamespace @@ -2084,6 +2085,7 @@ def test_perform_redaction_strips_the_matched_text_from_the_bedrock_guardrail_tr PROMPT: Final = "my ssn is 123-45-6789\nand my pin is 4321, keep both safe" +LONG_PROMPT: Final = f"my ssn is 123-45-6789 'keep it' and \"safe\" {'and more ' * 200}until the end" @pytest.mark.parametrize( @@ -2101,14 +2103,21 @@ def test_scrub_request_text_replaces_each_way_an_error_quotes_the_prompt(quoted: @pytest.mark.parametrize( "value", - [PROMPT, [PROMPT], {"content": PROMPT}, [{"content": PROMPT}]], - ids=["str", "list", "dict", "nested_under_12_characters"], + [PROMPT, [PROMPT], {"content": PROMPT}, [{"content": PROMPT}], [{"content": LONG_PROMPT}], [LONG_PROMPT]], + ids=[ + "str", + "list", + "dict", + "nested_under_12_characters", + "text_longer_than_the_error", + "text_longer_than_the_error_cut_inside_an_escaped_quote", + ], ) def test_scrub_request_text_replaces_both_sides_of_a_pydantic_cut_input_value(value: object) -> None: with pytest.raises(ValidationError, match="input_value=") as raised: TypeAdapter(int).validate_python(value) - scrubbed: Final = scrub_request_text(str(raised.value), request_texts(PROMPT)) + scrubbed: Final = scrub_request_text(str(raised.value), request_texts(value)) assert "my ssn" not in scrubbed and "4321" not in scrubbed assert "redacted-by-litellm...redacted-by-litellm" in scrubbed @@ -2130,6 +2139,25 @@ def test_scrub_request_text_replaces_the_whole_text_when_the_search_is_too_large assert scrub_request_text("x" * 5_000, texts) == "redacted-by-litellm" +def test_scrub_request_text_never_copies_a_request_text_longer_than_the_error() -> None: + error: Final = "litellm.RateLimitError: rate limit reached for gpt-4o" + image_urls: Final = request_texts([f"data:image/png;base64,{index}{'A' * 1_000_000}" for index in range(4)]) + already_tracing: Final = tracemalloc.is_tracing() + if not already_tracing: + tracemalloc.start() + tracemalloc.reset_peak() + baseline, _ = tracemalloc.get_traced_memory() + try: + scrubbed: Final = scrub_request_text(error, image_urls) + _, peak = tracemalloc.get_traced_memory() + finally: + if not already_tracing: + tracemalloc.stop() + + assert scrubbed == error + assert peak - baseline < 1_000_000, f"scrubbing a short error allocated {peak - baseline} bytes, copying the images" + + def test_request_texts_reads_content_and_skips_structural_fields() -> None: texts: Final = request_texts( [