mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
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
This commit is contained in:
parent
9a68573e9b
commit
519ecb67c8
8 changed files with 304 additions and 29 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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", {})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 <PERSON>"}]
|
||||
|
||||
|
||||
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", "<PERSON>")
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 <PERSON> 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 <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])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue