mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(logging): redact nested callback input and provider response fields (#45923)
* fix(logging): redact nested callback input and provider response fields * fix(logging): keep shared request body for per-callback redaction and drop hidden provider fields The deprecated per-callback message_logging=False path ran perform_redaction on the shared model_call_details, which now also redacted the nested proxy_server_request body. The spend log Request column and previous_response_id replay read that body, so one opted-out callback blanked them for everyone. That path now keeps the shared body and hidden params, and redact_streaming_responses_for_custom_logger hands the opted-out callback its own redacted litellm_params. Redacted litellm_params also clear provider_specific_fields copied into metadata.hidden_params (Anthropic thinking text), in litellm_params and in the nested request body, and the nested body redacts system and instructions. * fix(logging): keep replayed instructions and router usage marker under redaction The shared request body keeps system and instructions again, so spend logs store them and previous_response_id follow-ups replay the real instructions. Per-callback copies still redact them. Provider-specific fields in metadata.hidden_params are cleared in place, so the metadata dict keeps its identity and the router's counted-usage marker stays visible to deployment_callback_on_success (no double tpm/rpm count). * fix(logging): never copy shared metadata dicts during redaction The shared redaction pass now leaves litellm_params metadata dicts as they are, so a hidden_params write from the other success handler between the in-place clear and the copy can no longer move the router usage marker to a stale dict. The key-presence predicate change is reverted as unneeded. * refactor(logging): redact nested request body and provider fields with the existing message rules Run perform_redaction on a copy of proxy_server_request.body, the same call the spend log already makes, and redact reasoning inside provider_specific_fields the way the message's own reasoning_content and thinking_blocks are redacted, instead of clearing every provider field. Opaque fields such as Gemini thought_signatures stay, and the metadata dicts and streaming objects are handled as on main. --------- Co-authored-by: moyai-devin-berriai[bot] <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
moyai-devin-berriai[bot]
GitHub
parent
719c44db3f
commit
0b3fcc197d
4 changed files with 728 additions and 7 deletions
|
|
@ -108,6 +108,7 @@ def _redact_choice_content(choice):
|
|||
choice.message.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.message, "thinking_blocks"):
|
||||
choice.message.thinking_blocks = None
|
||||
_redact_provider_specific_fields(choice.message)
|
||||
_redact_tool_calls(getattr(choice.message, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.message, "function_call", None))
|
||||
elif isinstance(choice, litellm.utils.StreamingChoices):
|
||||
|
|
@ -117,10 +118,39 @@ def _redact_choice_content(choice):
|
|||
choice.delta.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.delta, "thinking_blocks"):
|
||||
choice.delta.thinking_blocks = None
|
||||
_redact_provider_specific_fields(choice.delta)
|
||||
_redact_tool_calls(getattr(choice.delta, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.delta, "function_call", None))
|
||||
|
||||
|
||||
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 _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
|
||||
|
||||
|
||||
def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None:
|
||||
for message in (choice.get("message"), choice.get("delta")):
|
||||
if (
|
||||
(readable := _as_mapping(message)) is not None
|
||||
and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None
|
||||
and isinstance(message, dict)
|
||||
):
|
||||
message["provider_specific_fields"] = _redacted_provider_specific_fields(fields)
|
||||
|
||||
|
||||
def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Copy of provider_specific_fields with its reasoning redacted the same way as the message's own."""
|
||||
return {
|
||||
**fields,
|
||||
**({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}),
|
||||
**{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields},
|
||||
}
|
||||
|
||||
|
||||
def _redact_responses_api_output(output_items):
|
||||
"""Helper to redact ResponsesAPIResponse output items."""
|
||||
for output_item in output_items:
|
||||
|
|
@ -228,6 +258,7 @@ def _redact_tool_calls_dict(message: Mapping[str, object]) -> None:
|
|||
def _redact_model_response_dict_choices(choices, redacted_str: str):
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
_redact_provider_specific_fields_dict(choice)
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
if choice["message"].get("content") is not None:
|
||||
choice["message"]["content"] = redacted_str
|
||||
|
|
@ -261,8 +292,37 @@ def _redacted_baseline_metadata(metadata: Mapping[str, object]) -> Mapping[str,
|
|||
return {**metadata, "_autorouter_baseline_route": replace(route, request_parameters=None)}
|
||||
|
||||
|
||||
def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]:
|
||||
request: Final = params.get("proxy_server_request")
|
||||
def _redacted_proxy_server_request(request: Mapping[str, object], *, redact_body: bool) -> Mapping[str, object]:
|
||||
"""Copy of proxy_server_request; redact_body redacts its body the same way as the top-level request."""
|
||||
body: Final = _as_mapping(request.get("body"))
|
||||
if not redact_body or body is None:
|
||||
return without_classifier_audit(request)
|
||||
redacted_body: Final = {**body}
|
||||
perform_redaction(redacted_body, result=None)
|
||||
return {**without_classifier_audit(request), "body": redacted_body}
|
||||
|
||||
|
||||
def _redact_hidden_provider_specific_fields(params: Mapping[str, object]) -> None:
|
||||
"""Redact the provider_specific_fields copied into metadata hidden_params, keeping each metadata dict."""
|
||||
for metadata in (params.get("metadata"), params.get("litellm_metadata")):
|
||||
if (
|
||||
(readable := _as_mapping(metadata)) is not None
|
||||
and (hidden_params := _as_mapping(readable.get("hidden_params"))) is not None
|
||||
and (fields := _as_mapping(hidden_params.get("provider_specific_fields"))) is not None
|
||||
and isinstance(metadata, dict)
|
||||
):
|
||||
metadata["hidden_params"] = {
|
||||
**hidden_params,
|
||||
"provider_specific_fields": _redacted_provider_specific_fields(fields),
|
||||
}
|
||||
|
||||
|
||||
def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body: bool = True) -> dict[str, object]:
|
||||
"""
|
||||
Copy of litellm_params without classifier audit fields or baseline route request parameters.
|
||||
redact_request_body also redacts the nested proxy_server_request body.
|
||||
"""
|
||||
request: Final = _as_mapping(params.get("proxy_server_request"))
|
||||
return {
|
||||
**params,
|
||||
**{
|
||||
|
|
@ -270,7 +330,11 @@ def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]:
|
|||
for key, value in params.items()
|
||||
if key in ("metadata", "litellm_metadata") and isinstance(value, Mapping)
|
||||
},
|
||||
**({"proxy_server_request": without_classifier_audit(request)} if isinstance(request, Mapping) else {}),
|
||||
**(
|
||||
{"proxy_server_request": _redacted_proxy_server_request(request, redact_body=redact_request_body)}
|
||||
if request is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -280,12 +344,17 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
|
||||
redact_streaming_responses=False skips the in-place redaction of the shared streaming
|
||||
response entries; per-callback redaction hands each opted-out callback its own redacted
|
||||
copy via redact_streaming_responses_for_custom_logger instead.
|
||||
copy via redact_streaming_responses_for_custom_logger instead. It also leaves the shared
|
||||
proxy request body and hidden provider fields as they are, since the spend log reads them.
|
||||
"""
|
||||
# Redact model_call_details
|
||||
params: Final = model_call_details.get("litellm_params")
|
||||
if isinstance(params, Mapping):
|
||||
model_call_details["litellm_params"] = redacted_litellm_params(params)
|
||||
params: Final = _as_mapping(model_call_details.get("litellm_params"))
|
||||
if params is not None:
|
||||
if redact_streaming_responses:
|
||||
_redact_hidden_provider_specific_fields(params)
|
||||
model_call_details["litellm_params"] = redacted_litellm_params(
|
||||
params, redact_request_body=redact_streaming_responses
|
||||
)
|
||||
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
|
|
|
|||
47
tests/integration/_support/redaction_probe.py
Normal file
47
tests/integration/_support/redaction_probe.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
def _payload(kwargs: Mapping[str, object], response_obj: object) -> str:
|
||||
params: Final = kwargs.get("litellm_params")
|
||||
request: Final = params.get("proxy_server_request") if isinstance(params, Mapping) else None
|
||||
standard: Final = kwargs.get("standard_logging_object")
|
||||
return json.dumps(
|
||||
{
|
||||
"request": request,
|
||||
"response": response_obj.model_dump() if isinstance(response_obj, BaseModel) else response_obj,
|
||||
"standard_response": standard.get("response") if isinstance(standard, Mapping) else None,
|
||||
},
|
||||
default=str,
|
||||
)
|
||||
|
||||
|
||||
class SyncRedactionProbe(CustomLogger):
|
||||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
with httpx.Client(trust_env=False) as client:
|
||||
response: Final = client.post(os.environ["REDACTION_PROBE_URL"], content=_payload(kwargs, response_obj))
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
class RedactionProbe(CustomLogger):
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
async with httpx.AsyncClient(trust_env=False) as client:
|
||||
response: Final = await client.post(
|
||||
os.environ["REDACTION_PROBE_URL"], content=_payload(kwargs, response_obj)
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
probe: Final = RedactionProbe()
|
||||
485
tests/integration/observability/test_custom_logger_redaction.py
Normal file
485
tests/integration/observability/test_custom_logger_redaction.py
Normal file
|
|
@ -0,0 +1,485 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, held
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.redaction_probe import RedactionProbe, SyncRedactionProbe
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from redis import Redis
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
INPUT: Final = "private-input-redaction-canary"
|
||||
REASONING: Final = "private-reasoning-redaction-canary"
|
||||
OUTPUT: Final = "private-output-redaction-canary"
|
||||
SIGNATURE: Final = "opaque-thought-signature-canary"
|
||||
INSTRUCTIONS: Final = "Reply in French. redaction-instructions-canary"
|
||||
BRIDGED: Final = "bridged-redaction-canary"
|
||||
USAGE_DEPLOYMENT: Final = f"redaction-usage-{uuid.uuid4().hex}"
|
||||
REDACT: Final = {"x-litellm-enable-message-redaction": "true"}
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
|
||||
def _provider(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
body: Final = json.loads(request.body)
|
||||
if request.target == "/v1/embeddings":
|
||||
assert body["input"] == INPUT, body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"object": "list",
|
||||
"model": body["model"],
|
||||
"data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}],
|
||||
"usage": {"prompt_tokens": 10, "total_tokens": 10},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
if request.target == "/v1/completions":
|
||||
assert body["prompt"] == INPUT, body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "cmpl-redaction-probe",
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": body["model"],
|
||||
"choices": [{"text": OUTPUT, "index": 0, "finish_reason": "stop", "logprobs": None}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
if request.target == "/v1/responses":
|
||||
assert body["input"] == INPUT, body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "resp-redaction-probe",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"model": body["model"],
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg-redaction-probe",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": OUTPUT, "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
if request.target == "/v1/messages":
|
||||
if BRIDGED not in request.body.decode():
|
||||
assert body["messages"] == [{"role": "user", "content": INPUT}], body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "msg-redaction-probe",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": body["model"],
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": REASONING, "signature": "synthetic-signature"},
|
||||
{"type": "text", "text": OUTPUT},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
assert body["messages"] == [{"role": "user", "content": INPUT}], body
|
||||
payload: Final = {
|
||||
"id": "chatcmpl-redaction-" + body.get("user", "probe"),
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": OUTPUT,
|
||||
"provider_specific_fields": {
|
||||
"reasoning_content": REASONING,
|
||||
"thinking_blocks": [{"thinking": REASONING}],
|
||||
"reasoningContentBlocks": [{"reasoningText": {"text": REASONING}}],
|
||||
"citations": [{"cited_text": INPUT}],
|
||||
"thought_signatures": [SIGNATURE],
|
||||
},
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
if body.get("stream"):
|
||||
choice: Final = payload["choices"][0]
|
||||
chunk: Final = {
|
||||
**payload,
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{"index": 0, "delta": choice["message"], "finish_reason": None}],
|
||||
}
|
||||
finish: Final = {
|
||||
**payload,
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"data: {json.dumps(event)}\n\n".encode() for event in (chunk, finish))
|
||||
+ (b"data: [DONE]\n\n",),
|
||||
)
|
||||
return Reply(body=json.dumps(payload).encode())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def redaction_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Wire, Wire]]:
|
||||
tmp_path: Final = tmp_path_factory.mktemp("custom-logger-redaction")
|
||||
with (
|
||||
wire_server(_provider) as provider,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = tmp_path / "redaction.yaml"
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_base": provider.url if model.startswith("anthropic/") else provider.url + "/v1",
|
||||
"api_key": "synthetic-provider-key",
|
||||
},
|
||||
}
|
||||
for name, model in (
|
||||
("redaction-probe", "openai/gpt-4o-mini"),
|
||||
("redaction-text", "openai/gpt-3.5-turbo-instruct"),
|
||||
("redaction-embedding", "openai/text-embedding-3-large"),
|
||||
("redaction-responses", "openai/gpt-5.6-sol"),
|
||||
("redaction-messages", "anthropic/claude-sonnet-5-5"),
|
||||
)
|
||||
]
|
||||
+ [
|
||||
{
|
||||
"model_name": "redaction-usage",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-5-5",
|
||||
"api_base": provider.url,
|
||||
"api_key": "synthetic-provider-key",
|
||||
"rpm": 100000,
|
||||
"tpm": 100000000,
|
||||
},
|
||||
"model_info": {"id": USAGE_DEPLOYMENT},
|
||||
}
|
||||
],
|
||||
"litellm_settings": {
|
||||
"callbacks": ["integration._support.redaction_probe.probe"],
|
||||
"enable_redis_auth_cache": True,
|
||||
"cache": True,
|
||||
"cache_params": {
|
||||
"type": "redis",
|
||||
"host": "os.environ/REDIS_HOST",
|
||||
"port": "os.environ/REDIS_PORT",
|
||||
},
|
||||
},
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"disable_model_info_refresh": True,
|
||||
"store_prompts_in_spend_logs": True,
|
||||
},
|
||||
"router_settings": {"redis_host": "os.environ/REDIS_HOST", "redis_port": "os.environ/REDIS_PORT"},
|
||||
}
|
||||
)
|
||||
)
|
||||
with owned_proxy(
|
||||
rig, tmp_path, {"REDACTION_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2
|
||||
) as proxy:
|
||||
yield proxy, provider, sink
|
||||
|
||||
|
||||
@pytest.mark.parametrize("redact", [False, True])
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
def test_custom_logger_redacts_nested_input_and_provider_reasoning(
|
||||
redaction_proxy: tuple[Gateway, Wire, Wire], redact: bool, stream: bool
|
||||
) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "redaction-probe",
|
||||
"messages": [{"role": "user", "content": INPUT}],
|
||||
"stream": stream,
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
headers={"x-litellm-enable-message-redaction": "true"} if redact else {},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
first: Final = (
|
||||
json.loads(next(line[6:] for line in response.text.splitlines() if line.startswith("data: ")))
|
||||
if stream
|
||||
else response.json()
|
||||
)
|
||||
message: Final = first["choices"][0]["delta" if stream else "message"]
|
||||
assert message["content"] == OUTPUT, message
|
||||
assert REASONING in json.dumps(message["provider_specific_fields"]), message
|
||||
received: Final = eventually(sink.drain, lambda values: len(values) == 1, seconds=20)
|
||||
callback: Final = json.loads(received[0].body)
|
||||
body: Final = callback["request"]["body"]
|
||||
assert "proxy_server_request" not in body, body
|
||||
assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1
|
||||
logged: Final = json.dumps(callback)
|
||||
assert SIGNATURE in logged, callback
|
||||
if redact:
|
||||
assert INPUT not in logged and REASONING not in logged and OUTPUT not in logged, callback
|
||||
else:
|
||||
assert INPUT in logged and REASONING in logged and OUTPUT in logged, callback
|
||||
|
||||
|
||||
@pytest.mark.parametrize("redact", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
("route", "field", "model"),
|
||||
[("/v1/completions", "prompt", "redaction-text"), ("/v1/embeddings", "input", "redaction-embedding")],
|
||||
)
|
||||
def test_custom_logger_redacts_prompt_and_embedding_input(
|
||||
redaction_proxy: tuple[Gateway, Wire, Wire], route: str, field: str, model: str, redact: bool
|
||||
) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
route,
|
||||
{"model": model, field: INPUT, "cache": {"no-cache": True}},
|
||||
headers={"x-litellm-enable-message-redaction": "true"} if redact else {},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
if field == "input":
|
||||
assert response.json()["data"][0]["embedding"] == [0.25, 0.5]
|
||||
else:
|
||||
assert response.json()["choices"][0]["text"] == OUTPUT
|
||||
received: Final = eventually(sink.drain, bool, seconds=20)
|
||||
assert len(received) == 1, received
|
||||
callback: Final = json.loads(received[0].body)
|
||||
assert "proxy_server_request" not in callback["request"]["body"]
|
||||
assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1
|
||||
assert (INPUT in json.dumps(callback)) is not redact, callback
|
||||
|
||||
|
||||
@pytest.mark.parametrize("redact", [False, True])
|
||||
@pytest.mark.parametrize("native_messages", [False, True])
|
||||
def test_native_messages_and_responses_logging_redaction(
|
||||
redaction_proxy: tuple[Gateway, Wire, Wire], native_messages: bool, redact: bool
|
||||
) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
body: Final = (
|
||||
{"model": "redaction-messages", "messages": [{"role": "user", "content": INPUT}], "max_tokens": 64}
|
||||
if native_messages
|
||||
else {"model": "redaction-responses", "input": INPUT}
|
||||
)
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/messages" if native_messages else "/v1/responses",
|
||||
{**body, "cache": {"no-cache": True}},
|
||||
headers={"x-litellm-enable-message-redaction": "true"} if redact else {},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
if native_messages:
|
||||
assert response.json()["content"][1]["text"] == OUTPUT
|
||||
assert response.json()["content"][0]["thinking"] == REASONING
|
||||
else:
|
||||
assert response.json()["output"][0]["content"][0]["text"] == OUTPUT
|
||||
received: Final = eventually(sink.drain, bool, seconds=20)
|
||||
assert len(received) == 1, received
|
||||
logged: Final = received[0].body.decode()
|
||||
assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1
|
||||
assert (INPUT in logged) is not redact, logged
|
||||
assert (OUTPUT in logged) is not redact, logged
|
||||
if native_messages:
|
||||
assert (REASONING in logged) is not redact, logged
|
||||
|
||||
|
||||
def test_concurrent_redacted_and_unredacted_callbacks_remain_isolated(
|
||||
redaction_proxy: tuple[Gateway, Wire, Wire],
|
||||
) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
users: Final = tuple(f"redaction-chaos-{index}" for index in range(4))
|
||||
|
||||
def send(index: int):
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "redaction-probe",
|
||||
"messages": [{"role": "user", "content": INPUT}],
|
||||
"user": users[index],
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
headers={"x-litellm-enable-message-redaction": "true"} if index % 2 else {},
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
responses: Final = tuple(pool.map(send, range(4)))
|
||||
for response in responses:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == OUTPUT
|
||||
assert REASONING in json.dumps(response.json()["choices"][0]["message"])
|
||||
batches: Final[list[Request]] = [] # mutable-ok: drain consumes the sink queue, so polls retain previous events
|
||||
|
||||
def observed() -> tuple[Request, ...]:
|
||||
batches.extend(sink.drain())
|
||||
return tuple(batches)
|
||||
|
||||
events: Final = eventually(observed, lambda values: len(values) == 4, seconds=20)
|
||||
callbacks: Final = tuple(json.loads(event.body) for event in events)
|
||||
assert {callback["request"]["body"]["user"] for callback in callbacks} == set(users)
|
||||
for callback in callbacks:
|
||||
index: Final = users.index(callback["request"]["body"]["user"])
|
||||
logged: Final = json.dumps(callback)
|
||||
assert (INPUT in logged) is (index % 2 == 0), callback
|
||||
assert (REASONING in logged) is (index % 2 == 0), callback
|
||||
assert (OUTPUT in logged) is (index % 2 == 0), callback
|
||||
assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 4
|
||||
|
||||
|
||||
def _drained(wire: Wire, count: int) -> tuple[Request, ...]:
|
||||
batches: Final[list[Request]] = [] # mutable-ok: drain consumes the queue, so polls retain previous events
|
||||
|
||||
def observed() -> tuple[Request, ...]:
|
||||
batches.extend(wire.drain())
|
||||
return tuple(batches)
|
||||
|
||||
return eventually(observed, lambda values: len(values) >= count, seconds=30)
|
||||
|
||||
|
||||
def test_redacted_responses_follow_up_replays_original_instructions(
|
||||
redaction_proxy: tuple[Gateway, Wire, Wire],
|
||||
) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
end_user: Final = f"redaction-replay-{uuid.uuid4().hex}"
|
||||
first: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{
|
||||
"model": "redaction-messages",
|
||||
"instructions": INSTRUCTIONS,
|
||||
"input": f"{BRIDGED} turn one",
|
||||
"user": end_user,
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
headers=REDACT,
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
eventually(
|
||||
lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user = %s', (end_user,)),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert len(_drained(provider, 1)) == 1
|
||||
follow_up: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{
|
||||
"model": "redaction-messages",
|
||||
"previous_response_id": first.json()["id"],
|
||||
"input": f"{BRIDGED} turn two",
|
||||
"user": end_user,
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
headers=REDACT,
|
||||
)
|
||||
assert follow_up.status_code == 200, follow_up.text
|
||||
sent: Final = tuple(json.loads(request.body) for request in _drained(provider, 1))
|
||||
assert len(_drained(sink, 2)) == 2
|
||||
assert len(sent) == 1, sent
|
||||
assert len(sent[0]["messages"]) == 3, sent[0]
|
||||
system: Final = json.dumps(sent[0].get("system"))
|
||||
assert INSTRUCTIONS in system and REDACTED_BY_LITELLM not in system, sent[0]
|
||||
|
||||
|
||||
def test_redacted_router_usage_is_counted_once(redaction_proxy: tuple[Gateway, Wire, Wire]) -> None:
|
||||
proxy, provider, sink = redaction_proxy
|
||||
requests: Final = 48
|
||||
|
||||
def send(index: int):
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "redaction-usage",
|
||||
"messages": [{"role": "user", "content": f"{BRIDGED} usage {index}"}],
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
headers=REDACT,
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
responses: Final = tuple(pool.map(send, range(requests)))
|
||||
for response in responses:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["usage"]["total_tokens"] == 15, response.json()
|
||||
|
||||
def counted_tokens() -> int:
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
keys: Final = tuple(cache.scan_iter(f"global_router:{USAGE_DEPLOYMENT}:*:tpm:*"))
|
||||
return sum(int(float(cache.get(key) or 0)) for key in keys)
|
||||
|
||||
_drained(sink, requests)
|
||||
held(counted_tokens, lambda total: total == 15 * requests, holding=3, seconds=30)
|
||||
assert len(_drained(provider, requests)) == requests
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("redact", [False, True])
|
||||
async def test_sdk_success_callback_redaction_over_http(
|
||||
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, redact: bool
|
||||
) -> None:
|
||||
with wire_server(_provider) as provider, wire_server(lambda request: Reply()) as sink:
|
||||
monkeypatch.setenv("REDACTION_PROBE_URL", sink.url)
|
||||
messages: Final = [{"role": "user", "content": INPUT}]
|
||||
request: Final = {"body": {"messages": messages, "model": "gpt-4o-mini"}}
|
||||
callback: Final = RedactionProbe().async_log_success_event if asynchronous else SyncRedactionProbe()
|
||||
parameters: Final = {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"messages": messages,
|
||||
"api_base": provider.url + "/v1",
|
||||
"api_key": "synthetic-sdk-provider-key",
|
||||
"proxy_server_request": request,
|
||||
"turn_off_message_logging": redact,
|
||||
"success_callback": [callback],
|
||||
}
|
||||
response: Final = (
|
||||
await litellm.acompletion(**parameters)
|
||||
if asynchronous
|
||||
else await asyncio.to_thread(litellm.completion, **parameters)
|
||||
)
|
||||
assert response.choices[0].message.content == OUTPUT
|
||||
assert REASONING in json.dumps(response.choices[0].message.provider_specific_fields)
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
delivered: Final = await asyncio.to_thread(eventually, sink.drain, bool, seconds=20)
|
||||
assert len(delivered) == 1, delivered
|
||||
logged: Final = delivered[0].body.decode()
|
||||
assert request["body"]["messages"] == messages
|
||||
if redact:
|
||||
assert INPUT not in logged and REASONING not in logged and OUTPUT not in logged, logged
|
||||
else:
|
||||
assert INPUT in logged and REASONING in logged and OUTPUT in logged, logged
|
||||
|
|
@ -6,6 +6,7 @@ but litellm_params["litellm_metadata"] is None.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -36,12 +37,131 @@ from litellm.responses.main import mock_responses_api_response
|
|||
from litellm.types.router import BaselineRouteStamp
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
ResponsesAPIResponse,
|
||||
StandardLoggingPayload,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
||||
SECRET: Final = "private-redaction-canary"
|
||||
SIGNATURE: Final = "opaque-thought-signature"
|
||||
|
||||
|
||||
def _provider_fields() -> dict[str, object]:
|
||||
return {
|
||||
"reasoning_content": SECRET,
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": SECRET}],
|
||||
"reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}],
|
||||
"citations": [{"cited_text": SECRET}],
|
||||
"thought_signatures": [SIGNATURE],
|
||||
}
|
||||
|
||||
|
||||
def _request_params(metadata: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
"metadata": metadata,
|
||||
"proxy_server_request": {
|
||||
"url": "http://localhost/v1/chat/completions",
|
||||
"body": {
|
||||
"model": "redaction-model",
|
||||
"system": "kept-system",
|
||||
"messages": [{"role": "user", "content": SECRET}],
|
||||
"metadata": metadata,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ["object", "dict_message", "dict_delta"])
|
||||
def test_redaction_reuses_message_rules_for_request_body_and_provider_fields(shape: str) -> None:
|
||||
response: Final = ModelResponse(
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": SECRET, "provider_specific_fields": _provider_fields()},
|
||||
}
|
||||
]
|
||||
)
|
||||
result: Final = (
|
||||
response
|
||||
if shape == "object"
|
||||
else {"choices": [{"index": 0, "delta": {"content": SECRET, "provider_specific_fields": _provider_fields()}}]}
|
||||
if shape == "dict_delta"
|
||||
else response.model_dump()
|
||||
)
|
||||
original_result: Final = copy.deepcopy(result)
|
||||
params: Final = _request_params({"session_id": "session"})
|
||||
original_params: Final = copy.deepcopy(params)
|
||||
details: Final[dict[str, object]] = {"litellm_params": params}
|
||||
|
||||
redacted: Final = perform_redaction(details, result)
|
||||
|
||||
assert SECRET not in str(redacted) and SECRET not in str(details), (redacted, details)
|
||||
assert SIGNATURE in str(redacted)
|
||||
body: Final = details["litellm_params"]["proxy_server_request"]["body"]
|
||||
assert body["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
assert body["system"] == "kept-system" and body["model"] == "redaction-model"
|
||||
assert body["metadata"] is params["metadata"]
|
||||
assert result == original_result and params == original_params
|
||||
|
||||
|
||||
def test_redaction_redacts_reasoning_copied_into_hidden_params_in_place() -> None:
|
||||
hidden_params: Final = {"model_id": "deployment", "provider_specific_fields": _provider_fields()}
|
||||
original_hidden_params: Final = copy.deepcopy(hidden_params)
|
||||
metadata: Final[dict[str, object]] = {"session_id": "session", "hidden_params": hidden_params}
|
||||
params: Final = _request_params(metadata)
|
||||
details: Final[dict[str, object]] = {"litellm_params": params}
|
||||
|
||||
perform_redaction(details, None)
|
||||
|
||||
assert details["litellm_params"]["metadata"] is metadata
|
||||
assert SECRET not in str(details), details
|
||||
assert metadata["hidden_params"]["model_id"] == "deployment"
|
||||
assert metadata["hidden_params"]["provider_specific_fields"]["thought_signatures"] == [SIGNATURE]
|
||||
assert hidden_params == original_hidden_params
|
||||
|
||||
|
||||
def test_per_callback_redaction_leaves_shared_request_body_and_hidden_params() -> None:
|
||||
metadata: Final[dict[str, object]] = {"hidden_params": {"provider_specific_fields": _provider_fields()}}
|
||||
params: Final = _request_params(metadata)
|
||||
original_params: Final = copy.deepcopy(params)
|
||||
details: Final[dict[str, object]] = {"litellm_params": params}
|
||||
|
||||
perform_redaction(details, None, redact_streaming_responses=False)
|
||||
|
||||
assert details["litellm_params"] == original_params
|
||||
|
||||
|
||||
def test_logger_level_redaction_redacts_its_copy_of_the_request_body() -> None:
|
||||
params: Final = _request_params({"session_id": "session"})
|
||||
original_params: Final = copy.deepcopy(params)
|
||||
details: Final[dict[str, object]] = {"litellm_params": params}
|
||||
|
||||
redacted: Final = CustomLogger(
|
||||
turn_off_message_logging=True
|
||||
).redact_standard_logging_payload_from_model_call_details(details)
|
||||
|
||||
assert SECRET not in str(redacted["litellm_params"]), redacted
|
||||
assert details["litellm_params"] == original_params
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["complete_streaming_response", "async_complete_streaming_response"])
|
||||
def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_chunk_dict(key: str) -> None:
|
||||
fields: Final = _provider_fields()
|
||||
response: Final = ModelResponseStream(
|
||||
choices=[
|
||||
litellm.utils.StreamingChoices(delta=litellm.utils.Delta(content=SECRET, provider_specific_fields=fields))
|
||||
]
|
||||
)
|
||||
details: Final[dict[str, object]] = {"stream": True, key: response}
|
||||
|
||||
perform_redaction(details, None)
|
||||
|
||||
assert SECRET not in str(response.choices[0].delta), response
|
||||
assert response.choices[0].delta.provider_specific_fields["thought_signatures"] == [SIGNATURE]
|
||||
assert fields == _provider_fields()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_global_redaction():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue