From bf52898fec09c753b80cf20d2f504ae3c19c6014 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:20:51 -0700 Subject: [PATCH] fix(guardrails): retry a handler package caught mid-import, skip unscannable during_call call types, and warn once per call type A package found while it is still being imported is now treated as failed, so the next lookup retries it instead of caching a partial map. during_call now skips a call type outside CallTypes with a warning, as pre_call and post_call do, instead of answering HTTP 500 on reads like GET /v1/responses/{id}. The skip warning is keyed on the guardrail, the call type, and the hook, so resource ids in the path no longer repeat it. Route and call type values are stripped of line breaks before they are logged. --- litellm/llms/__init__.py | 15 +++- .../unified_guardrail/unified_guardrail.py | 50 ++++++++--- .../test_guardrail_translation_discovery.py | 47 ++++++++-- .../test_openai_moderation_streaming.py | 5 +- .../test_unified_guardrail.py | 86 ++++++++++++++----- 5 files changed, 161 insertions(+), 42 deletions(-) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 7a9eeed6ebd..4bde909f3d0 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -5,7 +5,7 @@ import time from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass from itertools import chain -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger @@ -130,6 +130,10 @@ def _is_absent(module_name: str) -> bool: return False +def _still_importing(module: ModuleType) -> bool: + return attribute_or_none(attribute_or_none(module, "__spec__"), "_initializing") is True + + def _import_guardrail_translations( module_path: str, ) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage: @@ -140,6 +144,8 @@ def _import_guardrail_translations( reason=f"{type(e).__name__}: {e}", missing_dependency=isinstance(e, ModuleNotFoundError) and _is_absent(e.name or ""), ) + if _still_importing(module): + return _UnavailablePackage(reason=f"{module_path} was still being imported", missing_dependency=False) declared: Final = attribute_or_none(module, "guardrail_translation_mappings") return declared if isinstance(declared, dict) else _NO_MAPPINGS @@ -239,12 +245,15 @@ def load_guardrail_translation_mappings( return discovery.mappings -def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTranslation"]: +def get_guardrail_translation_mapping( + call_type: CallTypes, now: Callable[[], float] = time.monotonic +) -> type["BaseTranslation"]: """ Get the guardrail translation handler for a given call type. Args: call_type: The type of call (e.g., completion, acompletion, anthropic_messages) + now: Clock that decides whether a package that failed to import is due for another try Returns: The translation handler class for the given call type @@ -252,7 +261,7 @@ def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTransla Raises: ValueError: If no translation mapping exists for the given call type """ - mappings: Final = load_guardrail_translation_mappings() + mappings: Final = load_guardrail_translation_mappings(now) if call_type not in mappings: raise ValueError( f"No guardrail translation mapping found for call_type: {call_type}. " diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 1d60e79b93f..769d960bfa9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -11,6 +11,7 @@ import contextlib import copy import json from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, field from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable @@ -238,24 +239,30 @@ def _resolved_call_type(call_type: str | None) -> CallTypes | None: return None +def _without_line_breaks(value: str | None) -> str: + return str(value).replace("\r", "").replace("\n", "") + + def _log_left_unscanned( guardrail_name: str | None, request_route: str | None, call_type: str | None, consequence: str, ) -> None: + route: Final = _without_line_breaks(request_route) + shown_call_type: Final = _without_line_breaks(call_type) if _resolved_call_type(call_type) is not None: verbose_proxy_logger.warning( "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s.", guardrail_name, - request_route, - call_type, + route, + shown_call_type, consequence, ) return unscannable, remedy = ( ( - f"call type '{call_type}' is not one litellm can scan", + f"call type '{shown_call_type}' is not one litellm can scan", "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES", ) if call_type @@ -264,14 +271,32 @@ def _log_left_unscanned( verbose_proxy_logger.warning( "Guardrail '%s' selected for route '%s' but %s, so no guardrail can run on that route; %s. %s.", guardrail_name, - request_route, + route, unscannable, consequence, remedy, ) -_warn_left_unscanned_once: Final = lru_cache(maxsize=_UNSCANNED_WARNING_KEYS)(_log_left_unscanned) +@dataclass(frozen=True, slots=True) +class _UnscannedWarning: + guardrail_name: str | None + call_type_or_route: str | None + consequence: str + request_route: str | None = field(compare=False) + call_type: str | None = field(compare=False) + + +def _log_unscanned_warning(warning: _UnscannedWarning) -> None: + _log_left_unscanned( + guardrail_name=warning.guardrail_name, + request_route=warning.request_route, + call_type=warning.call_type, + consequence=warning.consequence, + ) + + +_warn_left_unscanned_once: Final = lru_cache(maxsize=_UNSCANNED_WARNING_KEYS)(_log_unscanned_warning) def _warn_left_unscanned( @@ -281,10 +306,13 @@ def _warn_left_unscanned( consequence: str, ) -> None: _warn_left_unscanned_once( - guardrail_name=guardrail_to_apply.guardrail_name, - request_route=user_api_key_dict.request_route, - call_type=call_type, - consequence=consequence, + _UnscannedWarning( + guardrail_name=guardrail_to_apply.guardrail_name, + call_type_or_route=call_type or user_api_key_dict.request_route, + consequence=consequence, + request_route=user_api_key_dict.request_route, + call_type=call_type, + ) ) @@ -393,8 +421,8 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() - resolved_call_type: Final = CallTypes(call_type) - if resolved_call_type not in mappings: + resolved_call_type: Final = _resolved_call_type(call_type) + if resolved_call_type is None or resolved_call_type not in mappings: _warn_left_unscanned( guardrail_to_apply=guardrail_to_apply, user_api_key_dict=user_api_key_dict, diff --git a/tests/unit/llms/test_guardrail_translation_discovery.py b/tests/unit/llms/test_guardrail_translation_discovery.py index f0b577a6222..543313c6a51 100644 --- a/tests/unit/llms/test_guardrail_translation_discovery.py +++ b/tests/unit/llms/test_guardrail_translation_discovery.py @@ -3,7 +3,6 @@ import importlib.machinery import importlib.util import logging import sys -import time from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager from types import ModuleType @@ -60,6 +59,29 @@ def raising_on_import(module_path: str, error: BaseException) -> Iterator[Raisin yield loader +class DiscoveringDuringItsOwnImport(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def __init__(self, module_path: str) -> None: + self.module_path = module_path + self.seen_mid_import: llms_package.GuardrailTranslationDiscovery | None = None + + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> importlib.machinery.ModuleSpec | None: + if fullname != self.module_path: + return None + return importlib.util.spec_from_loader(fullname, self) + + def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType | None: + return None + + def exec_module(self, module: ModuleType) -> None: + self.seen_mid_import = llms_package._discover((self.module_path,), already_found={}) + module.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler} + + @contextmanager def unimportable(module_path: str) -> Iterator[None]: with pytest.MonkeyPatch.context() as mp: @@ -172,11 +194,9 @@ def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announce def test_lookup_recovers_after_a_failed_discovery() -> None: with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): with pytest.raises(ValueError, match="acompletion"): - llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) + llms_package.get_guardrail_translation_mapping(CallTypes.acompletion, now=at(0.0)) - llms_package.load_guardrail_translation_mappings(now=at(time.monotonic() + RETRY_SECONDS)) - - assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) is not None + assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion, now=at(RETRY_SECONDS)) is not None def test_an_mcp_package_that_fails_to_import_is_reported_and_retried(caplog: pytest.LogCaptureFixture) -> None: @@ -232,6 +252,23 @@ class _StandInMCPHandler: pass +def test_a_package_caught_mid_import_is_retried_instead_of_cached_as_complete() -> None: + module_path: Final = "mid_import_guardrail_translation_stub" + loader: Final = DiscoveringDuringItsOwnImport(module_path) + with pytest.MonkeyPatch.context() as mp: + mp.setattr(sys, "meta_path", [loader, *sys.meta_path]) + importlib.import_module(module_path) + mid_import: Final = loader.seen_mid_import + assert mid_import is not None + retried: Final = llms_package._discover(mid_import.unavailable, already_found=mid_import.mappings) + mp.delitem(sys.modules, module_path) + + assert tuple(mid_import.unavailable) == (module_path,) + assert CallTypes.call_mcp_tool not in mid_import.mappings + assert not retried.unavailable + assert retried.mappings[CallTypes.call_mcp_tool] is _StandInMCPHandler + + def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch: pytest.MonkeyPatch) -> None: colliding: Final = ModuleType("litellm.llms.colliding_stub.guardrail_translation") colliding.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index e4a6f355c0f..0132e666dff 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -1,6 +1,5 @@ import os import sys -import time from collections.abc import AsyncIterator from typing import Final from unittest.mock import MagicMock, patch @@ -543,10 +542,10 @@ def reset_guardrail_translation_discovery(monkeypatch: pytest.MonkeyPatch) -> No async def test_moderation_still_runs_after_a_failed_translation_discovery( reset_guardrail_translation_discovery: None, ) -> None: - a_retry_interval_ago: Final = time.monotonic() - llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS + long_before_any_clock_reading: Final = float("-inf") with pytest.MonkeyPatch.context() as poison: poison.setitem(sys.modules, "litellm.llms.openai.chat.guardrail_translation", None) - poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: a_retry_interval_ago) + poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: long_before_any_clock_reading) assert CallTypes.acompletion not in poisoned with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index ecf61bea77d..bb34dc9253a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3272,46 +3272,92 @@ class TestUnscannedRequestIsAnnounced: ), self._warnings(caplog) @pytest.mark.asyncio - async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch: pytest.MonkeyPatch) -> None: + @pytest.mark.parametrize( + ("call_type", "forged_text"), + [ + pytest.param(CallTypes.acompletion.value, "/v1/responses/resp_1ERROR forged entry", id="handled-call-type"), + pytest.param("not_a_call_type\r\nERROR forged entry", "not_a_call_typeERROR forged entry", id="call-type-outside-the-enum"), + ], + ) + async def test_line_breaks_in_the_route_or_call_type_cannot_forge_log_lines( + self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, call_type: str, forged_text: str + ) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + + with self._capturing(caplog): + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", request_route="/v1/responses/resp_1\r\nERROR forged entry" + ), + cache=DualCache(), + data=self._request(RecordingGuardrail()), + call_type=call_type, + ) + + warnings: Final = self._warnings(caplog) + assert any(forged_text in message for message in warnings), warnings + assert not any("\r" in message or "\n" in message for message in warnings), warnings + + @pytest.mark.asyncio + async def test_during_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) guardrail: Final = RecordingGuardrail() - with pytest.raises(ValueError, match="acompact_responses"): - await UnifiedLLMGuardrails().async_moderation_hook( + with self._capturing(caplog): + returned: Final = await UnifiedLLMGuardrails().async_moderation_hook( data=self._request(guardrail), user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses/compact"), call_type="acompact_responses", ) assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "call type 'acompact_responses' is not one litellm can scan" in message + and "skipping during-call scanning" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + async def _pre_call_on(self, guardrail: CustomGuardrail, request_route: str, call_type: str | None) -> None: + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route=request_route), + cache=DualCache(), + data=self._request(guardrail), + call_type=call_type, + ) @pytest.mark.asyncio - async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: + async def test_a_call_type_with_no_handler_warns_once_however_many_resource_ids_reach_it(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) guardrail: Final = RecordingGuardrail() with self._capturing(caplog): - for _ in range(5): - await UnifiedLLMGuardrails().async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses/resp_123"), - cache=DualCache(), - data=self._request(guardrail), - call_type="aget_responses", - ) - await UnifiedLLMGuardrails().async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), - cache=DualCache(), - data=self._request(guardrail), - call_type="aget_responses", - ) + for response_id in ("resp_1", "resp_2", "resp_3", "resp_1"): + await self._pre_call_on(guardrail, f"/v1/responses/{response_id}", "aget_responses") + await self._pre_call_on(guardrail, "/v1/responses/resp_1", "adelete_responses") unscanned: Final = [ message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message ] assert len(unscanned) == 2, unscanned - assert "/v1/responses/resp_123" in unscanned[0] - assert "aget_responses" in unscanned[0] - assert "/v1/moderations" in unscanned[1] + assert "/v1/responses/resp_1" in unscanned[0] and "aget_responses" in unscanned[0] + assert "adelete_responses" in unscanned[1] + + @pytest.mark.asyncio + async def test_routes_without_a_call_type_each_warn_once(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail: Final = RecordingGuardrail() + + with self._capturing(caplog): + for request_route in ("/custom/route_a", "/custom/route_b", "/custom/route_a"): + await self._pre_call_on(guardrail, request_route, None) + + unscanned: Final = [ + message for message in self._distinct_warnings(caplog) if "its call type could not be resolved" in message + ] + assert len(unscanned) == 2, unscanned + assert "/custom/route_a" in unscanned[0] + assert "/custom/route_b" in unscanned[1] @pytest.mark.asyncio async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: