diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 7c6c39abb76..c2cca5739cf 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -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"] = "" diff --git a/tests/integration/_support/redaction_probe.py b/tests/integration/_support/redaction_probe.py new file mode 100644 index 00000000000..2a344daa73c --- /dev/null +++ b/tests/integration/_support/redaction_probe.py @@ -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() diff --git a/tests/integration/observability/test_custom_logger_redaction.py b/tests/integration/observability/test_custom_logger_redaction.py new file mode 100644 index 00000000000..adc0eb9103c --- /dev/null +++ b/tests/integration/observability/test_custom_logger_redaction.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 6fbcba170f1..d14f236f765 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -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():