This commit is contained in:
quietjules 2026-10-31 06:10:03 +00:00 • committed by GitHub
commit 857fdb079a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 1263 additions and 64 deletions

View file

@ -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.

View file

@ -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

View file

@ -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

View file

@ -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", {})

View file

@ -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"))

View file

@ -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

View file

@ -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,
)

View file

@ -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,

View file

@ -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

View file

@ -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,

View file

@ -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:

View file

@ -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 <PERSON>"}]
_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", "<PERSON>")
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",

View file

@ -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"""

View file

@ -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(

View file

@ -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 <PERSON>")),
):
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 <PERSON>"}]]

View file

@ -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 <PERSON> 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 <PERSON>"}]
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}"

View file

@ -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": "<PERSON> " * 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

View file

@ -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 <PERSON>"}]}, 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 <PERSON>")],
)
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 <PERSON>")],
)
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 <PERSON>")],
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

View file

@ -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