fix(logging): scrub request text as whole words without re's pattern cache, and rebuild an exception's frames when no traceback was logged

litellm/litellm_core_utils/redact_messages.py

@@ -10,11 +10,13 @@
 import asyncio
 import contextlib
 import copy
+import functools
 import inspect
 import itertools
 import json
 import re
-from collections.abc import Collection, Iterable, Mapping, Sequence
+import traceback
+from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any, Final
 
@@ -432,6 +434,11 @@ _ESCAPED_CHARACTERS: Final = ('"', "\\")
 _INPUT_VALUE_LABEL: Final = "input_value="
 _INPUT_TYPE_LABEL: Final = ", input_type="
 _MAX_SCRUB_COMPARISONS: Final = 10_000_000
+_MAX_ECHO_OCCURRENCES: Final = 10_000
+_QUOTES: Final = ("'", '"')
+_ESCAPE_TAIL: Final = re.compile(r"\\(?:[abfnrtv]|x[0-9a-fA-F]{2}|u[0-9a-fA-F]{4}|U[0-9a-fA-F]{8})\Z")
+_DIRECT_CAUSE_LINE: Final = "\nThe above exception was the direct cause of the following exception:\n\n"
+_CONTEXT_LINE: Final = "\nDuring handling of the above exception, another exception occurred:\n\n"
 
 
 def _child_values(value: object) -> tuple[object, ...]:
@@ -498,14 +505,83 @@ def scrub_request_text(text: str, texts: Collection[str]) -> str:
     )
     label_quotes: Final = tuple(form for form in echoes if _INPUT_VALUE_LABEL in form)
     other_quotes: Final = tuple(form for form in echoes if _INPUT_VALUE_LABEL not in form)
-    return _with_echoes_replaced(_with_input_values_replaced(_with_echoes_replaced(text, label_quotes)), other_quotes)
+    with_label_quotes_replaced: Final = _with_echoes_replaced(text, label_quotes)
+    scrubbed: Final = (
+        None
+        if with_label_quotes_replaced is None
+        else _with_echoes_replaced(_with_input_values_replaced(with_label_quotes_replaced), other_quotes)
+    )
+    return REDACTED_BY_LITELLM if scrubbed is None else scrubbed
 
 
-def _with_echoes_replaced(text: str, echoes: Collection[str]) -> str:
-    """``text`` with one pass that replaces each of ``echoes``, in order, by ``REDACTED_BY_LITELLM``."""
-    if not echoes:
-        return text
-    return re.sub("|".join(_echo_pattern(form) for form in echoes), REDACTED_BY_LITELLM, text)
+def _with_echoes_replaced(text: str, echoes: Collection[str]) -> str | None:
+    """
+    ``text`` with each quote of ``echoes`` replaced by ``REDACTED_BY_LITELLM`` in one left-to-right pass, the
+    longest where several start at one place. ``None`` when ``echoes`` turn up more often than a failure path
+    should look through.
+    """
+    found: Final = tuple(
+        itertools.islice(
+            itertools.chain.from_iterable(_quote_spans(text, form) for form in echoes), _MAX_ECHO_OCCURRENCES + 1
+        )
+    )
+    if len(found) > _MAX_ECHO_OCCURRENCES:
+        return None
+    ends: Final = dict(sorted((span for span in found if _stands_alone(text, *span)), key=lambda span: span[1]))
+    starts: Final = sorted(ends)
+    frontiers: Final = itertools.accumulate(starts, functools.partial(_frontier_after, ends), initial=0)
+    chosen: Final = tuple(start for start, frontier in zip(starts, frontiers) if start >= frontier)
+    resumes: Final = (0, *(ends[start] for start in chosen))
+    kept_and_replaced: Final = itertools.chain.from_iterable(
+        (text[resume:start], REDACTED_BY_LITELLM) for resume, start in zip(resumes, chosen)
+    )
+    return "".join((*kept_and_replaced, text[resumes[-1] :]))
+
+
+def _frontier_after(ends: Mapping[int, int], frontier: int, start: int) -> int:
+    """Where the next replacement may start once the quote at ``start`` is passed, replacing it if it starts at or
+    past ``frontier``."""
+    return ends[start] if start >= frontier else frontier
+
+
+def _quote_spans(text: str, form: str) -> Iterator[tuple[int, int]]:
+    """Where each quote of ``form`` in ``text`` starts and ends, overlaps included. A form short enough to occur in
+    ordinary wording only counts between quotes."""
+    if len(form) >= _MIN_UNQUOTED_TEXT_LENGTH:
+        return ((start, start + len(form)) for start in _found(text, form))
+    quoted_starts: Final = itertools.chain.from_iterable(
+        _found(text, f"{opening}{form}{closing}") for opening, closing in itertools.product(_QUOTES, repeat=2)
+    )
+    return ((start + 1, start + 1 + len(form)) for start in quoted_starts)
+
+
+def _found(text: str, piece: str) -> Iterator[int]:
+    """Every index ``piece`` starts at in ``text``, overlaps included."""
+    return itertools.takewhile(
+        lambda start: start >= 0,
+        itertools.accumulate(
+            itertools.repeat(None), lambda start, _: text.find(piece, start + 1), initial=text.find(piece)
+        ),
+    )
+
+
+def _stands_alone(text: str, start: int, end: int) -> bool:
+    """Whether ``text[start:end]`` is not part of a longer word. An escape such as ``\\n`` before it counts as a break."""
+    return (not _is_word_character(text[start]) or _is_word_break(text, start)) and (
+        not _is_word_character(text[end - 1]) or end == len(text) or not _is_word_character(text[end])
+    )
+
+
+def _is_word_break(text: str, index: int) -> bool:
+    return (
+        index == 0
+        or not _is_word_character(text[index - 1])
+        or _ESCAPE_TAIL.search(text, max(0, index - 10), index) is not None
+    )
+
+
+def _is_word_character(character: str) -> bool:
+    return character.isascii() and (character.isalnum() or character == "_")
 
 
 def _with_input_values_replaced(text: str) -> str:
