feat(logging): add helpers that scrub a request's own text out of failure error text

litellm/litellm_core_utils/redact_messages.py

@@ -10,11 +10,14 @@
 import asyncio
 import copy
 import inspect
-from collections.abc import Mapping, Sequence
+import itertools
+import json
+import re
+from collections.abc import Collection, Mapping, Sequence
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any, Final
 
-from pydantic import JsonValue, TypeAdapter, ValidationError
+from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
 
 import litellm
 from litellm.constants import REDACTED_BY_LITELLM
@@ -23,6 +26,9 @@ from litellm.litellm_core_utils.classifier_logging import without_classifier_aud
 from litellm.litellm_core_utils.core_helpers import (
     get_metadata_variable_name_from_kwargs,
 )
+from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
+    initialize_standard_callback_dynamic_params,
+)
 from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
 from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items
 from litellm.llms.vertex_ai.common_utils import (
@@ -31,7 +37,7 @@ from litellm.llms.vertex_ai.common_utils import (
 )
 from litellm.secret_managers.main import str_to_bool
 from litellm.types.router import BaselineRouteStamp
-from litellm.types.utils import StandardCallbackDynamicParams
+from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayloadErrorInformation
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import (
@@ -222,6 +228,10 @@ def _as_mapping(value: object) -> Mapping[str, object] | None:
     return value if isinstance(value, Mapping) else None  # pyright: ignore[reportUnknownVariableType]  # keys are read one by one
 
 
+def _as_items(value: object) -> tuple[object, ...] | None:
+    return tuple(value) if isinstance(value, (list, tuple)) else None  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # items are walked one by one
+
+
 def _redact_provider_specific_fields(message: object) -> None:
     if (fields := _as_mapping(getattr(message, "provider_specific_fields", None))) is not None:
         setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields))  # noqa: B010  # typed as object
