From c85e9afd2a25e6f4c5ccc532924462ff081ce34a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:26:59 -0700 Subject: [PATCH] fix(guardrails): bound the failure-log mask and keep billing on the sent prompt Bound the failure-log mask at 1 s and redact the request when it runs out, so a hung Presidio cannot hold every failed request's error response. Run the sync failure handler when the caller is cancelled during the mask, so router cooldowns still happen. Redact unauthenticated rejections without calling Presidio. Estimate provider-billed failure usage on the prompt sent to the model, not the masked log copy. Mark Guardrails AI as masking the logged request, since its logging hook rewrites the prompt like Presidio's --- litellm/constants.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 69 ++++++++++-- litellm/litellm_core_utils/redact_messages.py | 4 +- .../guardrails_ai/guardrails_ai.py | 4 +- litellm/proxy/utils.py | 14 ++- litellm/utils.py | 13 ++- .../test_litellm_logging.py | 102 ++++++++++++++---- .../guardrails_ai/test_guardrails_ai.py | 40 +++++++ .../guardrail_hooks/test_presidio.py | 13 ++- ..._utils_model_creation_and_error_logging.py | 32 +++++- .../test_post_call_failure_hook.py | 62 ++++++++++- 11 files changed, 303 insertions(+), 51 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 42ad624026c..515ea843f18 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -166,6 +166,7 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 +FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 95814989322..c12a67ed7c0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1,6 +1,7 @@ # What is this? ## Common Utility file for Logging handler # Logging function -> log the exact model details + what's being sent | Non-Blocking +import asyncio import copy import datetime import functools @@ -42,6 +43,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, @@ -126,6 +128,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, ) @@ -566,6 +569,7 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_ORPHANED_FAILURE_LOG_MASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: keeps orphaned masks alive def _provider_response_id(source: object) -> str | None: @@ -648,6 +652,7 @@ class Logging(LiteLLMLoggingBaseClass): baseline_observation: "CapturedBaselineObservation | None" = None _on_failure_payload_built: "Callable[[], None] | None" = None _failure_log_masked: bool = False + _unmasked_messages: object = None @property def _defer_async_logging(self) -> bool: @@ -2599,8 +2604,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: - await self.async_mask_failure_log_request() - 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: @@ -4185,6 +4191,27 @@ class Logging(LiteLLMLoggingBaseClass): def _failure_log_hook_kwargs(self) -> dict[str, object]: # mutable-ok: logging hooks mask a plain dict in place return {**self.model_call_details, "messages": copy.deepcopy(self.model_call_details.get("messages"))} + @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 _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") + return True + + def _redact_failure_log_request(self) -> None: + self.model_call_details = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + + def redact_failure_log_request(self) -> None: + masks_this_request: Final = any( + _runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers() + ) + if masks_this_request 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 %s", @@ -4193,11 +4220,18 @@ class Logging(LiteLLMLoggingBaseClass): ) self._handle_callback_failure(callback=guardrail) - async def async_mask_failure_log_request(self) -> None: - if self._failure_log_masked: - return - self._failure_log_masked = True - for guardrail in self._failure_log_maskers(): + 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: + 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 _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None: + for guardrail in maskers: try: if not _runs_on_logging_only(guardrail, self.model_call_details): continue @@ -4207,10 +4241,25 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: # noqa: BLE001 # a failing mask must not stop failure logging self._report_failure_log_mask_error(guardrail) - def mask_failure_log_request(self) -> None: - if self._failure_log_masked or is_event_loop_running(): + async def async_mask_failure_log_request(self, on_cancelled: Callable[[], 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: + _ORPHANED_FAILURE_LOG_MASKS.add(mask) + mask.add_done_callback(_ORPHANED_FAILURE_LOG_MASKS.discard) + if on_cancelled is not None: + mask.add_done_callback(lambda _: on_cancelled()) + raise + + def mask_failure_log_request(self) -> None: + if is_event_loop_running() or not self._claim_failure_log_mask(): return - self._failure_log_masked = True for guardrail in self._failure_log_maskers(): try: if not _runs_on_logging_only(guardrail, self.model_call_details): diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..f1eae1f82c7 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/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/utils.py b/litellm/proxy/utils.py index 9f58c267231..33ae99361d7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1592,6 +1592,11 @@ def _failure_fields_to_lift( if _logging_obj is None: return _EMPTY_LIFT _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) + _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) @@ -1599,7 +1604,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, @@ -1608,7 +1613,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 ) @@ -3947,8 +3952,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 @@ -4059,6 +4065,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 0366f31baf7..4e792240ca1 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 @@ -2428,13 +2429,11 @@ 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) - await logging_obj.async_mask_failure_log_request() - 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) + 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 06c74bf26f8..5ef152e1f76 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, 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 ClassVar, Final, Literal +from typing import ClassVar, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -8341,6 +8341,7 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur _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): @@ -8350,16 +8351,20 @@ class _NameMaskingGuardrail(CustomGuardrail): super().__init__(guardrail_name="name-masker", event_hook=event_hook) self.scans = 0 - def _mask(self, kwargs: dict) -> dict: + def _mask(self, kwargs: dict[str, object]) -> dict[str, object]: self.scans += 1 - for message in kwargs["messages"]: + 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, result: object, call_type: str) -> tuple[dict, object]: + 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, result: object, call_type: str) -> tuple[dict, object]: + 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 @@ -8369,36 +8374,53 @@ class _UnflaggedNameMaskingGuardrail(_NameMaskingGuardrail): class _RaisingMaskingGuardrail(_NameMaskingGuardrail): - def _mask(self, kwargs: dict) -> dict: + 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 _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.logged: list[tuple[list[dict], list[dict]]] = [] + self.logged: list[tuple[object, object]] = [] - def _record(self, kwargs: dict) -> None: - self.logged.append( - (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) - ) + 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, response_obj: object, start_time: object, end_time: object) -> None: + 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, response_obj: object, start_time: object, end_time: object + 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[list[dict]]]: +) -> tuple[_FailureLogRecorder, list[object]]: recorder: Final = _FailureLogRecorder() - sync_logged: Final[list[list[dict]]] = [] + sync_logged: Final[list[object]] = [] - def record_sync_failure(kwargs: dict, completion_response: object, start_time: object, end_time: object) -> None: + 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"): @@ -8408,8 +8430,17 @@ def _register_failure_logging( return recorder, sync_logged -async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict]) -> None: - error: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") +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) + + +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" @@ -8471,6 +8502,41 @@ async def test_failure_is_still_logged_when_the_masking_guardrail_raises( assert len(sync_logged) == 1 +@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() + _, 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)), timeout=10) + + assert sync_logged == [_MASKED_LOG] + + +@pytest.mark.asyncio +async def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: + recorder, sync_logged = _register_failure_logging(monkeypatch, _BlockingNameMaskingGuardrail()) + + with pytest.raises(litellm.AuthenticationError): + await asyncio.wait_for( + litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()), + timeout=10, + ) + await asyncio.wait_for(_until(lambda: bool(recorder.logged)), timeout=10) + + 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: 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 fc1daf02879..d6fd03f82df 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 @@ -34,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=""): @@ -4429,14 +4429,13 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use class _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.logged: list[tuple[list[dict], list[dict]]] = [] + self.logged: list[tuple[object, object]] = [] async def async_log_failure_event( - self, kwargs: dict, response_obj: object, start_time: object, end_time: object + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self.logged.append( - (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) - ) + 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: 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 f2151572f49..c6e77d979e3 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 @@ -3357,6 +3357,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 f880d3548f7..1571c29afaf 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 @@ -7,18 +7,20 @@ import asyncio 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 @@ -499,6 +501,62 @@ 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 async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch