diff --git a/litellm/constants.py b/litellm/constants.py index c17e776f54b..dd82a750cd2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -170,6 +170,7 @@ 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") +FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index dfe508f3630..c87d0d3de47 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -298,6 +298,7 @@ class CustomGuardrail(CustomLogger): use_native_lifecycle_hooks: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + masks_logged_request: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None logging_only_continue_on_input_failure: bool diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..18a6db07f7c 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, 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 924603dc7e8..18a573ea434 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2,9 +2,11 @@ ## Common Utility file for Logging handler # Logging function -> log the exact model details + what's being sent | Non-Blocking import asyncio +import concurrent.futures import copy import datetime import functools +import itertools import json import math import os @@ -43,6 +45,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, + FAILURE_LOG_MASK_TIMEOUT_SECONDS, INTERNAL_CALLBACK_NAMES, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -70,6 +73,7 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.langtrace import langtrace_trace_endpoint from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.classifier_logging import ( classifier_audit_fields, classifier_input_snapshot, @@ -129,6 +133,7 @@ from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redacted_model_call_details_for_custom_logger, + redacted_request_entries, redacted_response_obj_for_custom_logger, should_redact_message_logging, ) @@ -136,7 +141,7 @@ from litellm.litellm_core_utils.served_output_texts import ( SERVED_OUTPUT_TEXTS_KEY, overlay_served_output_texts, ) -from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.litellm_core_utils.thread_pool_executor import executor, failure_log_mask_executor from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils @@ -144,6 +149,7 @@ from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse from litellm.types.extraction import ExtractionResponse +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.integrations.s3_v2 import S3PartitionGranularity from litellm.types.interactions import ( InteractionsAPIResponse, @@ -572,6 +578,62 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata" _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set() # mutable-ok: keeps orphaned tasks alive + + +_ABSENT: Final = object() +_LOGGED_PROMPT_PATHS: Final = ( + ("messages",), + ("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] != 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, replacement: object +) -> dict[str, object]: # mutable-ok: becomes the Logging object's model_call_details + if raw is None: + return dict(details) + return dict( + functools.reduce( + lambda fields, path: _with_prompt_copy_replaced(fields, path, raw, replacement), + _LOGGED_PROMPT_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) + + _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -626,6 +688,14 @@ def _timestamp_seconds(moment: object) -> float | None: return None +def _masks_synchronously(guardrail: CustomGuardrail) -> bool: + return next(cls for cls in type(guardrail).__mro__ if "logging_hook" in vars(cls)) is not CustomLogger + + +def _runs_on_logging_only(guardrail: CustomGuardrail, data: Mapping[str, object]) -> bool: + return guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.logging_only) is True + + def same_attempt_marker(first: object, second: object) -> bool: return first is second or first == second @@ -654,6 +724,9 @@ class Logging(LiteLLMLoggingBaseClass): baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None _on_failure_payload_built: "Callable[[], None] | None" = None + _failure_log_masked: bool = False + _unmasked_messages: object = None + _failure_logged_messages: object = None @property def _defer_async_logging(self) -> bool: @@ -2702,7 +2775,9 @@ class Logging(LiteLLMLoggingBaseClass): sync_sdk: Final = self.is_sync_litellm_request(litellm_params) passthrough: Final = self.call_type == CallTypes.pass_through.value if sync_sdk and not prefer_async_handlers and not passthrough: - self.failure_handler(exception, traceback_exception) + run_failure_handler: Final = functools.partial(self.failure_handler, exception, traceback_exception) + await self.async_mask_failure_log_request(on_cancelled=run_failure_handler) + run_failure_handler() return try: @@ -3502,28 +3577,20 @@ class Logging(LiteLLMLoggingBaseClass): hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) for callback in hook_callbacks: if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - if ( - callback.should_run_guardrail( - data=call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): + if not _runs_on_logging_only(callback, call_details): continue self.shadow_eval_request_snapshot = None call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict kwargs=call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) elif isinstance(callback, CustomLogger): call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict kwargs=call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) if not for_async_request: self.model_call_details = call_details @@ -4094,22 +4161,14 @@ class Logging(LiteLLMLoggingBaseClass): for callback in hook_callbacks: try: if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): + if not _runs_on_logging_only(callback, self.model_call_details): continue self.shadow_eval_request_snapshot = None self.model_call_details, result = await callback.async_logging_hook( kwargs=self.model_call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) elif isinstance(callback, CustomLogger): result = redact_message_input_output_from_custom_logger( @@ -4118,7 +4177,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details, result = await callback.async_logging_hook( kwargs=self.model_call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) verbose_logger.error( @@ -4295,6 +4354,180 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) + @property + def _logging_hook_call_type(self) -> str: + return cast(str, self.call_type) # cast-ok: Logging.__init__ takes call_type untyped and callers pass names + + def _async_failure_callbacks(self) -> Sequence[str | Callable[..., object] | CustomLogger]: + return self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_async_failure_callbacks, + global_callbacks=litellm._async_failure_callback, + ) + + def _failure_log_maskers(self) -> tuple[CustomGuardrail, ...]: + callbacks: Final = ( + *self.get_combined_callback_list(self.dynamic_failure_callbacks, litellm.failure_callback), + *self._async_failure_callbacks(), + ) + return tuple( + callback + for callback in dict.fromkeys(callbacks) + if isinstance(callback, CustomGuardrail) and callback.masks_logged_request + ) + + @property + def messages_sent_to_model(self) -> object: + return self._unmasked_messages if self._failure_log_masked else self.model_call_details.get("messages") + + def failure_logged_proxy_server_request(self, proxy_server_request: Mapping[str, object]) -> Mapping[str, object]: + if not self._failure_log_masked: + return proxy_server_request + return _with_prompt_copy_replaced( + proxy_server_request, ("body", "messages"), self._unmasked_messages, self._failure_logged_messages + ) + + def _claim_failure_log_mask(self) -> bool: + if self._failure_log_masked: + return False + self._failure_log_masked = True + self._unmasked_messages = self.model_call_details.get("messages") + self._failure_logged_messages = self._unmasked_messages + return True + + def _log_failure_with_messages(self, details: Mapping[str, object], logged_messages: object) -> None: + self.model_call_details = _with_prompt_copies_replaced(details, self._unmasked_messages, logged_messages) + self._failure_logged_messages = logged_messages + + def _redact_failure_log_request(self) -> None: + redacted: Final = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + self._log_failure_with_messages(redacted, redacted["messages"]) + + @property + def masks_failure_log(self) -> bool: + return any( + _runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers() + ) + + def redact_failure_log_request(self) -> None: + if self.masks_failure_log and self._claim_failure_log_mask(): + self._redact_failure_log_request() + + def _report_failure_log_mask_error(self, guardrail: CustomGuardrail) -> None: + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred masking the failure log with %s, so its request " + "is redacted %s", + guardrail.guardrail_name, + traceback.format_exc(), + ) + self._handle_callback_failure(callback=guardrail) + + def _redact_failure_log_request_after_timeout(self) -> None: + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Masking the failure log took over %ss, so its request is redacted", + FAILURE_LOG_MASK_TIMEOUT_SECONDS, + ) + self._redact_failure_log_request() + + async def _mask_failure_log_request_with(self, maskers: tuple[CustomGuardrail, ...]) -> None: + try: + await asyncio.wait_for(self._run_failure_log_maskers(maskers), timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + self._redact_failure_log_request_after_timeout() + + async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None: + masked: Final = await self._masked_failure_log_details(self.model_call_details, maskers) + if masked is None: + self._redact_failure_log_request() + return + self._log_failure_with_messages(self.model_call_details, masked.get("messages")) + + async def _masked_failure_log_details( + self, details: Mapping[str, object] | None, maskers: tuple[CustomGuardrail, ...] + ) -> Mapping[str, object] | None: + if details is None or not maskers: + return details + return await self._masked_failure_log_details( + await self._async_failure_log_mask_step(details, maskers[0]), maskers[1:] + ) + + async def _async_failure_log_mask_step( + self, details: Mapping[str, object], guardrail: CustomGuardrail + ) -> Mapping[str, object] | None: + if not _runs_on_logging_only(guardrail, details): + return details + try: + hooked_details: Final = _METADATA_MAPPING.validate_python( + ( + await guardrail.async_logging_hook( + kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type + ) + )[0] + ) + except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging + self._report_failure_log_mask_error(guardrail) + return None + return hooked_details + + async def async_mask_failure_log_request( + self, + on_cancelled: Callable[[], None] | None = None, + on_cancelled_async: Callable[[], Awaitable[None]] | None = None, + ) -> None: + if not self._claim_failure_log_mask(): + return + maskers: Final = self._failure_log_maskers() + if not maskers: + return + mask: Final = asyncio.create_task(self._mask_failure_log_request_with(maskers)) + try: + await asyncio.shield(mask) + except asyncio.CancelledError: + _keep_alive_until_done(mask) + if on_cancelled is not None: + mask.add_done_callback(lambda _: on_cancelled()) + if on_cancelled_async is not None: + mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async()))) + raise + + def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = failure_log_mask_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 = mask_executor.submit( + functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details + ) + try: + masked: Final = masking.result(timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS) + except concurrent.futures.TimeoutError: + _ = masking.cancel() + 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._log_failure_with_messages(self.model_call_details, masked.get("messages")) + + def _sync_failure_log_mask_step( + 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 + if not _masks_synchronously(guardrail): + return None + try: + hooked: Final = guardrail.logging_hook( + kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type + ) + except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging + self._report_failure_log_mask_error(guardrail) + return None + return hooked[0] + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time @@ -4420,6 +4653,7 @@ class Logging(LiteLLMLoggingBaseClass): is_sync_request: Final = self.is_sync_litellm_request(litellm_params) try: + self.mask_failure_log_request() start_time, end_time = self._failure_handler_helper_fn( exception=exception, traceback_exception=traceback_exception, @@ -4619,6 +4853,8 @@ class Logging(LiteLLMLoggingBaseClass): await self.special_failure_handlers(exception=exception) if not self.should_run_logging(event_type="async_failure"): # prevent double logging return + self.mark_logging_complete(event_type="async_failure") + await self.async_mask_failure_log_request() start_time, end_time = self._failure_handler_helper_fn( exception=exception, traceback_exception=traceback_exception, @@ -4627,14 +4863,10 @@ class Logging(LiteLLMLoggingBaseClass): ) if self._on_failure_payload_built is not None: self._on_failure_payload_built() - callbacks: Final = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_async_failure_callbacks, - global_callbacks=litellm._async_failure_callback, - ) + callbacks: Final = self._async_failure_callbacks() result: Final = None # result sent to all loggers, init this to None incase it's not created - self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: litellm_params = self.model_call_details.get("litellm_params", {}) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..746a66cf058 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -67,11 +67,11 @@ def redacted_model_call_details_for_custom_logger(model_call_details: dict, cust return model_call_details return { **{key: value for key, value in model_call_details.items() if key != SERVED_OUTPUT_TEXTS_KEY}, - **_redacted_entries_for_opted_out_logger(model_call_details), + **redacted_request_entries(model_call_details), } -def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, object]) -> Mapping[str, object]: +def redacted_request_entries(model_call_details: Mapping[str, object]) -> Mapping[str, object]: optional_params: Final = _as_mapping(model_call_details.get("optional_params")) litellm_params: Final = _as_mapping(model_call_details.get("litellm_params")) standard_logging_object: Final = _as_mapping(model_call_details.get("standard_logging_object")) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 484edd869cb..f60cae2e840 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, @@ -2509,7 +2510,8 @@ class CustomStreamWrapper: """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" if self._failure_logging is None: return - await asyncio.wait((self._failure_logging,), timeout=timeout) + mask_wait: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS if self.logging_obj.masks_failure_log else 0.0 + await asyncio.wait((self._failure_logging,), timeout=None if timeout is None else timeout + mask_wait) def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier diff --git a/litellm/litellm_core_utils/thread_pool_executor.py b/litellm/litellm_core_utils/thread_pool_executor.py index f989f20247f..00614769932 100644 --- a/litellm/litellm_core_utils/thread_pool_executor.py +++ b/litellm/litellm_core_utils/thread_pool_executor.py @@ -80,3 +80,7 @@ executor: Final = BoundedLoggingThreadPoolExecutor( max_workers=MAX_THREADS, max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS, ) +failure_log_mask_executor: Final = BoundedLoggingThreadPoolExecutor( + max_workers=MAX_THREADS, + max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS, +) diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index cf6592a3e58..fadd91f02d9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,7 +7,7 @@ import json import os -from typing import TYPE_CHECKING, Final, Literal, TypedDict +from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypedDict from fastapi import HTTPException @@ -48,6 +48,8 @@ class GuardrailsAIResponsePreCall(TypedDict): class GuardrailsAI(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + def __init__( self, guard_name: str, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 7e2f6bfec33..ef8313123cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -17,7 +17,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime from itertools import accumulate, count, pairwise -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -238,6 +238,7 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -389,7 +390,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async with self._session_lock: if self._http_session is None or self._http_session.closed: self._http_session = aiohttp.ClientSession() - yield self._http_session + session: Final = self._http_session + yield session else: # Background thread/loop -> use loop-bound session cache # This avoids "attached to a different loop" or "no running event loop" errors diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9e237ffcb57..b2c0d75e11e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1708,12 +1708,24 @@ def _failure_fields_to_lift( request_data (it is not serialisable), so the caller merges these fields onto request_data first: the request start and first-handoff instants for duration and preprocessing latency, the call type, recovered or estimated - usage for token counts, and the standard logging object for deployment - attribution on failed-request spend logs.""" + usage for token counts, the standard logging object for deployment + attribution on failed-request spend logs, and the request body with the + prompt the failure log masked.""" _logging_obj: Final = request_data.get("litellm_logging_obj") if _logging_obj is None: return _EMPTY_LIFT _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) + _proxy_server_request: Final = request_data.get("proxy_server_request") + _logged_proxy_server_request: Final = ( + _logging_obj.failure_logged_proxy_server_request(_as_str_mapping(_proxy_server_request)) + if isinstance(_logging_obj, Logging) and isinstance(_proxy_server_request, Mapping) + else None + ) + _sent_model_call_details: Final = ( + {**_model_call_details, "messages": _logging_obj.messages_sent_to_model} + if isinstance(_logging_obj, Logging) + else _model_call_details + ) _first_handoff: Final = _model_call_details.get("first_api_call_start_time") _litellm_params: Final = getattr(_logging_obj, "litellm_params", None) _base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details) @@ -1721,7 +1733,7 @@ def _failure_fields_to_lift( _success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued _billed: Final = provider_billed and not _success_already_logged _usage_to_lift: Final = _failure_usage_to_lift( - model_call_details=_model_call_details, + model_call_details=_sent_model_call_details, request_body=request_data, dispatched=_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, provider_billed=_billed, @@ -1730,7 +1742,7 @@ def _failure_fields_to_lift( dispatched_call_type=dispatched_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) + _estimated_prompt_usage(_sent_model_call_details, request_data, dispatched_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -1747,6 +1759,7 @@ def _failure_fields_to_lift( ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), ("billed_failed_attempts", _earlier_attempts or None), + ("proxy_server_request", _logged_proxy_server_request), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) @@ -4100,8 +4113,9 @@ class ProxyLogging: litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) set_up_by_request: Final = litellm_logging_obj is not None request_fields: Final = cast(Mapping[str, object], request_data) # cast-ok: request bodies are str-keyed JSON + authenticated: Final = error_type != ProxyErrorTypes.auth_error logged_metadata_keys: Final = _proxy_only_error_logged_metadata_keys( - authenticated=error_type != ProxyErrorTypes.auth_error, past_pre_call=set_up_by_request + authenticated=authenticated, past_pre_call=set_up_by_request ) user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( user_api_key_dict=user_api_key_dict @@ -4210,6 +4224,8 @@ class ProxyLogging: input=input, api_key="", ) + if not authenticated: + litellm_logging_obj.redact_failure_log_request() await self._dispatch_proxy_only_failure_handlers( litellm_logging_obj=litellm_logging_obj, diff --git a/litellm/utils.py b/litellm/utils.py index 8fe3fd84fbd..d588014824c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10,6 +10,7 @@ import binascii import contextvars import copy import datetime +import functools import hashlib import inspect import io @@ -2535,12 +2536,16 @@ def client(original_function): end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) - try: - logging_obj.failure_handler( - e, traceback_exception, start_time, end_time, for_async_request=True - ) # DO NOT MAKE THREADED - router retry fallback relies on this! - except Exception as e: - raise e + run_failure_handler: Final = functools.partial( + logging_obj.failure_handler, e, traceback_exception, start_time, end_time, for_async_request=True + ) + await logging_obj.async_mask_failure_log_request( + on_cancelled=run_failure_handler, + on_cancelled_async=functools.partial( + logging_obj.async_failure_handler, e, traceback_exception, start_time, end_time + ), + ) + run_failure_handler() # DO NOT MAKE THREADED - router retry fallback relies on this! try: await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time) except Exception as e: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e952127a32c..8470fa4c5d4 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import concurrent.futures import copy import datetime import importlib.abc @@ -10,11 +11,11 @@ import sys import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object from importlib.machinery import ModuleSpec from types import MappingProxyType, ModuleType -from typing import Final, Literal +from typing import ClassVar, Final, Literal, ParamSpec, TypeVar, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -31,6 +32,7 @@ from litellm._logging import session_id_var, trace_id_var, verbose_logger from litellm._service_logger import ServiceLogging from litellm.caching.caching import DualCache from litellm.constants import ( + FAILURE_LOG_MASK_TIMEOUT_SECONDS, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, LOGGING_WORKER_MAX_TIME_PER_COROUTINE, @@ -38,6 +40,7 @@ from litellm.constants import ( SENTRY_PII_DENYLIST, ) from litellm.cost_calculator import ocr_batch_cost +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -52,12 +55,14 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.llm_response_utils.response_metadata import update_response_metadata 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 from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.rerank import RerankResponse from litellm.types.utils import ( @@ -8775,6 +8780,531 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() +_RAW_PROMPT: Final = "My name is John Smith" +_RAW_LOG: Final = [{"role": "user", "content": _RAW_PROMPT}] +_MASKED_LOG: Final = [{"role": "user", "content": "My name is "}] +_REDACTED_LOG: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] + + +class _NameMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + def __init__(self, event_hook: GuardrailEventHooks = GuardrailEventHooks.logging_only) -> None: + super().__init__(guardrail_name="name-masker", event_hook=event_hook) + self.scans = 0 + + def _mask(self, kwargs: dict[str, object]) -> dict[str, object]: + self.scans += 1 + for message in cast(list[dict[str, str]], kwargs["messages"]): + message["content"] = message["content"].replace("John Smith", "") + return kwargs + + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return self._mask(kwargs), result + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + await asyncio.sleep(0) + return self._mask(kwargs), result + + +class _UnflaggedNameMaskingGuardrail(_NameMaskingGuardrail): + masks_logged_request: ClassVar[bool] = False + + +class _RaisingMaskingGuardrail(_NameMaskingGuardrail): + def _mask(self, kwargs: dict[str, object]) -> dict[str, object]: + raise RuntimeError("masking failed") + + +class _BlockingNameMaskingGuardrail(_NameMaskingGuardrail): + def __init__(self) -> None: + super().__init__() + self.scanning: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.scanning.set() + await self.release.wait() + return self._mask(kwargs), result + + +class _HangingSyncMaskingGuardrail(_NameMaskingGuardrail): + def __init__(self) -> None: + super().__init__() + self.release: Final = threading.Event() + + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + _ = self.release.wait(timeout=30) + return self._mask(kwargs), result + + +class _AsyncOnlyNameMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return await _NameMaskingGuardrail().async_logging_hook(kwargs, result, call_type) + + +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[object, object]] = [] + + def _record(self, kwargs: dict[str, object]) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"]))) + + def log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + +def _register_failure_logging( + monkeypatch: pytest.MonkeyPatch, guardrail: CustomGuardrail +) -> tuple[_FailureLogRecorder, list[object]]: + recorder: Final = _FailureLogRecorder() + sync_logged: Final[list[object]] = [] + + def record_sync_failure( + kwargs: dict[str, object], completion_response: object, start_time: object, end_time: object + ) -> None: + sync_logged.append(copy.deepcopy(kwargs["messages"])) + + for name in ("input_callback", "success_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder]) + monkeypatch.setattr(litellm, "failure_callback", [record_sync_failure]) + return recorder, sync_logged + + +def _bad_key_error() -> litellm.AuthenticationError: + return litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") + + +async def _until(condition: Callable[[], bool]) -> None: + while not condition(): + await asyncio.sleep(0) + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock(scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, None]]) -> None: + loop: Final = _SkippableClockLoop() + try: + loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +class _OutlastedTimeBoundFuture(concurrent.futures.Future[_T]): + def __init__(self, work: concurrent.futures.Future[_T], on_wait: Callable[[float | None], None]) -> None: + super().__init__() + self._work: Final = work + self._on_wait: Final = on_wait + + def result(self, timeout: float | None = None) -> _T: + self._on_wait(timeout) + raise concurrent.futures.TimeoutError + + def cancel(self) -> bool: + return self._work.cancel() + + +class _TimeBoundOutlastingExecutor(concurrent.futures.Executor): + def __init__(self, workers: concurrent.futures.Executor) -> None: + self._workers: Final = workers + self.waits: tuple[float | None, ...] = () + + def _record_wait(self, timeout: float | None) -> None: + self.waits = (*self.waits, timeout) + + def submit(self, fn: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> concurrent.futures.Future[_T]: + return _OutlastedTimeBoundFuture(self._workers.submit(fn, *args, **kwargs), self._record_wait) + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + self._workers.shutdown(wait=wait, cancel_futures=cancel_futures) + + +async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict[str, str]]) -> None: + error: Final = _bad_key_error() + request: Final = ( + litellm.acompletion(model="gpt-4o-mini", messages=messages, mock_response=error) + if call == "acompletion" + else asyncio.to_thread(litellm.completion, model="gpt-4o-mini", messages=messages, mock_response=error) + ) + with pytest.raises(litellm.AuthenticationError): + await request + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_log_carries_the_prompt_a_logging_only_guardrail_masked( + monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"] +) -> None: + guardrail: Final = _NameMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + messages: Final = copy.deepcopy(_RAW_LOG) + + await _fail_chat_call(call, messages) + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + assert sync_logged == [_MASKED_LOG] + assert messages == _RAW_LOG + assert guardrail.scans == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_guardrail", + [_UnflaggedNameMaskingGuardrail, lambda: _NameMaskingGuardrail(event_hook=GuardrailEventHooks.pre_call)], + ids=["logging_only_without_masking_opt_in", "masking_opt_in_in_pre_call_mode"], +) +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_log_keeps_the_prompt_when_no_logging_only_guardrail_masks_it( + monkeypatch: pytest.MonkeyPatch, + make_guardrail: Callable[[], _NameMaskingGuardrail], + call: Literal["completion", "acompletion"], +) -> None: + guardrail: Final = make_guardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + + await _fail_chat_call(call, copy.deepcopy(_RAW_LOG)) + + assert recorder.logged == [(_RAW_LOG, _RAW_LOG)] + assert sync_logged == [_RAW_LOG] + assert guardrail.scans == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_is_logged_redacted_when_the_masking_guardrail_raises( + monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"] +) -> None: + recorder, sync_logged = _register_failure_logging(monkeypatch, _RaisingMaskingGuardrail()) + + await _fail_chat_call(call, copy.deepcopy(_RAW_LOG)) + + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + assert sync_logged == [_REDACTED_LOG] + + +@pytest.mark.asyncio +async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mid_mask( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + call: Final = asyncio.create_task( + litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()) + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + guardrail.release.set() + await asyncio.wait_for(_until(lambda: bool(sync_logged) and bool(recorder.logged)), timeout=10) + + assert sync_logged == [_MASKED_LOG] + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + +def _sync_failure_logging(call_id: str, guardrail: CustomGuardrail, recorder: _FailureLogRecorder) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id=call_id, + function_id=call_id, + dynamic_failure_callbacks=[guardrail, recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound() -> None: + guardrail: Final = _HangingSyncMaskingGuardrail() + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging("hanging-mask-failure", guardrail, recorder) + mask_executor: Final = _TimeBoundOutlastingExecutor( + BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=1) + ) + + try: + logging_obj.mask_failure_log_request(mask_executor=mask_executor) + finally: + guardrail.release.set() + mask_executor.shutdown(wait=True) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert mask_executor.waits == (FAILURE_LOG_MASK_TIMEOUT_SECONDS,) + 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 = _sync_failure_logging("full-executor-failure", _NameMaskingGuardrail(), recorder) + + 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)] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_is_logged_redacted_when_the_masking_guardrail_has_only_an_async_hook() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging( + "async-only-mask-failure", + _AsyncOnlyNameMaskingGuardrail( + guardrail_name="async-only-name-masker", event_hook=GuardrailEventHooks.logging_only + ), + recorder, + ) + + logging_obj.mask_failure_log_request() + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_mask_still_queued_at_its_time_bound_never_runs() -> None: + guardrail: Final = _NameMaskingGuardrail() + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging("queued-mask-failure", guardrail, recorder) + mask_executor: Final = _TimeBoundOutlastingExecutor( + BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2) + ) + busy_worker_released: Final = threading.Event() + _ = mask_executor.submit(busy_worker_released.wait, 30) + + try: + logging_obj.mask_failure_log_request(mask_executor=mask_executor) + finally: + busy_worker_released.set() + mask_executor.shutdown(wait=True) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert guardrail.scans == 0 + 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(2026, 1, 1), + 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 + + +async def _fail_once_the_mask_outlasts_its_time_bound( + clock: _SkippableClockLoop, guardrail: _BlockingNameMaskingGuardrail, recorder: _FailureLogRecorder +) -> None: + call: Final = asyncio.create_task( + litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()) + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + clock.skip(FAILURE_LOG_MASK_TIMEOUT_SECONDS) + with pytest.raises(litellm.AuthenticationError): + await asyncio.wait_for(call, timeout=10) + await asyncio.wait_for(_until(lambda: bool(recorder.logged)), timeout=10) + + +def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + + _run_on_skippable_clock(lambda clock: _fail_once_the_mask_outlasts_its_time_bound(clock, guardrail, recorder)) + + assert sync_logged == [_REDACTED_LOG] + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_concurrent_failure_notifications_log_one_masked_failure() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="concurrent-masked-failure", + function_id="concurrent-masked-failure", + dynamic_async_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + error: Final = litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini") + + await asyncio.gather(*(logging_obj.async_failure_handler(error, "tb") for _ in range(3))) + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + +class _DeploymentSuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[str, str, str]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((payload["model_id"] or "", payload["model_group"] or "", payload["api_base"])) + + +def _deployment_params(model_id: str, model_group: str) -> dict[str, object]: + return { + "acompletion": True, + "api_base": f"https://{model_group}.example.com", + "metadata": {"model_group": model_group, "model_info": {"id": model_id}}, + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_fallback_started_while_the_failure_log_is_masked_logs_its_own_deployment() -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + failures: Final = _FailureLogRecorder() + successes: Final = _DeploymentSuccessRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="fallback-during-mask", + function_id="fallback-during-mask", + dynamic_async_failure_callbacks=[guardrail, failures], + dynamic_async_success_callbacks=[successes], + ) + logging_obj.update_environment_variables(litellm_params=_deployment_params("rejected", "primary"), optional_params={}) + rejected_attempt_log: Final = asyncio.create_task( + logging_obj.async_failure_handler(litellm.RateLimitError("429", "openai", "gpt-4o-mini"), "tb") + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + + logging_obj.update_environment_variables(litellm_params=_deployment_params("served", "backup"), optional_params={}) + guardrail.release.set() + await asyncio.wait_for(rejected_attempt_log, timeout=10) + await logging_obj.async_success_handler( + result=ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + ) + ) + + assert failures.logged == [(_MASKED_LOG, _MASKED_LOG)] + assert successes.logged == [("served", "backup", "https://backup.example.com")] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_sync_sdk_failure_dispatched_from_async_code_logs_the_masked_prompt() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=True, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="sync-sdk-masked-failure", + function_id="sync-sdk-masked-failure", + dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await logging_obj.dispatch_failure_handlers(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 4d1699c7fe4..b4e0edba67f 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -4,15 +4,16 @@ import itertools import json import os import queue +import selectors import subprocess import sys import threading import time -from collections.abc import Mapping, Sequence +from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType -from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast +from typing import AsyncIterator, Callable, ClassVar, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -23,7 +24,9 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var +from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging @@ -35,6 +38,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( ChoiceLogprobs, CompletionTokensDetailsWrapper, @@ -3478,7 +3482,9 @@ def _stuck_stream_wrapper( provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], call_type: Literal["completion", "acompletion"], recorder: _StuckStreamFailureRecorder, + masker: CustomGuardrail | None = None, ) -> CustomStreamWrapper: + failure_callbacks: Final[list[CustomLogger]] = [recorder] if masker is None else [masker, recorder] logging_obj: Final = Logging( model="test-model", messages=[{"role": "user", "content": "hi"}], @@ -3487,8 +3493,8 @@ def _stuck_stream_wrapper( start_time=datetime.datetime(2026, 1, 1), litellm_call_id=f"stuck-stream-{call_type}", function_id=f"stuck-stream-{call_type}", - dynamic_failure_callbacks=[recorder], - dynamic_async_failure_callbacks=[recorder], + dynamic_failure_callbacks=failure_callbacks, + dynamic_async_failure_callbacks=failure_callbacks, ) logging_obj.update_environment_variables( model="test-model", @@ -3503,6 +3509,83 @@ def _stuck_stream_wrapper( ) +class _BlockedStuckStreamFailureRecorder(_StuckStreamFailureRecorder): + def __init__(self, release: asyncio.Event) -> None: + super().__init__() + self.release = release + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + await self.release.wait() + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + + +class _FailureLogMasker(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return kwargs, result + + +class _VirtualClockSelector(selectors.DefaultSelector): + now: float = 0.0 + + def select(self, timeout: float | None = None) -> list[tuple[selectors.SelectorKey, int]]: + if timeout is None: + return super().select(None) + ready: Final = super().select(0) + if not ready: + self.now += timeout + return ready + + +class _VirtualClockLoop(asyncio.SelectorEventLoop): + def __init__(self) -> None: + self.clock: Final = _VirtualClockSelector() + super().__init__(self.clock) + + def time(self) -> float: + return self.clock.now + + +def _on_virtual_clock(scenario: Coroutine[object, object, float]) -> float: + loop: Final = _VirtualClockLoop() + try: + return loop.run_until_complete(scenario) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _seconds_waited_for_a_blocked_failure_log(masker: CustomGuardrail | None) -> float: + release: Final = asyncio.Event() + wrapper: Final = _stuck_stream_wrapper( + _StuckProvider().async_stream(), "acompletion", _BlockedStuckStreamFailureRecorder(release), masker + ) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + clock: Final = asyncio.get_running_loop() + started: Final = clock.time() + await wrapper.wait_for_failure_logging(timeout=0.1) + waited: Final = clock.time() - started + release.set() + await wrapper.wait_for_failure_logging(timeout=None) + return waited + + +def test_mid_stream_failure_wait_extends_by_the_mask_time_only_when_the_failure_log_is_masked() -> None: + assert _on_virtual_clock(_seconds_waited_for_a_blocked_failure_log(None)) == pytest.approx(0.1) + assert _on_virtual_clock( + _seconds_waited_for_a_blocked_failure_log( + _FailureLogMasker(guardrail_name="failure-log-masker", event_hook=GuardrailEventHooks.logging_only) + ) + ) == pytest.approx(0.1 + FAILURE_LOG_MASK_TIMEOUT_SECONDS) + + @pytest.mark.asyncio async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" diff --git a/tests/unit/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py index 084a0fef9db..8c5d3fd45be 100644 --- a/tests/unit/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -646,6 +646,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): mock_logging_obj.update_environment_variables = MagicMock() mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() mock_logging_obj.async_failure_handler = AsyncMock() + mock_logging_obj.async_mask_failure_log_request = AsyncMock() with ( patch( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py index c6a7fbd7cab..608acc01865 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py @@ -227,3 +227,43 @@ async def test_guardrails_ai_process_input(): # Should keep original content when no output fields are present assert result["messages"][0]["content"] == "Original message" + + +@pytest.mark.asyncio +async def test_guardrails_ai_logging_only_masks_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.guardrails.guardrail_hooks.guardrails_ai.guardrails_ai import GuardrailsAIResponse + from litellm.types.guardrails import GuardrailEventHooks + + logged: list[object] = [] + + class _FailureLogRecorder(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + logged.append(kwargs["messages"]) + + guardrail = GuardrailsAI( + guardrail_name="pii-guard", + guard_name="pii-guard", + api_base="http://test.example.com", + event_hook=GuardrailEventHooks.logging_only, + ) + for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", [guardrail, _FailureLogRecorder()]) + + with patch.object( + guardrail, + "make_guardrails_ai_api_request", + AsyncMock(return_value=GuardrailsAIResponse(validatedOutput="My name is ")), + ): + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "My name is John Smith"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + + assert logged == [[{"role": "user", "content": "My name is "}]] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index e312378969a..b6edd273d20 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -11,7 +11,7 @@ import re from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager -from typing import Final, Literal +from typing import Final, Literal, cast from unittest.mock import MagicMock, patch import aiohttp @@ -26,6 +26,7 @@ import litellm from litellm import mock_completion from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.presidio import ( OPTIONAL_PresidioPIIMasking, @@ -33,7 +34,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import ( ) from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.proxy.guardrails.guardrail_hooks.presidio import PresidioAnalyzeRequest, PresidioAnalyzeResponseItem -from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices +from litellm.types.utils import Choices, Delta, Message, ModelResponse, StandardLoggingPayload, StreamingChoices def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""): @@ -4425,6 +4426,45 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use assert later[3]["content"] == "Now compare against too." +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[object, object]] = [] + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"]))) + + +async def test_logging_only_presidio_masks_the_prompt_in_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed call is logged with the prompt a logging_only Presidio guardrail masked, the same as a + successful one. The analyzer and anonymizer are the in-process fake handed over through api_base.""" + recorder: Final = _FailureLogRecorder() + for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + async with TestServer(_fake_presidio_app()) as server: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base=str(server.make_url("/")), + presidio_anonymizer_api_base=str(server.make_url("/")), + pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK}, + logging_only=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder]) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "My name is John Smith"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + await guardrail._close_http_session() + + masked: Final = [{"role": "user", "content": "My name is "}] + assert recorder.logged == [(masked, masked)] + + @pytest.mark.asyncio @pytest.mark.parametrize("surface", ["mcp_arguments", "mcp_result", "llm_output"]) @pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK]) @@ -4857,11 +4897,17 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]: ) +_PENDING_PRESIDIO_REPLIES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: keeps pending replies alive + + class _InMemoryPresidioTransport(asyncio.Transport): - def __init__(self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None: + def __init__( + self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int + ) -> None: super().__init__() self._protocol: Final = protocol self._analyzed: Final = analyzed + self._replies_wait_for: Final = replies_wait_for self._received = b"" self._closing = False @@ -4881,7 +4927,12 @@ class _InMemoryPresidioTransport(asyncio.Transport): + f"Content-Length: {len(reply)}\r\n\r\n".encode() + reply ) - asyncio.get_running_loop().call_soon(self._protocol.data_received, response) + _PENDING_PRESIDIO_REPLIES.add(asyncio.get_running_loop().create_task(self._reply_once_analyzed(response))) + + async def _reply_once_analyzed(self, response: bytes) -> None: + while self._analyzed.qsize() < self._replies_wait_for: + await asyncio.sleep(0.01) + self._protocol.data_received(response) def _reply(self, path: str, body: bytes) -> tuple[PresidioAnalyzeResponseItem, ...] | _PresidioAnonymizeReply: if path == "/analyze": @@ -4922,15 +4973,16 @@ class _InMemoryPresidioTransport(asyncio.Transport): class _InMemoryPresidioConnector(aiohttp.BaseConnector): - def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None: + def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int) -> None: super().__init__() self._analyzed: Final = analyzed + self._replies_wait_for: Final = replies_wait_for async def _create_connection( # pyright: ignore[reportImplicitOverride] # aiohttp's connector extension point self, req: aiohttp.ClientRequest, traces: Sequence[object], timeout: aiohttp.ClientTimeout ) -> ResponseHandler: protocol: Final = ResponseHandler(asyncio.get_running_loop()) - protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed)) + protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed, self._replies_wait_for)) return protocol @@ -4939,17 +4991,37 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna def _guardrail_with_in_memory_presidio( - analyzed: asyncio.Queue[PresidioAnalyzeRequest], + analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int = 0 ) -> OPTIONAL_PresidioPIIMasking: guardrail: Final = OPTIONAL_PresidioPIIMasking( pii_entities_config=_BLOCK_CARD_MASK_EMAIL, presidio_analyzer_api_base="http://presidio-analyzer.test/", presidio_anonymizer_api_base="http://presidio-anonymizer.test/", ) - guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed)) + guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed, replies_wait_for)) return guardrail +@pytest.mark.asyncio +async def test_presidio_scans_every_message_of_a_request_at_once() -> None: + texts: Final = (f"My email is {_EMAIL}", f"Write to {_EMAIL} today") + analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() + guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, replies_wait_for=len(texts)) + try: + _ = await asyncio.wait_for( + guardrail.async_logging_hook( + kwargs={"messages": [{"role": "user", "content": text} for text in texts]}, + result=None, + call_type="acompletion", + ), + timeout=5, + ) + finally: + await guardrail._close_http_session() + + assert sorted(request.get("text") or "" for request in _drain(analyzed)) == sorted(texts) + + @pytest.mark.asyncio async def test_check_pii_raises_blocked_entity_for_card() -> None: text: Final = f"My credit card number is {_CARD_NUMBER} and my email is {_EMAIL}" diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 0a483840336..292a2e09627 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -4,7 +4,7 @@ import json import smtplib import traceback from collections.abc import Mapping -from typing import Final +from typing import ClassVar, Final import httpx import pytest @@ -3379,6 +3379,36 @@ def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object: return logging_obj +class _LongPlaceholderMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return {**kwargs, "messages": [{"role": "user", "content": " " * 40}]}, result + + +@pytest.mark.asyncio +async def test_a_billed_failure_is_estimated_on_the_prompt_sent_not_the_masked_failure_log(monkeypatch): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + + def billed_prompt_tokens(logging_obj: object) -> int: + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) + return lifted["combined_usage_object"].prompt_tokens + + unmasked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={})) + masked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={})) + monkeypatch.setattr( + litellm, + "_async_failure_callback", + [_LongPlaceholderMaskingGuardrail(guardrail_name="masker", event_hook=GuardrailEventHooks.logging_only)], + ) + await masked.async_mask_failure_log_request() + + assert masked.model_call_details["messages"] != ONE_USER_MESSAGE + assert billed_prompt_tokens(masked) == billed_prompt_tokens(unmasked) + + def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch): """``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat completions) next to the router-stamped ``metadata.model_info.id``; the failure row must diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 4386d5aec94..72b24189116 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -8,18 +8,20 @@ import json from collections.abc import Callable from datetime import datetime from types import MappingProxyType -from typing import Final +from typing import ClassVar, Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException import litellm -from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY, REDACTED_BY_LITELLM from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CachingDetails @@ -500,6 +502,181 @@ async def test_proxy_only_failure_log_keeps_caller_tags_only_after_auth( assert recorded[0]["standard_logging_object"]["request_tags"] == logged_tags +class _PromptMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + def __init__(self) -> None: + super().__init__(guardrail_name="prompt-masker", event_hook=GuardrailEventHooks.logging_only) + self.scans = 0 + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.scans += 1 + return {**kwargs, "messages": [{"role": "user", "content": "My name is "}]}, result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error_type", "masker_scans", "logged_content"), + [(ProxyErrorTypes.auth_error, 0, REDACTED_BY_LITELLM), (None, 1, "My name is ")], +) +async def test_proxy_only_failure_log_masks_after_auth_and_redacts_before_it( + proxy_logging, + monkeypatch: pytest.MonkeyPatch, + error_type: ProxyErrorTypes | None, + masker_scans: int, + logged_content: str, +) -> None: + from litellm.proxy import proxy_server + + recorded: Final[list[dict[str, object]]] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + recorded.append(kwargs) + + masker: Final = _PromptMaskingGuardrail() + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [masker, _RecordingLogger()]) + proxy_logging.alert_types = [] + + await proxy_logging.post_call_failure_hook( + request_data={"model": "internal-model", "messages": [{"role": "user", "content": "My name is John Smith"}]}, + original_exception=HTTPException(status_code=401, detail="Invalid proxy server token passed"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + error_type=error_type, + route="/chat/completions", + ) + + logged_messages: Final = [{"role": "user", "content": logged_content}] + assert masker.scans == masker_scans + assert [(kwargs["messages"], kwargs["standard_logging_object"]["messages"]) for kwargs in recorded] == [ + (logged_messages, logged_messages) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error_type", "logged_content"), + [(ProxyErrorTypes.auth_error, REDACTED_BY_LITELLM), (None, "My name is ")], +) +async def test_failure_hooks_get_the_request_body_with_the_prompt_the_failure_log_masked( + proxy_logging, monkeypatch: pytest.MonkeyPatch, error_type: ProxyErrorTypes | None, logged_content: str +) -> None: + from litellm.proxy import proxy_server + + received_bodies: Final[list[object]] = [] + + class _SpendLogLikeHook(CustomLogger): + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + received_bodies.append(request_data["proxy_server_request"]) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [_PromptMaskingGuardrail(), _SpendLogLikeHook()]) + proxy_logging.alert_types = [] + messages: Final = [{"role": "user", "content": "My name is John Smith"}] + + await proxy_logging.post_call_failure_hook( + request_data={ + "model": "internal-model", + "messages": messages, + "proxy_server_request": { + "url": "http://localhost/chat/completions", + "body": {"model": "internal-model", "messages": list(messages)}, + }, + }, + original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + error_type=error_type, + route="/chat/completions", + ) + + assert received_bodies == [ + { + "url": "http://localhost/chat/completions", + "body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]}, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("maskers", "logged_content"), + [((), "My name is John Smith"), ((_PromptMaskingGuardrail(),), "My name is ")], + ids=["without_a_masker", "with_a_masker"], +) +async def test_failure_hooks_get_the_request_body_the_failure_log_carried_after_the_router_retried_the_call( + proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, maskers: tuple[CustomGuardrail, ...], logged_content: str +) -> None: + from litellm.proxy import proxy_server + + received_bodies: Final[list[object]] = [] + + class _SpendLogLikeHook(CustomLogger): + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + received_bodies.append(request_data["proxy_server_request"]) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [*maskers, _SpendLogLikeHook()]) + proxy_logging.alert_types = [] + messages: Final = [{"role": "user", "content": "My name is John Smith"}] + router: Final = litellm.Router( + model_list=[ + { + "model_name": "internal-model", + "litellm_params": { + "model": "openai/gpt-4.1", + "api_key": "sk-test", + "mock_response": "litellm.RateLimitError", + }, + } + ], + num_retries=1, + retry_after=0, + ) + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime(2026, 1, 1), + model="internal-model", + messages=messages, + litellm_call_id="retried-call", + ) + with pytest.raises(litellm.RateLimitError): + await router.acompletion(model="internal-model", messages=messages, litellm_logging_obj=logging_obj) + + await proxy_logging.post_call_failure_hook( + request_data={ + "model": "internal-model", + "messages": messages, + "litellm_logging_obj": logging_obj, + "proxy_server_request": {"body": {"model": "internal-model", "messages": messages}}, + }, + original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + route="/chat/completions", + ) + + assert received_bodies == [ + {"body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]}} + ] + + @pytest.mark.asyncio async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 9af92e7f9ca..bc475eba095 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -52,6 +52,7 @@ def _make_logging_obj( logging_obj.get_chat_completion_prompt.return_value = prompt_return logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return) logging_obj.async_failure_handler = AsyncMock() + logging_obj.async_mask_failure_log_request = AsyncMock() logging_obj.model_call_details = {} return logging_obj