fix(logging): wholesale-redact complete_input_dict and scrub rerank/OCR keys in body snapshot

This commit is contained in:
michelligabriele 2026-06-08 23:08:54 +02:00
parent 11da9d0a4b
commit e223a02c8e
No known key found for this signature in database
2 changed files with 162 additions and 23 deletions

View file

@ -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):

View file

@ -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)