@@ -525,12 +601,6 @@ def _line_with_input_value_replaced(line: str) -> str:
     return f"{line[:value_start]}{REDACTED_BY_LITELLM}{line[type_start:] if type_start >= 0 else ''}"
 
 
-def _echo_pattern(form: str) -> str:
-    """A quote of ``form``. A form short enough to occur in ordinary wording only counts between quotes."""
-    escaped: Final = re.escape(form)
-    return escaped if len(form) >= _MIN_UNQUOTED_TEXT_LENGTH else f"(?<=['\"]){escaped}(?=['\"])"
-
-
 def scrubbed_error_information(
     error_information: StandardLoggingPayloadErrorInformation, texts: Collection[str]
 ) -> StandardLoggingPayloadErrorInformation:
@@ -588,6 +658,41 @@ def scrubbed_exception_text(model_call_details: Mapping[str, object]) -> tuple[s
     return (message if isinstance(message, str) else _exception_text(logged_traceback_text), logged_traceback_text)
 
 
+def scrubbed_exception_record(model_call_details: Mapping[str, object]) -> tuple[str, str] | None:
+    """
+    The message and stack trace to record for the failure's exception when redaction scrubbed the request's text
+    out of the logged error, as ``scrubbed_exception_text`` reads them. A failure that logged no traceback gets the
+    exception's own frames, with ``REDACTED_BY_LITELLM`` for the text of each exception it chains.
+    """
+    scrubbed: Final = scrubbed_exception_text(model_call_details)
+    exception: Final = model_call_details.get("exception")
+    if scrubbed is None or scrubbed[1] or not isinstance(exception, BaseException):
+        return scrubbed
+    return scrubbed[0], _exception_stack(exception, scrubbed[0], frozenset())
+
+
+def _exception_stack(exception: BaseException, message: str, newer: frozenset[int]) -> str:
+    """``exception``'s traceback as Python prints it, with ``message`` for its text and ``REDACTED_BY_LITELLM`` for
+    the text of each exception it chains. ``newer`` holds the exceptions already printed after it."""
+    printed: Final = newer | {id(exception)}
+    caused: Final = exception.__cause__ is not None
+    chained: Final = (
+        exception.__cause__ if caused else None if exception.__suppress_context__ else exception.__context__
+    )
+    earlier: Final = (
+        ""
+        if chained is None or id(chained) in printed
+        else _exception_stack(chained, REDACTED_BY_LITELLM, printed) + (_DIRECT_CAUSE_LINE if caused else _CONTEXT_LINE)
+    )
+    frames: Final = "".join(traceback.format_tb(exception.__traceback__))
+    header: Final = "Traceback (most recent call last):\n" if frames else ""
+    kind: Final = type(exception)
+    name: Final = (
+        kind.__qualname__ if kind.__module__ in ("__main__", "builtins") else f"{kind.__module__}.{kind.__qualname__}"
+    )
+    return f"{earlier}{header}{frames}{name}: {message}\n" if message else f"{earlier}{header}{frames}{name}\n"
+
+
 def _exception_text(traceback_text: str) -> str:
     """The lines after the last indented frame line of ``traceback_text``."""
     lines: Final = reversed(traceback_text.rstrip().splitlines())

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -11,6 +11,7 @@ import importlib
 import json
 import os
 import threading
+import traceback
 import tracemalloc
 from collections.abc import AsyncIterator, Callable, Mapping
 from datetime import datetime
@@ -37,6 +38,7 @@ from litellm.litellm_core_utils.redact_messages import (
     request_texts,
     scrub_request_text,
     scrubbed_error_information_for_request,
+    scrubbed_exception_record,
     scrubbed_exception_text,
     should_redact_failed_request,
     should_redact_message_logging,
@@ -2444,6 +2446,71 @@ def test_scrub_request_text_replaces_a_short_text_only_between_quotes() -> None:
     assert scrubbed == "Invalid value 'redacted-by-litellm' and \"redacted-by-litellm\"; 432 is too short"
 
 
+@pytest.mark.parametrize(
+    ("prompt", "error"),
+    [
+        ("Error", "litellm.RateLimitError: Rate limit reached"),
+        ("test", "The model `gpt-4-latest` does not exist"),
+        ("text", "litellm.ContextWindowExceededError: too long"),
+        ("test", "No deployment named test_model"),
+    ],
+    ids=["word_end", "word_start", "word_middle", "before_underscore"],
+)
+def test_scrub_request_text_leaves_a_prompt_that_only_occurs_inside_a_longer_word(prompt: str, error: str) -> None:
+    assert scrub_request_text(error, request_texts(prompt)) == error
+
+
+@pytest.mark.parametrize(
+    ("error", "scrubbed"),
+    [
+        ("Error", "redacted-by-litellm"),
+        ("Invalid content: Error.", "Invalid content: redacted-by-litellm."),
+        ('content: "first\\nError"', 'content: "first\\nredacted-by-litellm"'),
+        ('content: "\\u00e9Error"', 'content: "\\u00e9redacted-by-litellm"'),
+        ("content: 'caf\\xe9Error'", "content: 'caf\\xe9redacted-by-litellm'"),
+        ("内容错误Error请重试", "内容错误redacted-by-litellm请重试"),
+    ],
+    ids=["whole_text", "whole_word", "after_newline_escape", "after_unicode_escape", "after_hex_escape", "non_ascii"],
+)
+def test_scrub_request_text_replaces_a_prompt_that_stands_as_its_own_word(error: str, scrubbed: str) -> None:
+    assert scrub_request_text(error, request_texts("Error")) == scrubbed
+
+
+def test_scrub_request_text_replaces_the_longest_prompt_that_starts_at_one_place() -> None:
+    scrubbed: Final = scrub_request_text("bad: secret code list", request_texts("secret", "secret code", "code list"))
+
+    assert scrubbed == "bad: redacted-by-litellm list"
+
+
+@pytest.mark.parametrize(
+    ("echoes", "scrubbed"),
+    [(10_000, "bad: redacted-by-litellm " * 10_000), (10_001, "redacted-by-litellm")],
+    ids=["at_the_limit", "past_the_limit"],
+)
+def test_scrub_request_text_replaces_the_whole_text_when_the_prompt_repeats_past_the_limit(
+    echoes: int, scrubbed: str
+) -> None:
+    assert scrub_request_text("bad: secret " * echoes, request_texts("secret")) == scrubbed
+
+
+def test_scrub_request_text_keeps_no_copy_of_the_prompt_after_it_returns() -> None:
+    prompt: Final = "".join(f"secret {index} " for index in range(20_000))
+    error: Final = f"Invalid content: '{prompt}'"
+    already_tracing: Final = tracemalloc.is_tracing()
+    if not already_tracing:
+        tracemalloc.start()
+    baseline, _ = tracemalloc.get_traced_memory()
+    try:
+        scrubbed: Final = scrub_request_text(error, request_texts(prompt))
+        kept, _ = tracemalloc.get_traced_memory()
+    finally:
+        if not already_tracing:
+            tracemalloc.stop()
+
+    assert scrubbed == "Invalid content: 'redacted-by-litellm'"
+    assert kept - baseline < len(prompt) // 2, f"scrubbing kept {kept - baseline} bytes alive after returning"
+
+
 def test_scrub_request_text_replaces_the_whole_text_when_the_search_is_too_large() -> None:
     texts: Final = request_texts([f"message number {index}" for index in range(5_000)])
 
@@ -2572,6 +2639,76 @@ def test_scrubbed_exception_text_reads_the_message_off_the_traceback_without_a_l
     )
 
 
+class _ProviderError(Exception):
+    pass
+
+
+def _raise_chained(chain: tuple[BaseException, ...]) -> None:
+    """Raise the last of ``chain`` while handling the one before it, and so on down to the first."""
+    if len(chain) == 1:
+        raise chain[0]
+    try:
+        _raise_chained(chain[:-1])
+    except BaseException as earlier:
+        raise chain[-1] from earlier
+
+
+def _raised(chain: tuple[BaseException, ...]) -> BaseException:
+    with pytest.raises(type(chain[-1])) as raised:
+        _raise_chained(chain)
+    return raised.value
+
+
+def _failure(exception: BaseException, error_information: Mapping[str, str]) -> dict[str, object]:
+    return {"exception": exception, "standard_logging_object": {"error_information": error_information}}
+
+
+SCRUBBED_ERROR: Final = {"error_message": "bad: redacted-by-litellm", "traceback": ""}
+
+
+def test_scrubbed_exception_record_prints_the_exceptions_own_frames_without_a_logged_traceback() -> None:
+    exception: Final = _raised(
+        (ValueError("redacted-by-litellm"), KeyError("redacted-by-litellm"), _ProviderError("bad: redacted-by-litellm"))
+    )
+    printed: Final = "".join(traceback.format_exception(exception))
+    scrubbed: Final = scrubbed_exception_record(_failure(exception, SCRUBBED_ERROR))
+
+    assert scrubbed == (
+        "bad: redacted-by-litellm",
+        printed.replace("KeyError: 'redacted-by-litellm'", "KeyError: redacted-by-litellm"),
+    )
+    assert scrubbed is not None and "in _raise_chained" in scrubbed[1]
+
+
+def test_scrubbed_exception_record_replaces_the_text_of_each_exception_the_failure_chains() -> None:
+    exception: Final = _raised((ValueError(PROMPT), _ProviderError("bad: redacted-by-litellm")))
+    scrubbed: Final = scrubbed_exception_record(_failure(exception, SCRUBBED_ERROR))
+
+    assert scrubbed is not None and "123-45-6789" not in scrubbed[1]
+    assert "ValueError: redacted-by-litellm\n\nThe above exception was the direct cause" in scrubbed[1]
+    assert scrubbed[1].endswith("_ProviderError: bad: redacted-by-litellm\n")
+
+
+def test_scrubbed_exception_record_stops_at_an_exception_chained_back_to_itself() -> None:
+    first: Final = ValueError("redacted-by-litellm")
+    second: Final = _ProviderError("bad: redacted-by-litellm")
+    first.__context__ = second
+    second.__context__ = first
+
+    assert scrubbed_exception_record(_failure(second, SCRUBBED_ERROR)) == (
+        "bad: redacted-by-litellm",
+        "".join(traceback.format_exception(second)),
+    )
+
+
+def test_scrubbed_exception_record_keeps_a_logged_traceback_and_skips_an_unscrubbed_error() -> None:
+    exception: Final = _raised((_ProviderError(f"bad: {PROMPT}"),))
+    logged: Final = {"error_message": "bad: redacted-by-litellm", "traceback": "tb redacted-by-litellm"}
+
+    assert scrubbed_exception_record(_failure(exception, logged)) == (logged["error_message"], logged["traceback"])
+    assert scrubbed_exception_record(_failure(exception, {"error_message": "bad", "traceback": ""})) is None
+
+
 def test_scrubbed_error_information_for_request_scrubs_only_when_the_request_redacts() -> None:
     error_information: Final = StandardLoggingPayloadErrorInformation(
         error_code="400", error_class="BadRequestError", error_message=f"bad: {PROMPT}", traceback=f"tb {PROMPT}"