fix(logging): scrub pydantic cut sides that hold the prompt's own quotes

litellm/litellm_core_utils/redact_messages.py

@@ -379,8 +379,9 @@ _ESCAPED_CHARACTERS: Final = ('"', "\\")
 # 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_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}}}(?=['\"])")
+_CUT: Final = "..."
+_CUT_STARTS: Final = re.compile(f"(?={re.escape(_CUT)})")
+_QUOTES: Final = "'\""
 _MAX_SCRUB_COMPARISONS: Final = 10_000_000
 
 
@@ -437,13 +438,18 @@ def _cut_sides(text: str) -> str:
     return f"{text[:_CUT_SIDE_LENGTH]}{quotes}{text[-_CUT_SIDE_LENGTH:]}"
 
 
-def _cut_echoes(text: str, forms: Collection[str]) -> tuple[tuple[str, str], ...]:
+def _cut_echoes(text: str, forms: Collection[str]) -> tuple[tuple[str, str], ...] | None:
     """
     Each side of a quoted value ``text`` cuts short with ``...`` that starts or ends one of ``forms``, as
-    pydantic shows a long input value, with its pattern anchored to the quote and the cut.
+    pydantic shows a long input value, with its pattern anchored to the quote and the cut. A side can hold
+    the request's own quotes, so every quote within a side's length of the cut is tried as its opening or
+    closing one. ``None`` when checking every side against ``forms`` is too large a search for a failure path.
     """
-    heads: Final = frozenset(match.group() for match in _CUT_ECHO_HEAD.finditer(text))
-    tails: Final = frozenset(match.group() for match in _CUT_ECHO_TAIL.finditer(text))
+    cuts: Final = tuple(match.start() for match in _CUT_STARTS.finditer(text))
+    heads: Final = frozenset(itertools.chain.from_iterable(_cut_heads(text, cut) for cut in cuts))
+    tails: Final = frozenset(itertools.chain.from_iterable(_cut_tails(text, cut) for cut in cuts))
+    if (len(heads) + len(tails)) * len(forms) > _MAX_SCRUB_COMPARISONS:
+        return None
     return (
         *(
             (head, f"(?<=['\"]){re.escape(head)}(?=\\.\\.\\.)")
@@ -458,6 +464,23 @@ def _cut_echoes(text: str, forms: Collection[str]) -> tuple[tuple[str, str], ...
     )
 
 
+def _cut_heads(text: str, cut: int) -> tuple[str, ...]:
+    """Each run of ``text`` from just after a quote up to the ``...`` at ``cut``, no longer than a cut side."""
+    return tuple(
+        text[quote + 1 : cut] for quote in range(max(0, cut - _CUT_SIDE_LENGTH - 1), cut - 1) if text[quote] in _QUOTES
+    )
+
+
+def _cut_tails(text: str, cut: int) -> tuple[str, ...]:
+    """Each run of ``text`` from the end of the ``...`` at ``cut`` up to just before a quote, no longer than a cut side."""
+    start: Final = cut + len(_CUT)
+    return tuple(
+        text[start:quote]
+        for quote in range(start + 1, min(len(text), start + _CUT_SIDE_LENGTH + 1))
+        if text[quote] in _QUOTES
+    )
+
+
 def scrub_request_text(text: str, texts: Collection[str]) -> str:
     """
     ``text`` with every quote of the request's own ``texts`` replaced by ``REDACTED_BY_LITELLM``, so an
@@ -474,8 +497,11 @@ def scrub_request_text(text: str, texts: Collection[str]) -> str:
     cut_forms: Final = frozenset(
         itertools.chain.from_iterable(_echo_forms(_cut_sides(each)) for each in texts if len(each) > len(text))
     )
+    cut_echoes: Final = _cut_echoes(text, forms | cut_forms)
+    if cut_echoes is None:
+        return REDACTED_BY_LITELLM
     echoes: Final = sorted(
-        (*((form, _echo_pattern(form)) for form in forms if form in text), *_cut_echoes(text, forms | cut_forms)),
+        (*((form, _echo_pattern(form)) for form in forms if form in text), *cut_echoes),
         key=lambda echo: len(echo[0]),
         reverse=True,
     )

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -10,6 +10,7 @@ import copy
 import importlib
 import json
 import os
+import random
 import threading
 import tracemalloc
 from collections.abc import AsyncIterator, Callable, Mapping
@@ -2137,6 +2138,42 @@ def test_scrub_request_text_replaces_both_sides_of_a_pydantic_cut_input_value(va
     assert "redacted-by-litellm...redacted-by-litellm" in scrubbed
 
 
+SINGLE_QUOTED_PROMPT: Final = "account 'secret' " + "x" * 1_000 + " end 'private'"
+BOTH_QUOTED_PROMPT: Final = 'said "secret" it\'s ' + "x" * 1_000 + ' it\'s "private"'
+DOTTED_PROMPT: Final = "account secret is there." + "x" * 1_000 + ".private secret is here"
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        SINGLE_QUOTED_PROMPT,
+        [SINGLE_QUOTED_PROMPT],
+        {"content": SINGLE_QUOTED_PROMPT},
+        BOTH_QUOTED_PROMPT,
+        [BOTH_QUOTED_PROMPT],
+        {"content": BOTH_QUOTED_PROMPT},
+        DOTTED_PROMPT,
+    ],
+    ids=[
+        "single_quotes_str",
+        "single_quotes_list",
+        "single_quotes_dict",
+        "both_quotes_str",
+        "both_quotes_list",
+        "both_quotes_dict",
+        "dots_beside_both_cuts",
+    ],
+)
+def test_scrub_request_text_replaces_pydantic_cut_sides_that_hold_the_prompts_own_quotes(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(value))
+
+    assert not any(word in scrubbed for word in ("account", "said", "secret", "private"))
+    assert "redacted-by-litellm...redacted-by-litellm" in scrubbed
+
+
 def test_scrub_request_text_leaves_an_error_that_does_not_quote_the_prompt() -> None:
     assert scrub_request_text("Incorrect API key provided", request_texts(PROMPT, "hi")) == "Incorrect API key provided"
 
@@ -2153,6 +2190,14 @@ 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_replaces_the_whole_text_when_the_cut_side_search_is_too_large() -> None:
+    rng: Final = random.Random(0)
+    error: Final = "".join(rng.choice("'ab") if index % 10 else "..." for index in range(10_000))
+    texts: Final = request_texts([f'say "hi" {index}' for index in range(800)])
+
+    assert scrub_request_text(error, 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)])