From 519ecb67c848fe2ae3ee5d59cf0c3df69a4f69f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:10:04 -0700 Subject: [PATCH] fix(guardrails): mask the prompt in failure logs for logging_only Presidio A Presidio guardrail in logging_only mode masked PII only in success logs. On a failed call every failure logger got the raw prompt in kwargs["messages"] and in standard_logging_object["messages"] Guardrails that mask the logged request now opt in with masks_logged_request (Presidio does). The failure path runs their logging hook once per call on a copy of the messages, before the standard logging payload is built, so the caller's messages stay untouched for retries and fallbacks --- litellm/integrations/custom_guardrail.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 105 +++++++--- .../guardrails/guardrail_hooks/presidio.py | 3 +- litellm/utils.py | 1 + .../test_litellm_logging.py | 180 +++++++++++++++++- .../unit/passthrough/test_passthrough_main.py | 1 + .../guardrail_hooks/test_presidio.py | 41 ++++ .../test_responses_prompt_management.py | 1 + 8 files changed, 304 insertions(+), 29 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..fc29742686f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -226,6 +226,7 @@ class CustomGuardrail(CustomLogger): use_native_lifecycle_hooks: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + masks_logged_request: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None logging_only_continue_on_input_failure: bool diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7f4b069ded6..95814989322 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -69,6 +69,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, @@ -140,6 +141,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, @@ -613,6 +615,10 @@ def _timestamp_seconds(moment: object) -> float | None: return None +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 @@ -641,6 +647,7 @@ 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 @property def _defer_async_logging(self) -> bool: @@ -2592,6 +2599,7 @@ 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) return @@ -3372,28 +3380,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 @@ -3968,22 +3968,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( @@ -3992,7 +3984,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( @@ -4169,6 +4161,66 @@ 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 + ) + + 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"))} + + 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", + guardrail.guardrail_name, + traceback.format_exc(), + ) + 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(): + try: + if not _runs_on_logging_only(guardrail, self.model_call_details): + continue + self.model_call_details, _ = await guardrail.async_logging_hook( + kwargs=self._failure_log_hook_kwargs(), result=None, call_type=self._logging_hook_call_type + ) + 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(): + 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): + continue + self.model_call_details, _ = guardrail.logging_hook( + kwargs=self._failure_log_hook_kwargs(), result=None, call_type=self._logging_hook_call_type + ) + except Exception: # noqa: BLE001 # a failing mask must not stop failure logging + self._report_failure_log_mask_error(guardrail) + 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 @@ -4291,6 +4343,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, @@ -4489,6 +4542,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, @@ -4497,14 +4552,10 @@ class Logging(LiteLLMLoggingBaseClass): ) if self._on_failure_payload_built is not None: self._on_failure_payload_built() - callbacks: Final = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_async_failure_callbacks, - global_callbacks=litellm._async_failure_callback, - ) + callbacks: Final = self._async_failure_callbacks() result: Final = None # result sent to all loggers, init this to None incase it's not created - self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: litellm_params = self.model_call_details.get("litellm_params", {}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 7e2f6bfec33..35f92403ef6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -17,7 +17,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime from itertools import accumulate, count, pairwise -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -238,6 +238,7 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 72c37c3760e..0366f31baf7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2428,6 +2428,7 @@ 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 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 299088a6c4f..06c74bf26f8 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 Final, Literal +from typing import ClassVar, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -38,6 +38,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 @@ -57,6 +58,7 @@ 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.utils import ( CallTypes, @@ -8336,6 +8338,182 @@ 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 "}] + + +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) -> dict: + self.scans += 1 + for message in 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]: + return self._mask(kwargs), result + + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, 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) -> dict: + raise RuntimeError("masking failed") + + +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[list[dict], list[dict]]] = [] + + def _record(self, kwargs: dict) -> None: + self.logged.append( + (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) + ) + + def log_failure_event(self, kwargs: dict, 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 + ) -> None: + self._record(kwargs) + + +def _register_failure_logging( + monkeypatch: pytest.MonkeyPatch, guardrail: CustomGuardrail +) -> tuple[_FailureLogRecorder, list[list[dict]]]: + recorder: Final = _FailureLogRecorder() + sync_logged: Final[list[list[dict]]] = [] + + def record_sync_failure(kwargs: dict, 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 + + +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") + 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_still_logged_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 len(recorder.logged) == 1 + assert len(sync_logged) == 1 + + +@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.now(), + 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)] + + +@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.now(), + litellm_call_id="sync-sdk-masked-failure", + function_id="sync-sdk-masked-failure", + dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await logging_obj.dispatch_failure_handlers(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", diff --git a/tests/unit/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py index 084a0fef9db..8c5d3fd45be 100644 --- a/tests/unit/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -646,6 +646,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): mock_logging_obj.update_environment_variables = MagicMock() mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() mock_logging_obj.async_failure_handler = AsyncMock() + mock_logging_obj.async_mask_failure_log_request = AsyncMock() with ( patch( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index e312378969a..fc1daf02879 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -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, @@ -4425,6 +4426,46 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use assert later[3]["content"] == "Now compare against too." +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[list[dict], list[dict]]] = [] + + async def async_log_failure_event( + self, kwargs: dict, response_obj: object, start_time: object, end_time: object + ) -> None: + self.logged.append( + (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) + ) + + +async def test_logging_only_presidio_masks_the_prompt_in_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed call is logged with the prompt a logging_only Presidio guardrail masked, the same as a + successful one. The analyzer and anonymizer are the in-process fake handed over through api_base.""" + recorder: Final = _FailureLogRecorder() + for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + async with TestServer(_fake_presidio_app()) as server: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base=str(server.make_url("/")), + presidio_anonymizer_api_base=str(server.make_url("/")), + pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK}, + logging_only=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder]) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "My name is John Smith"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + await guardrail._close_http_session() + + masked: Final = [{"role": "user", "content": "My name is "}] + assert recorder.logged == [(masked, masked)] + + @pytest.mark.asyncio @pytest.mark.parametrize("surface", ["mcp_arguments", "mcp_result", "llm_output"]) @pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK]) diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 9af92e7f9ca..bc475eba095 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -52,6 +52,7 @@ def _make_logging_obj( logging_obj.get_chat_completion_prompt.return_value = prompt_return logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return) logging_obj.async_failure_handler = AsyncMock() + logging_obj.async_mask_failure_log_request = AsyncMock() logging_obj.model_call_details = {} return logging_obj