@@ -360,6 +370,187 @@ def _redacted_responses_api_response(
     return _with_redacted_instructions(redacted) if redact_instructions else redacted
 
 
+_REQUEST_TEXT_KEYS: Final = ("messages", "input", "prompt", "system", "instructions")
+_STRUCTURAL_KEYS: Final = frozenset({"role", "type", "id", "tool_call_id", "name"})
+_MIN_UNQUOTED_TEXT_LENGTH: Final = 4
+_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}}}(?=['\"])")
+_MAX_SCRUB_COMPARISONS: Final = 10_000_000
+
+
+def _child_values(value: object) -> tuple[object, ...]:
+    if (mapping := _as_mapping(value)) is not None:
+        return tuple(item for key, item in mapping.items() if key not in _STRUCTURAL_KEYS)
+    if (items := _as_items(value)) is not None:
+        return items
+    if isinstance(value, BaseModel):
+        return (value.model_dump(),)
+    return ()
+
+
+def _next_level(level: tuple[object, ...], _: None) -> tuple[object, ...]:
+    return tuple(itertools.chain.from_iterable(map(_child_values, level)))
+
+
+def request_texts(*values: object) -> frozenset[str]:
+    """
+    The non-empty strings a request carries under ``values``, skipping structural fields such as ``role``
+    and ``type``. A value passed more than once is walked once.
+    """
+    distinct_values: Final = tuple({id(value): value for value in values}.values())
+    levels: Final = itertools.takewhile(
+        bool, itertools.accumulate(itertools.repeat(None), _next_level, initial=distinct_values)
+    )
+    return frozenset(value for value in itertools.chain.from_iterable(levels) if isinstance(value, str) and value)
+
+
+def failure_request_texts(model_call_details: Mapping[str, object], *originals: object) -> frozenset[str]:
+    """The request text a failure on this call can quote: its messages, input, prompt, system prompt and instructions."""
+    optional_params: Final = _as_mapping(model_call_details.get("optional_params")) or {}
+    return request_texts(
+        *originals,
+        *(model_call_details.get(key) for key in _REQUEST_TEXT_KEYS),
+        *(optional_params.get(key) for key in ("system", "instructions")),
+    )
+
+
+def _echo_forms(text: str) -> frozenset[str]:
+    """How an error can quote ``text``: verbatim, or JSON or repr escaped. Printable ASCII without a double quote
+    or a backslash escapes to itself."""
+    if text.isascii() and text.isprintable() and not any(character in text for character in _ESCAPED_CHARACTERS):
+        return frozenset((text,))
+    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
+    pydantic shows a long input value, with its pattern anchored to the quote and the cut.
+    """
+    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))
+    return (
+        *(
+            (head, f"(?<=['\"]){re.escape(head)}(?=\\.\\.\\.)")
+            for head in heads
+            if any(form.startswith(head) for form in forms)
+        ),
+        *(
+            (tail, f"(?<=\\.\\.\\.){re.escape(tail)}(?=['\"])")
+            for tail in tails
+            if any(form.endswith(tail) for form in forms)
+        ),
+    )
+
+
+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
+    error that does not echo the prompt keeps its wording. A search too large for a failure path
+    replaces the whole text.
+    """
+    if not text or not texts:
+        return text
+    if len(text) * len(texts) > _MAX_SCRUB_COMPARISONS:
+        return REDACTED_BY_LITELLM
+    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 | cut_forms)),
+        key=lambda echo: len(echo[0]),
+        reverse=True,
+    )
+    if not echoes:
+        return text
+    return re.sub("|".join(pattern for _, pattern in echoes), REDACTED_BY_LITELLM, text)
+
+
+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:
+    """``error_information`` with the request's ``texts`` scrubbed out of its message and traceback. The type,
+    status code and every other field carry no prompt and are left as they are."""
+    error_message: Final = error_information.get("error_message")
+    traceback_text: Final = error_information.get("traceback")
+    return {
+        **error_information,
+        "error_message": error_message if error_message is None else scrub_request_text(error_message, texts),
+        "traceback": traceback_text if traceback_text is None else scrub_request_text(traceback_text, texts),
+    }
+
+
+def _with_scrubbed_text(
+    mapping: Mapping[str, object], fields: tuple[str, ...], texts: Collection[str]
+) -> Mapping[str, object]:
+    return {
+        key: scrub_request_text(value, texts) if key in fields and isinstance(value, str) else value
+        for key, value in mapping.items()
+    }
+
+
+def scrubbed_error_fields(payload: Mapping[str, object], texts: Collection[str]) -> Mapping[str, object]:
+    """The ``error_str`` and ``error_information`` entries of a ``StandardLoggingPayload`` with the request's
+    ``texts`` scrubbed out, to merge over a copy of the payload. Absent entries are left out."""
+    error_str: Final = payload.get("error_str")
+    error_information: Final = _as_mapping(payload.get("error_information"))
+    return {
+        **({"error_str": scrub_request_text(error_str, texts)} if isinstance(error_str, str) else {}),
+        **(
+            {"error_information": _with_scrubbed_text(error_information, ("error_message", "traceback"), texts)}
+            if error_information is not None
+            else {}
+        ),
+    }
+
+
+def scrubbed_exception_text(model_call_details: Mapping[str, object]) -> tuple[str, str] | None:
+    """
+    The error message and traceback a failure logged when redaction scrubbed the request's text out of
+    either, for callbacks that would otherwise record the raw exception. ``None`` when nothing was scrubbed.
+    Without a logged error message, the message is the exception text that ends the traceback.
+    """
+    payload: Final = _as_mapping(model_call_details.get("standard_logging_object")) or {}
+    error_information: Final = _as_mapping(payload.get("error_information")) or {}
+    message: Final = error_information.get("error_message")
+    logged_traceback: Final = model_call_details.get("traceback_exception")
+    traceback_text: Final = (
+        logged_traceback if isinstance(logged_traceback, str) else error_information.get("traceback")
+    )
+    if not any(isinstance(text, str) and REDACTED_BY_LITELLM in text for text in (message, traceback_text)):
+        return None
+    logged_traceback_text: Final = traceback_text if isinstance(traceback_text, str) else ""
+    return (message if isinstance(message, str) else _exception_text(logged_traceback_text), logged_traceback_text)
+
+
+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())
+    return "\n".join(reversed(tuple(itertools.takewhile(lambda line: not line.startswith(" "), lines))))
+
+
 def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:
     """
     Return a copy of a ``StandardLoggingPayload`` with its messages and response redacted.
@@ -643,6 +834,48 @@ def should_redact_message_logging(model_call_details: dict) -> bool:
     return litellm.turn_off_message_logging is True
 
 
+def should_redact_failure_logging(model_call_details: Mapping[str, object]) -> bool:
+    """
+    ``should_redact_message_logging`` for the failure path, which must log the failure whatever the
+    request carried. Headers that are not a mapping carry no redaction header, so the request's own
+    setting or the global one decides.
+    """
+    try:
+        return should_redact_message_logging(dict(model_call_details))
+    except AttributeError:
+        dynamic_turn_off: Final = _get_turn_off_message_logging_from_dynamic_params(dict(model_call_details))
+        return dynamic_turn_off if dynamic_turn_off is not None else litellm.turn_off_message_logging is True
+
+
+def should_redact_failed_request(request_data: Mapping[str, object]) -> bool:
+    """
+    The redaction decision for proxy failure hooks, which get the request body but not its logging
+    object: the same metadata, headers and dynamic params the logging object would read. When the
+    logging object would reject the request's callback params, the global setting decides.
+    """
+    body: Final = dict(request_data)
+    try:
+        dynamic_params: Final = initialize_standard_callback_dynamic_params(body)
+    except ValueError:
+        return litellm.turn_off_message_logging is True
+    return should_redact_failure_logging(
+        {
+            "litellm_params": {key: body[key] for key in ("metadata", "litellm_metadata") if key in body},
+            "standard_callback_dynamic_params": dynamic_params,
+        }
+    )
+
+
+def scrubbed_error_information_for_request(
+    error_information: StandardLoggingPayloadErrorInformation, request_data: Mapping[str, object]
+) -> StandardLoggingPayloadErrorInformation:
+    """``error_information`` built from a proxy failure hook's exception, with the request's own text
+    scrubbed out when the request's message logging is redacted."""
+    if not should_redact_failed_request(request_data):
+        return error_information
+    return scrubbed_error_information(error_information, failure_request_texts(request_data))
+
+
 def redact_message_input_output_from_logging(model_call_details: dict, result, input: object | None = None) -> Any:
     """
     Removes messages, prompts, input, response from logging. This modifies the data in-place

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -11,7 +11,8 @@ import importlib
 import json
 import os
 import threading
-from collections.abc import AsyncIterator, Mapping
+import tracemalloc
+from collections.abc import AsyncIterator, Callable, Mapping
 from datetime import datetime
 from types import MappingProxyType, SimpleNamespace
 from typing import Final, Optional, Union
@@ -20,7 +21,7 @@ from unittest.mock import patch
 import httpx
 import pytest
 import pytest_asyncio
-from pydantic import JsonValue
+from pydantic import JsonValue, TypeAdapter, ValidationError
 
 import litellm
 from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
@@ -33,6 +34,11 @@ from litellm.litellm_core_utils.redact_messages import (
     redacted_model_call_details_for_custom_logger,
     redacted_response_obj_for_custom_logger,
     redacted_standard_logging_payload,
+    request_texts,
+    scrub_request_text,
+    scrubbed_error_information_for_request,
+    scrubbed_exception_text,
+    should_redact_failed_request,
     should_redact_message_logging,
 )
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@@ -43,6 +49,7 @@ from litellm.types.utils import (
     ModelResponseStream,
     ResponsesAPIResponse,
     StandardLoggingPayload,
+    StandardLoggingPayloadErrorInformation,
     TextCompletionResponse,
 )
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -2075,3 +2082,196 @@ def test_perform_redaction_strips_the_matched_text_from_the_bedrock_guardrail_tr
     assert "US_SSN" in redacted.model_dump_json(), "the finding stays, only its matched text goes"
     assert "123-45-6789" not in str(redacted_dict)
     assert "123-45-6789" in response.model_dump_json(), "the caller's own response is not redacted"
+
+
+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(
+    "quoted",
+    [PROMPT, json.dumps(PROMPT)[1:-1], repr(PROMPT)[1:-1]],
+    ids=["verbatim", "json", "repr"],
+)
+def test_scrub_request_text_replaces_each_way_an_error_quotes_the_prompt(quoted: str) -> None:
+    scrubbed: Final = scrub_request_text(f"Invalid content: <{quoted}> rejected", request_texts(PROMPT))
+
+    assert "123-45-6789" not in scrubbed and "4321" not in scrubbed
+    assert "redacted-by-litellm" in scrubbed
+    assert scrubbed.startswith("Invalid content: <") and scrubbed.endswith("> rejected")
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    ['say "hi" to Bob', "C:\\Users\\bob", "café for Bob", "tab\tfor Bob"],
+    ids=["double_quote", "backslash", "non_ascii", "control_character"],
+)
+@pytest.mark.parametrize("escaped", [json.dumps, repr], ids=["json", "repr"])
+def test_scrub_request_text_replaces_an_escaped_echo_of_a_prompt_that_escaping_changes(
+    prompt: str, escaped: Callable[[str], str]
+) -> None:
+    error: Final = f"Invalid content: <{escaped(prompt)[1:-1]}> rejected"
+
+    assert scrub_request_text(error, request_texts(prompt)) == "Invalid content: <redacted-by-litellm> rejected"
+
+
+@pytest.mark.parametrize(
+    "value",
+    [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(value))
+
+    assert "my ssn" not in scrubbed and "4321" not in scrubbed
+    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"
+
+
+def test_scrub_request_text_replaces_a_short_text_only_between_quotes() -> None:
+    scrubbed: Final = scrub_request_text("Invalid value '432' and \"432\"; 432 is too short", request_texts("432"))
+
+    assert scrubbed == "Invalid value 'redacted-by-litellm' and \"redacted-by-litellm\"; 432 is too short"
+
+
+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)])
+
+    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(
+        [
+            {"role": "assistant", "content": [{"type": "text", "text": "part one text"}]},
+            {"role": "tool", "tool_call_id": "call_123456", "name": "lookup", "content": "tool output"},
+            litellm.Message(role="assistant", content="from a model object"),
+        ]
+    )
+
+    assert texts == {"part one text", "tool output", "from a model object"}
+
+
+@pytest.mark.parametrize(
+    "request_data",
+    [
+        {"turn_off_message_logging": True},
+        {"metadata": {"turn_off_message_logging": True}},
+        {"metadata": {"headers": {"x-litellm-enable-message-redaction": "true"}}},
+    ],
+    ids=["body_flag", "metadata_flag", "enable_header"],
+)
+def test_should_redact_failed_request_follows_the_request_opt_ins(request_data: dict) -> None:
+    assert should_redact_failed_request(request_data) is True
+    assert should_redact_failed_request({}) is False
+
+
+@pytest.mark.parametrize(
+    "request_data",
+    [{"metadata": {"headers": "x"}}, {"turn_off_message_logging": "os.environ/REDACT"}],
+    ids=["non_mapping_headers", "env_reference"],
+)
+@pytest.mark.parametrize("global_redaction", [True, False], ids=["global_on", "global_off"])
+def test_should_redact_failed_request_falls_back_to_the_global_setting_for_an_undecidable_request(
+    monkeypatch: pytest.MonkeyPatch, request_data: dict, global_redaction: bool
+) -> None:
+    monkeypatch.setattr(litellm, "turn_off_message_logging", global_redaction)
+
+    assert should_redact_failed_request(request_data) is global_redaction
+
+
+def test_should_redact_failed_request_reads_the_request_setting_past_non_mapping_headers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "turn_off_message_logging", False)
+
+    assert should_redact_failed_request({"metadata": {"headers": "x"}, "turn_off_message_logging": True}) is True
+
+
+@pytest.mark.parametrize(
+    "error_information",
+    [
+        {"error_message": "bad: redacted-by-litellm", "traceback": "tb"},
+        {"error_message": "bad input", "traceback": "tb redacted-by-litellm"},
+    ],
+    ids=["message", "traceback_only"],
+)
+def test_scrubbed_exception_text_reports_a_scrub_of_either_field(error_information: dict) -> None:
+    logged: Final = scrubbed_exception_text({"standard_logging_object": {"error_information": error_information}})
+
+    assert logged == (error_information["error_message"], error_information["traceback"])
+    assert scrubbed_exception_text({"standard_logging_object": {"error_information": {"error_message": "bad"}}}) is None
+
+
+def test_scrubbed_exception_text_reports_a_scrub_of_the_logged_traceback() -> None:
+    model_call_details: Final = {
+        "traceback_exception": "ValueError: bad: redacted-by-litellm",
+        "standard_logging_object": {"error_information": {"error_message": "outer", "traceback": "frames"}},
+    }
+
+    assert scrubbed_exception_text(model_call_details) == ("outer", "ValueError: bad: redacted-by-litellm")
+
+
+def test_scrubbed_exception_text_reads_the_message_off_the_traceback_without_a_logged_message() -> None:
+    traceback_text: Final = (
+        "Traceback (most recent call last):\n"
+        '  File "main.py", line 1, in completion\n'
+        "    raise BadRequestError(message)\n"
+        "BadRequestError: Invalid content: 'redacted-by-litellm'\n"
+        "Received Messages=['redacted-by-litellm']\n"
+    )
+
+    assert scrubbed_exception_text({"traceback_exception": traceback_text}) == (
+        "BadRequestError: Invalid content: 'redacted-by-litellm'\nReceived Messages=['redacted-by-litellm']",
+        traceback_text,
+    )
+
+
+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}"
+    )
+    messages: Final = [{"role": "user", "content": PROMPT}]
+
+    scrubbed: Final = scrubbed_error_information_for_request(
+        error_information, {"messages": messages, "turn_off_message_logging": True}
+    )
+
+    assert scrubbed == {
+        **error_information,
+        "error_message": "bad: redacted-by-litellm",
+        "traceback": "tb redacted-by-litellm",
+    }
+    assert scrubbed_error_information_for_request(error_information, {"messages": messages}) == error_information