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