diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a9ae6da58da..5a42ee6c049 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -116,6 +116,7 @@ from litellm.litellm_core_utils.logging_worker import optional_callback_budget from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request from litellm.litellm_core_utils.redact_messages import ( + perform_redaction, redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redact_streaming_responses_for_custom_logger, @@ -3585,6 +3586,7 @@ class Logging(LiteLLMLoggingBaseClass): with fail_after(optional_callback_budget(LOGGING_WORKER_MAX_TIME_PER_COROUTINE, fraction=0.5)): return await self._run_async_success_hooks(callbacks, result) except TimeoutError: + from litellm.constants import X_LITELLM_DISABLE_CALLBACKS from litellm.proxy.spend_tracking.spend_event import ( SpendEventBuildError, build_spend_event_payload, @@ -3592,6 +3594,10 @@ class Logging(LiteLLMLoggingBaseClass): ) self.shadow_eval_request_snapshot = None + self.model_call_details = redact_streaming_responses_for_custom_logger( + self.model_call_details, CustomLogger(message_logging=False) + ) + perform_redaction(self.model_call_details, None) assembled_stream: Final = "async_complete_streaming_response" in self.model_call_details accounting_response: Final[object] = ( self.model_call_details["async_complete_streaming_response"] if assembled_stream else result @@ -3603,8 +3609,24 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.error("Cannot prepare body-free spend accounting after a logging hook timeout") return event args: Final = spend_event_callback_args(event) + mapping_adapter: Final = TypeAdapter(Mapping[str, object]) + empty_headers: Final = MappingProxyType[str, object]({}) + params: Final = mapping_adapter.validate_python(self.model_call_details["litellm_params"]) + request_data: Final = params.get("proxy_server_request") + request: Final = ( + mapping_adapter.validate_python(request_data) if request_data is not None else empty_headers + ) + header_data: Final = request.get("headers") + headers: Final = mapping_adapter.validate_python(header_data) if header_data is not None else empty_headers self.model_call_details = { **args.kwargs, + "litellm_params": { + **args.kwargs["litellm_params"], + "no-log": params.get("no-log", False), + "proxy_server_request": { + "headers": {X_LITELLM_DISABLE_CALLBACKS: headers.get(X_LITELLM_DISABLE_CALLBACKS)} + }, + }, **({"async_complete_streaming_response": args.response_obj} if assembled_stream else {}), } verbose_logger.warning("Logging hooks timed out, recording spend without request or response bodies") diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 25d5394c426..0a41585129b 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -13162,3 +13162,94 @@ async def test_invalid_accounting_after_a_hook_timeout_never_exports_unmasked_da assert "Cannot prepare body-free spend accounting after a logging hook timeout" in caplog.text finally: await worker.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("no_log", [False, True]) +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_masking_timeout_preserves_opt_outs_and_redacts_hidden_provider_content( + monkeypatch: pytest.MonkeyPatch, + no_log: bool, + metadata_field: str, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress + + recorded: Final[asyncio.Queue[str]] = asyncio.Queue() + exported: Final[asyncio.Queue[str]] = asyncio.Queue() + never: Final = asyncio.Event() + + async def store_event(line: bytes) -> None: + recorded.put_nowait(line.decode()) + + producer: Final = SpendEventProducer( + address=UnixAddress("/unused"), + on_unavailable="fallback", + buffer_size=1, + connect_timeout=1, + fallback=store_event, + ) + await producer.close(drain_timeout=1) + + class MaskingLogger(CustomLogger): + async def async_logging_hook( + self, + kwargs: dict[str, object], + result: object, + call_type: str, + ) -> tuple[dict[str, object], object]: + await never.wait() + return kwargs, result + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + exported.put_nowait(json.dumps(kwargs, default=str) + str(response_obj)) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "global_disable_no_log_param", False) + monkeypatch.setattr(litellm, "allow_dynamic_callback_disabling", True) + monkeypatch.setattr(proxy_server, "premium_user", True) + logging_obj: Final = _async_request_logging_obj( + "private-provider-content", [MaskingLogger(), _ProxyDBLogger(producer)] + ) + logging_obj.update_environment_variables( + litellm_params={ + "acompletion": True, + "no-log": no_log, + metadata_field: { + "hidden_params": { + "provider_specific_fields": { + "reasoning_content": "private-thinking-canary", + "thinking_blocks": [{"type": "thinking", "thinking": "private-thinking-canary"}], + "citations": [{"cited_text": "private-citation-canary"}], + } + } + }, + "proxy_server_request": {"headers": {"x-litellm-disable-callbacks": "dynamodb"}}, + }, + optional_params={}, + ) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse())) + try: + await worker.flush() + assert recorded.qsize() == int(not no_log) + assert exported.qsize() == int(not no_log) + payloads: Final = tuple(recorded.get_nowait() for _ in range(recorded.qsize())) + tuple( + exported.get_nowait() for _ in range(exported.qsize()) + ) + assert all( + "private-thinking-canary" not in payload and "private-citation-canary" not in payload + for payload in payloads + ) + assert not logging_obj.should_run_callback( + "dynamodb", logging_obj.model_call_details["litellm_params"], "async_success" + ) + finally: + await worker.stop()