diff --git a/litellm/constants.py b/litellm/constants.py index 515ea843f18..906995a98ae 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -165,8 +165,8 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( # Streams exceeding this duration are terminated with a Timeout error. # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") -MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 +MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS + 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 66d30ab2c74..efc9be4c1bb 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -504,7 +504,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, Any], object]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, object], object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ddcb7db06a2..0354b504790 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6,6 +6,7 @@ import concurrent.futures import copy import datetime import functools +import itertools import json import math import os @@ -573,10 +574,54 @@ _REDACTED_RESULT: Final = TypeAdapter(object) _ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set() # mutable-ok: keeps orphaned tasks alive -def _failure_log_hook_kwargs(details: Mapping[str, Any]) -> dict[str, object]: # mutable-ok: hooks mask in place +_ABSENT: Final = object() +_LOGGED_PROMPT_COPY_PATHS: Final = ( + ("input",), + ("optional_params", "messages"), + ("litellm_params", "proxy_server_request", "body", "messages"), +) + + +def _failure_log_hook_kwargs(details: Mapping[str, object]) -> dict[str, object]: # mutable-ok: hooks mask in place return {**details, "messages": copy.deepcopy(details.get("messages"))} +def _logged_field(container: object, key: str) -> object: + if not isinstance(container, Mapping): + return _ABSENT + return cast(Mapping[str, object], container).get(key, _ABSENT) # cast-ok: logged request dicts are str-keyed JSON + + +def _with_field(inner: object, step: tuple[object, str]) -> object: + container, key = step + return {**cast(Mapping[str, object], container), key: inner} # cast-ok: containers above a found copy are Mappings + + +def _with_prompt_copy_replaced( + details: Mapping[str, object], path: Sequence[str], raw: object, replacement: object +) -> Mapping[str, object]: + containers: Final = tuple(itertools.accumulate(path, _logged_field, initial=details)) + if containers[-1] is not raw: + return details + rebuilt: Final = functools.reduce(_with_field, reversed(tuple(zip(containers, path))), replacement) + return cast(Mapping[str, object], rebuilt) # cast-ok: the outermost rebuilt container is details + + +def _with_prompt_copies_replaced( + details: Mapping[str, object], raw: object +) -> dict[str, object]: # mutable-ok: becomes the Logging object's model_call_details + if raw is None: + return dict(details) + replacement: Final = details.get("messages") + return dict( + functools.reduce( + lambda fields, path: _with_prompt_copy_replaced(fields, path, raw, replacement), + _LOGGED_PROMPT_COPY_PATHS, + details, + ) + ) + + def _keep_alive_until_done(task: "asyncio.Future[None]") -> None: _ORPHANED_FAILURE_LOG_TASKS.add(task) task.add_done_callback(_ORPHANED_FAILURE_LOG_TASKS.discard) @@ -4210,7 +4255,9 @@ class Logging(LiteLLMLoggingBaseClass): return True def _redact_failure_log_request(self) -> None: - self.model_call_details = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + self.model_call_details = _with_prompt_copies_replaced( + {**self.model_call_details, **redacted_request_entries(self.model_call_details)}, self._unmasked_messages + ) def redact_failure_log_request(self) -> None: masks_this_request: Final = any( @@ -4255,6 +4302,7 @@ class Logging(LiteLLMLoggingBaseClass): self._report_failure_log_mask_error(guardrail) self._redact_failure_log_request() return + self.model_call_details = _with_prompt_copies_replaced(self.model_call_details, self._unmasked_messages) async def async_mask_failure_log_request( self, @@ -4277,13 +4325,13 @@ class Logging(LiteLLMLoggingBaseClass): mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async()))) raise - def mask_failure_log_request(self) -> None: + def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = executor) -> None: if is_event_loop_running() or not self._claim_failure_log_mask(): return maskers: Final = self._failure_log_maskers() if not maskers: return - masking: Final = executor.submit( + masking: Final = mask_executor.submit( functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details ) try: @@ -4291,14 +4339,17 @@ class Logging(LiteLLMLoggingBaseClass): except concurrent.futures.TimeoutError: self._redact_failure_log_request_after_timeout() return + except concurrent.futures.CancelledError: + self._redact_failure_log_request() + return if masked is None: self._redact_failure_log_request() return - self.model_call_details = dict(masked) + self.model_call_details = _with_prompt_copies_replaced(masked, self._unmasked_messages) def _sync_failure_log_mask_step( - self, details: Mapping[str, Any] | None, guardrail: CustomGuardrail - ) -> Mapping[str, Any] | None: + self, details: Mapping[str, object] | None, guardrail: CustomGuardrail + ) -> Mapping[str, object] | None: if details is None or not _runs_on_logging_only(guardrail, details): return details try: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 7e8cfdb0a94..98087838667 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.thread_pool_executor import BoundedLoggingThreadPoolExecutor from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck @@ -8551,6 +8552,86 @@ async def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_boun assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="full-executor-failure", + function_id="full-executor-failure", + dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + logging_obj.mask_failure_log_request( + mask_executor=BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=0) + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + +class _ProxyRequestRecorder(_FailureLogRecorder): + def __init__(self) -> None: + super().__init__() + self.request_copies: list[tuple[object, object]] = [] + + def _record(self, kwargs: dict[str, object]) -> None: + super()._record(kwargs) + litellm_params: Final = cast(dict[str, object], kwargs["litellm_params"]) + self.request_copies.append( + copy.deepcopy((kwargs["optional_params"], litellm_params["proxy_server_request"])) + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +@pytest.mark.parametrize("sync", [True, False]) +@pytest.mark.parametrize( + ("guardrail_type", "logged_messages"), + [(_NameMaskingGuardrail, _MASKED_LOG), (_RaisingMaskingGuardrail, _REDACTED_LOG)], +) +async def test_failure_log_masks_the_prompt_copies_a_proxy_rejection_logs( + sync: bool, guardrail_type: type[_NameMaskingGuardrail], logged_messages: list[dict[str, str]] +) -> None: + messages: Final = copy.deepcopy(_RAW_LOG) + guardrail: Final = guardrail_type() + recorder: Final = _ProxyRequestRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=messages, + stream=False, + call_type="completion" if sync else "acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="proxy-rejection-copies", + function_id="proxy-rejection-copies", + dynamic_failure_callbacks=[guardrail, recorder], + dynamic_async_failure_callbacks=[guardrail, recorder], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {"url": "/chat/completions", "body": {"messages": messages}}}, + optional_params={"messages": messages}, + ) + logging_obj.model_call_details["input"] = messages + error: Final = litellm.RateLimitError("429", "openai", "gpt-4o-mini") + + if sync: + await asyncio.to_thread(logging_obj.failure_handler, error, "tb") + else: + await logging_obj.async_failure_handler(error, "tb") + + assert len(recorder.request_copies) == 1 + optional_params, proxy_server_request = recorder.request_copies[0] + assert cast(dict[str, object], optional_params)["messages"] == logged_messages + assert _RAW_PROMPT not in str(proxy_server_request) + assert "/chat/completions" in str(proxy_server_request) + assert messages == _RAW_LOG + + @pytest.mark.asyncio async def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: recorder, sync_logged = _register_failure_logging(monkeypatch, _BlockingNameMaskingGuardrail())