diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..66d30ab2c74 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -504,7 +504,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, Any], object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c12a67ed7c0..ddcb7db06a2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2,6 +2,7 @@ ## Common Utility file for Logging handler # Logging function -> log the exact model details + what's being sent | Non-Blocking import asyncio +import concurrent.futures import copy import datetime import functools @@ -569,7 +570,16 @@ 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 +_ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set() # mutable-ok: keeps orphaned tasks alive + + +def _failure_log_hook_kwargs(details: Mapping[str, Any]) -> dict[str, object]: # mutable-ok: hooks mask in place + return {**details, "messages": copy.deepcopy(details.get("messages"))} + + +def _keep_alive_until_done(task: "asyncio.Future[None]") -> None: + _ORPHANED_FAILURE_LOG_TASKS.add(task) + task.add_done_callback(_ORPHANED_FAILURE_LOG_TASKS.discard) def _provider_response_id(source: object) -> str | None: @@ -4188,9 +4198,6 @@ class Logging(LiteLLMLoggingBaseClass): 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"))} - @property def messages_sent_to_model(self) -> object: return self._unmasked_messages if self._failure_log_masked else self.model_call_details.get("messages") @@ -4214,21 +4221,25 @@ class Logging(LiteLLMLoggingBaseClass): 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", + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred masking the failure log with %s, so its request " + "is redacted %s", guardrail.guardrail_name, traceback.format_exc(), ) self._handle_callback_failure(callback=guardrail) + def _redact_failure_log_request_after_timeout(self) -> None: + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Masking the failure log took over %ss, so its request is redacted", + FAILURE_LOG_MASK_TIMEOUT_SECONDS, + ) + self._redact_failure_log_request() + async def _mask_failure_log_request_with(self, maskers: tuple[CustomGuardrail, ...]) -> None: try: await asyncio.wait_for(self._run_failure_log_maskers(maskers), timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS) except asyncio.TimeoutError: - 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() + self._redact_failure_log_request_after_timeout() async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None: for guardrail in maskers: @@ -4236,12 +4247,20 @@ class Logging(LiteLLMLoggingBaseClass): 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 + kwargs=_failure_log_hook_kwargs(self.model_call_details), + result=None, + call_type=self._logging_hook_call_type, ) - except Exception: # noqa: BLE001 # a failing mask must not stop failure logging + except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging self._report_failure_log_mask_error(guardrail) + self._redact_failure_log_request() + return - async def async_mask_failure_log_request(self, on_cancelled: Callable[[], None] | None = None) -> None: + async def async_mask_failure_log_request( + self, + on_cancelled: Callable[[], None] | None = None, + on_cancelled_async: Callable[[], Awaitable[None]] | None = None, + ) -> None: if not self._claim_failure_log_mask(): return maskers: Final = self._failure_log_maskers() @@ -4251,24 +4270,45 @@ class Logging(LiteLLMLoggingBaseClass): try: await asyncio.shield(mask) except asyncio.CancelledError: - _ORPHANED_FAILURE_LOG_MASKS.add(mask) - mask.add_done_callback(_ORPHANED_FAILURE_LOG_MASKS.discard) + _keep_alive_until_done(mask) if on_cancelled is not None: mask.add_done_callback(lambda _: on_cancelled()) + if on_cancelled_async is not None: + mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async()))) raise def mask_failure_log_request(self) -> None: if is_event_loop_running() or not self._claim_failure_log_mask(): return - 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) + maskers: Final = self._failure_log_maskers() + if not maskers: + return + masking: Final = executor.submit( + functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details + ) + try: + masked: Final = masking.result(timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS) + except concurrent.futures.TimeoutError: + self._redact_failure_log_request_after_timeout() + return + if masked is None: + self._redact_failure_log_request() + return + self.model_call_details = dict(masked) + + def _sync_failure_log_mask_step( + self, details: Mapping[str, Any] | None, guardrail: CustomGuardrail + ) -> Mapping[str, Any] | None: + if details is None or not _runs_on_logging_only(guardrail, details): + return details + try: + hooked: Final = guardrail.logging_hook( + kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type + ) + except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging + self._report_failure_log_mask_error(guardrail) + return None + return hooked[0] def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 35f92403ef6..ef8313123cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -390,7 +390,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async with self._session_lock: if self._http_session is None or self._http_session.closed: self._http_session = aiohttp.ClientSession() - yield self._http_session + session: Final = self._http_session + yield session else: # Background thread/loop -> use loop-bound session cache # This avoids "attached to a different loop" or "no running event loop" errors diff --git a/litellm/utils.py b/litellm/utils.py index 4e792240ca1..ce0cd0c6914 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2432,7 +2432,12 @@ def client(original_function): 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) + await logging_obj.async_mask_failure_log_request( + on_cancelled=run_failure_handler, + on_cancelled_async=functools.partial( + logging_obj.async_failure_handler, e, traceback_exception, start_time, end_time + ), + ) run_failure_handler() # DO NOT MAKE THREADED - router retry fallback relies on this! try: await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 5ef152e1f76..7e8cfdb0a94 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8392,6 +8392,18 @@ class _BlockingNameMaskingGuardrail(_NameMaskingGuardrail): return self._mask(kwargs), result +class _HangingSyncMaskingGuardrail(_NameMaskingGuardrail): + def __init__(self) -> None: + super().__init__() + self.release: Final = threading.Event() + + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + _ = self.release.wait(timeout=30) + return self._mask(kwargs), result + + class _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() @@ -8491,15 +8503,15 @@ async def test_failure_log_keeps_the_prompt_when_no_logging_only_guardrail_masks @pytest.mark.asyncio @pytest.mark.parametrize("call", ["completion", "acompletion"]) -async def test_failure_is_still_logged_when_the_masking_guardrail_raises( +async def test_failure_is_logged_redacted_when_the_masking_guardrail_raises( monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"] ) -> None: recorder, sync_logged = _register_failure_logging(monkeypatch, _RaisingMaskingGuardrail()) await _fail_chat_call(call, copy.deepcopy(_RAW_LOG)) - assert len(recorder.logged) == 1 - assert len(sync_logged) == 1 + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + assert sync_logged == [_REDACTED_LOG] @pytest.mark.asyncio @@ -8507,7 +8519,7 @@ async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mi monkeypatch: pytest.MonkeyPatch, ) -> None: guardrail: Final = _BlockingNameMaskingGuardrail() - _, sync_logged = _register_failure_logging(monkeypatch, guardrail) + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) call: Final = asyncio.create_task( litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()) ) @@ -8517,9 +8529,26 @@ async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mi with pytest.raises(asyncio.CancelledError): await call guardrail.release.set() - await asyncio.wait_for(_until(lambda: bool(sync_logged)), timeout=10) + await asyncio.wait_for(_until(lambda: bool(sync_logged) and bool(recorder.logged)), timeout=10) assert sync_logged == [_MASKED_LOG] + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + +@pytest.mark.asyncio +async def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _HangingSyncMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + + try: + await asyncio.wait_for(_fail_chat_call("completion", copy.deepcopy(_RAW_LOG)), timeout=10) + finally: + guardrail.release.set() + + assert sync_logged == [_REDACTED_LOG] + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] @pytest.mark.asyncio diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index d6fd03f82df..b6edd273d20 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -4897,11 +4897,17 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]: ) +_PENDING_PRESIDIO_REPLIES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: keeps pending replies alive + + class _InMemoryPresidioTransport(asyncio.Transport): - def __init__(self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None: + def __init__( + self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int + ) -> None: super().__init__() self._protocol: Final = protocol self._analyzed: Final = analyzed + self._replies_wait_for: Final = replies_wait_for self._received = b"" self._closing = False @@ -4921,7 +4927,12 @@ class _InMemoryPresidioTransport(asyncio.Transport): + f"Content-Length: {len(reply)}\r\n\r\n".encode() + reply ) - asyncio.get_running_loop().call_soon(self._protocol.data_received, response) + _PENDING_PRESIDIO_REPLIES.add(asyncio.get_running_loop().create_task(self._reply_once_analyzed(response))) + + async def _reply_once_analyzed(self, response: bytes) -> None: + while self._analyzed.qsize() < self._replies_wait_for: + await asyncio.sleep(0.01) + self._protocol.data_received(response) def _reply(self, path: str, body: bytes) -> tuple[PresidioAnalyzeResponseItem, ...] | _PresidioAnonymizeReply: if path == "/analyze": @@ -4962,15 +4973,16 @@ class _InMemoryPresidioTransport(asyncio.Transport): class _InMemoryPresidioConnector(aiohttp.BaseConnector): - def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None: + def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int) -> None: super().__init__() self._analyzed: Final = analyzed + self._replies_wait_for: Final = replies_wait_for async def _create_connection( # pyright: ignore[reportImplicitOverride] # aiohttp's connector extension point self, req: aiohttp.ClientRequest, traces: Sequence[object], timeout: aiohttp.ClientTimeout ) -> ResponseHandler: protocol: Final = ResponseHandler(asyncio.get_running_loop()) - protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed)) + protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed, self._replies_wait_for)) return protocol @@ -4979,17 +4991,37 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna def _guardrail_with_in_memory_presidio( - analyzed: asyncio.Queue[PresidioAnalyzeRequest], + analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int = 0 ) -> OPTIONAL_PresidioPIIMasking: guardrail: Final = OPTIONAL_PresidioPIIMasking( pii_entities_config=_BLOCK_CARD_MASK_EMAIL, presidio_analyzer_api_base="http://presidio-analyzer.test/", presidio_anonymizer_api_base="http://presidio-anonymizer.test/", ) - guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed)) + guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed, replies_wait_for)) return guardrail +@pytest.mark.asyncio +async def test_presidio_scans_every_message_of_a_request_at_once() -> None: + texts: Final = (f"My email is {_EMAIL}", f"Write to {_EMAIL} today") + analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() + guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, replies_wait_for=len(texts)) + try: + _ = await asyncio.wait_for( + guardrail.async_logging_hook( + kwargs={"messages": [{"role": "user", "content": text} for text in texts]}, + result=None, + call_type="acompletion", + ), + timeout=5, + ) + finally: + await guardrail._close_http_session() + + assert sorted(request.get("text") or "" for request in _drain(analyzed)) == sorted(texts) + + @pytest.mark.asyncio async def test_check_pii_raises_blocked_entity_for_card() -> None: text: Final = f"My credit card number is {_CARD_NUMBER} and my email is {_EMAIL}"