From e223a02c8ead4b5b0082db846f669f8d1332cfff Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 8 Jun 2026 23:08:54 +0200 Subject: [PATCH] fix(logging): wholesale-redact complete_input_dict and scrub rerank/OCR keys in body snapshot --- litellm/litellm_core_utils/redact_messages.py | 33 +++- .../test_redact_messages.py | 152 +++++++++++++++--- 2 files changed, 162 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 8180c803b28..52a9fc1e978 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -165,16 +165,23 @@ def _redact_standard_logging_object(model_call_details: dict): standard_logging_object["response"] = {"text": redacted_str} -# Input-bearing keys that show up in request-body snapshots -# (proxy_server_request.body and additional_args.complete_input_dict). +# Input-bearing keys that show up in the proxy_server_request.body snapshot. +# This is the LiteLLM-API (OpenAI-compat) request body, so its input keys are +# a finite, stable set; it must stay key-based (not wholesale-redacted) because +# downstream consumers read named keys off it (`user`->Lago billing, +# `tools`->spend-log index, `input`/`messages`->Responses session). # "messages"/"prompt"/"input" are the OpenAI-style payload entry points; # "contents" is the Gemini/Vertex native user-turn field. +# "query"/"documents" are the rerank inputs and "document" the OCR input — +# all top-level keys the proxy preserves verbatim in this snapshot. # "system"/"system_prompt"/"instructions" are the provider-native top-level # system-prompt fields (Anthropic `system`, Responses API `instructions`); # "system_instruction"/"systemInstruction" are the Gemini/Vertex equivalents. # All carry user content just as the messages do. +# (additional_args.complete_input_dict is the provider-native wire body and is +# wholesale-redacted in _redact_additional_args_complete_input_dict instead.) def _redact_request_body_dict(body: dict): - """Scrub the input/system-prompt keys on a materialised request-body dict.""" + """Scrub the input/system-prompt keys on the proxy_server_request body dict.""" if "messages" in body: body["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] if "prompt" in body: @@ -185,6 +192,12 @@ def _redact_request_body_dict(body: dict): body["contents"] = [ {"role": "user", "parts": [{"text": "redacted-by-litellm"}]} ] + if "query" in body: # rerank + body["query"] = "redacted-by-litellm" + if "documents" in body: # rerank + body["documents"] = ["redacted-by-litellm"] + if "document" in body: # OCR + body["document"] = "redacted-by-litellm" for key in ( "system", "system_prompt", @@ -236,7 +249,19 @@ def _redact_additional_args_complete_input_dict(model_call_details: dict): if not isinstance(complete_input_dict, dict): return - _redact_request_body_dict(complete_input_dict) + # complete_input_dict is the provider-NATIVE wire body. User input lands + # under ~20+ provider-specific, often nested key shapes (Vertex embeddings + # `instances`, rerank `query`/`documents`, image `prompt`, audio `file`, + # OCR `document`, Bedrock invoke `inputText`, passthrough arbitrary bodies), + # so a key-allowlist cannot close the class. No consumer reads a named key + # from this dict (the OTel exporter iterates `.items()` generically; logging + # treats it as an opaque curl-command payload), so we wholesale-redact it, + # preserving only the non-input `model` for span/debug usefulness. + redacted: dict = {"redacted-by-litellm": True} + model = complete_input_dict.get("model") + if isinstance(model, str): + redacted["model"] = model + additional_args["complete_input_dict"] = redacted def _redact_model_response_dict_choices(choices, redacted_str: str): diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index aedccb56fce..a4e20493b39 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -504,10 +504,9 @@ class TestPerformRedaction: assert redacted.choices[0].message.content == "redacted-by-litellm" def test_redacts_additional_args_complete_input_dict_messages(self): - """perform_redaction must scrub the provider-native request payload - stashed on additional_args.complete_input_dict. Anthropic-shape - content blocks (list of dicts) and the OpenAI prompt / input keys - must all be wiped.""" + """perform_redaction wholesale-redacts the provider-native request + payload stashed on additional_args.complete_input_dict: every input + key is dropped and only the non-input `model` survives.""" details = { "additional_args": { "complete_input_dict": { @@ -532,9 +531,11 @@ class TestPerformRedaction: perform_redaction(details, None) cid = details["additional_args"]["complete_input_dict"] - assert cid["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] - assert cid["prompt"] == "" - assert cid["input"] == "" + assert cid == {"redacted-by-litellm": True, "model": "claude-3-7-sonnet"} + assert "messages" not in cid + assert "prompt" not in cid + assert "input" not in cid + assert "CANARY_INPUT_should_be_redacted" not in str(cid) def test_redact_additional_args_complete_input_dict_is_safe_when_missing(self): """No KeyError / TypeError when additional_args / complete_input_dict @@ -571,8 +572,8 @@ class TestPerformRedaction: def test_redacts_system_prompt_in_complete_input_dict(self): """Provider-native system-prompt keys (Anthropic `system`, the - `system_prompt` variant, Responses API `instructions`) must be - scrubbed from the wire-format request body too.""" + `system_prompt` variant, Responses API `instructions`) are dropped by + the wholesale redaction of the wire-format request body.""" details = { "additional_args": { "complete_input_dict": { @@ -589,15 +590,17 @@ class TestPerformRedaction: perform_redaction(details, None) cid = details["additional_args"]["complete_input_dict"] - assert cid["system"] == "redacted-by-litellm" - assert cid["system_prompt"] == "redacted-by-litellm" - assert cid["instructions"] == "redacted-by-litellm" + assert cid == {"redacted-by-litellm": True, "model": "claude-3-7-sonnet"} + assert "system" not in cid + assert "system_prompt" not in cid + assert "instructions" not in cid + assert "CANARY_SYSTEM_should_be_redacted" not in str(cid) def test_redacts_gemini_native_fields_in_complete_input_dict(self): """Gemini/Vertex wire-format requests carry the user turn in `contents` and the system prompt in `system_instruction` / `systemInstruction`, - not the OpenAI-style `messages` / `system` keys — these provider-native - fields must be scrubbed from the wire-format request snapshot too.""" + not the OpenAI-style `messages` / `system` keys — the wholesale + redaction drops them all, preserving only `model`.""" details = { "additional_args": { "complete_input_dict": { @@ -621,11 +624,11 @@ class TestPerformRedaction: perform_redaction(details, None) cid = details["additional_args"]["complete_input_dict"] - assert cid["contents"] == [ - {"role": "user", "parts": [{"text": "redacted-by-litellm"}]} - ] - assert cid["system_instruction"] == "redacted-by-litellm" - assert cid["systemInstruction"] == "redacted-by-litellm" + assert cid == {"redacted-by-litellm": True, "model": "gemini-2.0-flash"} + assert "contents" not in cid + assert "system_instruction" not in cid + assert "systemInstruction" not in cid + assert "CANARY" not in str(cid) def test_redacts_gemini_native_fields_in_proxy_server_request_body(self): """Same Gemini/Vertex native input fields can also land in the proxy @@ -656,3 +659,114 @@ class TestPerformRedaction: {"role": "user", "parts": [{"text": "redacted-by-litellm"}]} ] assert body["system_instruction"] == "redacted-by-litellm" + + def test_complete_input_dict_wholesale_redacts_embeddings_instances(self): + """Vertex embeddings carry user input under `instances` — a key the old + allowlist never enumerated. Wholesale redaction drops it regardless.""" + details = { + "additional_args": { + "complete_input_dict": { + "model": "text-embedding-004", + "instances": [{"content": "CANARY_INPUT_should_be_redacted"}], + } + } + } + + perform_redaction(details, None) + + cid = details["additional_args"]["complete_input_dict"] + assert cid == {"redacted-by-litellm": True, "model": "text-embedding-004"} + assert "instances" not in cid + assert "CANARY_INPUT_should_be_redacted" not in str(cid) + + def test_complete_input_dict_wholesale_redacts_rerank_query_documents(self): + """Cohere/Bedrock rerank carry user input under `query` / `documents` — + wholesale redaction of the native body drops both.""" + details = { + "additional_args": { + "complete_input_dict": { + "model": "rerank-english-v3.0", + "query": "CANARY_QUERY_should_be_redacted", + "documents": ["CANARY_DOC_should_be_redacted"], + } + } + } + + perform_redaction(details, None) + + cid = details["additional_args"]["complete_input_dict"] + assert cid == {"redacted-by-litellm": True, "model": "rerank-english-v3.0"} + assert "query" not in cid + assert "documents" not in cid + assert "CANARY" not in str(cid) + + def test_complete_input_dict_wholesale_redacts_passthrough_arbitrary(self): + """An arbitrary/passthrough provider key with no allowlist entry must + still be dropped — proving the redaction has no allowlist dependency.""" + details = { + "additional_args": { + "complete_input_dict": { + "model": "some-provider/some-model", + "totally_unknown_provider_key": "CANARY_should_be_redacted", + } + } + } + + perform_redaction(details, None) + + cid = details["additional_args"]["complete_input_dict"] + assert cid == { + "redacted-by-litellm": True, + "model": "some-provider/some-model", + } + assert "totally_unknown_provider_key" not in cid + assert "CANARY_should_be_redacted" not in str(cid) + + def test_redacts_rerank_keys_in_proxy_server_request_body(self): + """The keyed proxy body snapshot scrubs the rerank input keys + (`query` / `documents`) while leaving non-input keys (`model`, `user`) + intact for downstream consumers.""" + details = { + "litellm_params": { + "proxy_server_request": { + "body": { + "model": "rerank-english-v3.0", + "user": "user-123", + "query": "CANARY_QUERY_should_be_redacted", + "documents": ["CANARY_DOC_should_be_redacted"], + } + } + }, + } + + perform_redaction(details, None) + + body = details["litellm_params"]["proxy_server_request"]["body"] + assert body["query"] == "redacted-by-litellm" + assert body["documents"] == ["redacted-by-litellm"] + assert body["model"] == "rerank-english-v3.0" + assert body["user"] == "user-123" + + def test_redacts_ocr_document_in_proxy_server_request_body(self): + """The OCR route lands the user's `document` as a top-level key in the + proxy body snapshot — it must be scrubbed while `model` survives.""" + details = { + "litellm_params": { + "proxy_server_request": { + "body": { + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "CANARY_DOC_should_be_redacted", + }, + } + } + }, + } + + perform_redaction(details, None) + + body = details["litellm_params"]["proxy_server_request"]["body"] + assert body["document"] == "redacted-by-litellm" + assert body["model"] == "mistral/mistral-ocr-latest" + assert "CANARY_DOC_should_be_redacted" not in str(body)