From 52935639ada4f79cd1b5f783d8c1df08b40aefc8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:45:04 -0700 Subject: [PATCH 1/9] test(guardrails): cover a guardrail translation package that fails to import --- .../test_guardrail_translation_discovery.py | 205 +++++++++++++++ .../test_mcp_guardrail_handler.py | 4 +- .../mcp_server/test_mcp_server_manager.py | 4 +- tests/unit/proxy/guardrails/conftest.py | 12 + .../test_openai_moderation_streaming.py | 79 +++++- .../test_unified_guardrail.py | 244 ++++++++++++++++++ 6 files changed, 540 insertions(+), 8 deletions(-) create mode 100644 tests/unit/llms/test_guardrail_translation_discovery.py create mode 100644 tests/unit/proxy/guardrails/conftest.py diff --git a/tests/unit/llms/test_guardrail_translation_discovery.py b/tests/unit/llms/test_guardrail_translation_discovery.py new file mode 100644 index 00000000000..fa6ed3ec789 --- /dev/null +++ b/tests/unit/llms/test_guardrail_translation_discovery.py @@ -0,0 +1,205 @@ +import importlib.abc +import importlib.util +import logging +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from types import ModuleType + +import pytest + +import litellm.llms as llms_package +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.types.utils import CallTypes + +OPENAI_CHAT_TRANSLATION_MODULE = "litellm.llms.openai.chat.guardrail_translation" +MCP_TRANSLATION_MODULE = "litellm.proxy._experimental.mcp_server.guardrail_translation" + + +class RaisingLoader(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def __init__(self, module_path: str, error: BaseException) -> None: + self.module_path = module_path + self.error = error + + def find_spec(self, fullname: str, path=None, target=None): + if fullname != self.module_path: + return None + return importlib.util.spec_from_loader(fullname, self) + + def create_module(self, spec) -> ModuleType | None: + return None + + def exec_module(self, module: ModuleType) -> None: + raise self.error + + +@contextmanager +def raising_on_import(module_path: str, error: BaseException) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.delitem(sys.modules, module_path, raising=False) + mp.setattr(sys, "meta_path", [RaisingLoader(module_path, error), *sys.meta_path]) + yield + + +@contextmanager +def unimportable(module_path: str) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.setitem(sys.modules, module_path, None) + yield + + +@contextmanager +def capturing(caplog: pytest.LogCaptureFixture, level: int) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.setattr(verbose_logger, "propagate", True) + caplog.set_level(level, logger=verbose_logger.name) + yield + + +@pytest.fixture(autouse=True) +def reset_guardrail_translation_discovery(): + saved = llms_package.guardrail_translation_discovery + llms_package.guardrail_translation_discovery = None + yield + llms_package.guardrail_translation_discovery = saved + + +def test_discovery_reports_the_handler_package_it_could_not_import(): + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + discovery = llms_package.discover_guardrail_translations() + + assert tuple(discovery.unavailable) == (OPENAI_CHAT_TRANSLATION_MODULE,) + assert "None in sys.modules" in discovery.unavailable[OPENAI_CHAT_TRANSLATION_MODULE] + assert CallTypes.acompletion not in discovery.mappings + assert CallTypes.completion not in discovery.mappings + assert CallTypes.aembedding in discovery.mappings + + +def test_complete_discovery_reports_nothing_unavailable(): + discovery = llms_package.discover_guardrail_translations() + + assert not discovery.unavailable + assert CallTypes.acompletion in discovery.mappings + + +def test_the_next_lookup_retries_a_package_that_failed_to_import(): + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + partial = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion not in partial + assert CallTypes.aembedding in partial + + recovered = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion in recovered + assert CallTypes.completion in recovered + assert CallTypes.aembedding in recovered + assert not llms_package.guardrail_translation_discovery.unavailable + + +def test_a_complete_discovery_is_not_scanned_again(): + healthy = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion in healthy + + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + after_the_package_breaks = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion in after_the_package_breaks + assert not llms_package.guardrail_translation_discovery.unavailable + + +def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announced(caplog): + with capturing(caplog, logging.INFO), unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + for _ in range(3): + llms_package.load_guardrail_translation_mappings() + + errors = [record for record in caplog.records if record.levelno == logging.ERROR] + assert len(errors) == 1, [record.getMessage() for record in errors] + assert OPENAI_CHAT_TRANSLATION_MODULE in errors[0].getMessage() + assert "None in sys.modules" in errors[0].getMessage() + + caplog.clear() + with capturing(caplog, logging.INFO): + llms_package.load_guardrail_translation_mappings() + llms_package.load_guardrail_translation_mappings() + + recoveries = [record for record in caplog.records if "available again" in record.getMessage()] + assert len(recoveries) == 1 + assert OPENAI_CHAT_TRANSLATION_MODULE in recoveries[0].getMessage() + assert recoveries[0].levelno >= logging.WARNING + assert not [record for record in caplog.records if record.levelno == logging.ERROR] + + +def test_lookup_recovers_after_a_failed_discovery(): + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + with pytest.raises(ValueError, match="acompletion"): + llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) + + assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) is not None + + +def test_an_mcp_package_that_fails_to_import_is_reported_and_retried(): + with raising_on_import(MCP_TRANSLATION_MODULE, AttributeError("module 'mcp.types' has no attribute 'ToolCall'")): + partial = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.call_mcp_tool not in partial + assert CallTypes.acompletion in partial + assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,) + assert "AttributeError" in llms_package.guardrail_translation_discovery.unavailable[MCP_TRANSLATION_MODULE] + + recovered = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.call_mcp_tool in recovered + assert CallTypes.acompletion in recovered + assert not llms_package.guardrail_translation_discovery.unavailable + + +def test_an_install_without_the_mcp_server_is_discovered_once_and_quietly(caplog): + absent = ModuleNotFoundError("No module named 'mcp'", name="mcp") + + with capturing(caplog, logging.DEBUG), unimportable("mcp"), raising_on_import(MCP_TRANSLATION_MODULE, absent): + first = llms_package.load_guardrail_translation_mappings() + second = llms_package.load_guardrail_translation_mappings() + + assert first is second + assert CallTypes.call_mcp_tool not in first + assert CallTypes.acompletion in first + assert not llms_package.guardrail_translation_discovery.unavailable + assert not [record for record in caplog.records if record.levelno >= logging.WARNING] + + +def test_a_broken_mcp_dependency_is_reported_and_retried(caplog): + broken = ModuleNotFoundError("No module named 'mcp.types'", name="mcp.types") + + with capturing(caplog, logging.DEBUG), raising_on_import(MCP_TRANSLATION_MODULE, broken): + partial = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.call_mcp_tool not in partial + assert CallTypes.acompletion in partial + assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,) + errors = [record for record in caplog.records if record.levelno >= logging.ERROR] + assert len(errors) == 1, [record.getMessage() for record in caplog.records] + assert "mcp.types" in errors[0].getMessage() + + recovered = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.call_mcp_tool in recovered + assert not llms_package.guardrail_translation_discovery.unavailable + + +class _StandInMCPHandler: + pass + + +def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch): + colliding = ModuleType("litellm.llms.colliding_stub.guardrail_translation") + colliding.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler} + monkeypatch.setitem(sys.modules, colliding.__name__, colliding) + + discovery = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={}) + + assert discovery.mappings[CallTypes.call_mcp_tool] is MCPGuardrailTranslationHandler diff --git a/tests/unit/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/unit/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index f3e1bcf979f..819d45063e8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/unit/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -133,8 +133,8 @@ def restore_callbacks(monkeypatch): monkeypatch.setattr(litellm, "callbacks", litellm.callbacks) monkeypatch.setattr( litellm_llms, - "endpoint_guardrail_translation_mappings", - litellm_llms.endpoint_guardrail_translation_mappings, + "guardrail_translation_discovery", + litellm_llms.guardrail_translation_discovery, ) yield ProxyLogging._callback_capabilities_cache.clear() diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7315315a619..25dc477d259 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -19906,8 +19906,8 @@ def catalog_guardrail(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [guardrail]) monkeypatch.setattr( litellm_llms, - "endpoint_guardrail_translation_mappings", - litellm_llms.endpoint_guardrail_translation_mappings, + "guardrail_translation_discovery", + litellm_llms.guardrail_translation_discovery, ) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.slack_alerting_instance.send_alert = AsyncMock() diff --git a/tests/unit/proxy/guardrails/conftest.py b/tests/unit/proxy/guardrails/conftest.py new file mode 100644 index 00000000000..56a93c515bd --- /dev/null +++ b/tests/unit/proxy/guardrails/conftest.py @@ -0,0 +1,12 @@ +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, +) + + +@pytest.fixture(autouse=True) +def _forget_unscanned_warnings(): + unified_module._warn_left_unscanned_once.cache_clear() + yield + unified_module._warn_left_unscanned_once.cache_clear() 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 16f04073fae..971f5e29a0f 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,14 +1,24 @@ -import pytest -from unittest.mock import MagicMock, patch import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +import litellm.llms as llms_package +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) -from litellm.types.utils import ModelResponseStream, ModelResponse -from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import ( + CallTypes, + Delta, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) @pytest.mark.asyncio @@ -519,3 +529,64 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) + + +@pytest.fixture +def reset_guardrail_translation_discovery(): + saved = llms_package.guardrail_translation_discovery + llms_package.guardrail_translation_discovery = None + yield llms_package + llms_package.guardrail_translation_discovery = saved + + +@pytest.mark.asyncio +async def test_moderation_still_runs_after_a_failed_translation_discovery( + reset_guardrail_translation_discovery, +): + with pytest.MonkeyPatch.context() as poison: + poison.setitem(sys.modules, "litellm.llms.openai.chat.guardrail_translation", None) + poisoned = llms_package.load_guardrail_translation_mappings() + assert CallTypes.acompletion not in poisoned + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = MagicMock() + mock_mod_response.results = [] + + async def mock_stream(): + chunks_data = ["Hello", " ", "world", "!", " Goodbye"] + for i, content in enumerate(chunks_data): + yield ModelResponseStream( + model="gpt-4", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=content, role="assistant"), + finish_reason=("stop" if i == len(chunks_data) - 1 else None), + ) + ], + ) + + with patch.object(openai_guardrail, "async_make_request", return_value=mock_mod_response) as patched_make_request: + chunks_received = 0 + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), + response=mock_stream(), + request_data={ + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + }, + ): + chunks_received += 1 + + assert chunks_received == 5 + assert patched_make_request.await_count > 0, ( + "Moderation never ran: the failed discovery was cached and the streaming hook " + "passed every chunk through unscanned" + ) 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 3efbe6475b2..90c62531e75 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 @@ -3053,3 +3053,247 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none(): ) assert len(guardrail.apply_calls) == 1 + + +class TestUnscannedStreamIsAnnounced: + + @staticmethod + async def _drive(caplog, monkeypatch, request_route, mappings, response_chunks): + _patch_translation_mappings(monkeypatch, mappings) + + async def stream(): + for chunk in response_chunks: + yield chunk + + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") + unified_module.verbose_proxy_logger.addHandler(caplog.handler) + try: + chunks = [ + chunk + async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test", request_route=request_route + ), + response=stream(), + request_data={ + "guardrail_to_apply": RecordingGuardrail(), + "metadata": {"guardrails": ["recording-guardrail"]}, + }, + ) + ] + finally: + unified_module.verbose_proxy_logger.removeHandler(caplog.handler) + + return chunks, [ + record.getMessage() + for record in caplog.records + if record.levelno >= logging.WARNING + ] + + @pytest.mark.asyncio + async def test_warns_when_the_call_type_has_no_translation_handler( + self, caplog, monkeypatch + ): + chunks, warnings = await self._drive( + caplog, + monkeypatch, + request_route="/chat/completions", + mappings={CallTypes.aembedding: _NoopTranslation}, + response_chunks=[ + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content))] + ) + for content in ("a", "b", "c") + ], + ) + + assert len(chunks) == 3 + assert any( + "no guardrail translation handler" in message + and "Add a guardrail translation handler for that call type." in message + and "recording-guardrail" in message + and "/chat/completions" in message + for message in warnings + ), warnings + + @pytest.mark.asyncio + async def test_warns_when_the_call_type_cannot_be_resolved(self, caplog, monkeypatch): + chunks, warnings = await self._drive( + caplog, + monkeypatch, + request_route="/v1/not-a-mapped-route", + mappings=load_guardrail_translation_mappings(), + response_chunks=[{"event": "delta", "text": content} for content in ("a", "b", "c")], + ) + + assert len(chunks) == 3 + assert any( + "call type could not be resolved" in message + and "Add the route to API_ROUTE_TO_CALL_TYPES." in message + and "recording-guardrail" in message + for message in warnings + ), warnings + + +class TestUnscannedRequestIsAnnounced: + + @staticmethod + def _request(guardrail): + return { + "guardrail_to_apply": guardrail, + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello world"}], + } + + @staticmethod + @contextlib.contextmanager + def _capturing(caplog): + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") + unified_module.verbose_proxy_logger.addHandler(caplog.handler) + try: + yield + finally: + unified_module.verbose_proxy_logger.removeHandler(caplog.handler) + + @staticmethod + def _warnings(caplog): + return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] + + @staticmethod + def _distinct_warnings(caplog): + by_record = {id(record): record for record in caplog.records if record.levelno >= logging.WARNING} + return [record.getMessage() for record in by_record.values()] + + @pytest.mark.asyncio + async def test_pre_call_warns_when_the_call_type_has_no_translation_handler(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail = RecordingGuardrail() + data = self._request(guardrail) + + with self._capturing(caplog): + returned = await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), + cache=DualCache(), + data=data, + call_type=CallTypes.acompletion.value, + ) + + assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "no guardrail translation handler" in message + and "skipping pre-call scanning" in message + and "recording-guardrail" in message + and "/v1/moderations" in message + and "acompletion" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + @pytest.mark.asyncio + async def test_during_call_warns_when_the_call_type_has_no_translation_handler(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail = RecordingGuardrail() + data = self._request(guardrail) + + with self._capturing(caplog): + returned = await UnifiedLLMGuardrails().async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), + call_type=CallTypes.acompletion.value, + ) + + assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "skipping during-call scanning" in message and "recording-guardrail" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + @pytest.mark.asyncio + async def test_pre_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) + guardrail = RecordingGuardrail() + data = self._request(guardrail) + + with self._capturing(caplog): + returned = await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), + cache=DualCache(), + data=data, + call_type="not_a_call_type", + ) + + assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "call type 'not_a_call_type' is not one litellm can scan" in message + and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message + and "skipping pre-call scanning" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + @pytest.mark.asyncio + async def test_during_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) + guardrail = RecordingGuardrail() + data = self._request(guardrail) + + with self._capturing(caplog): + returned = await UnifiedLLMGuardrails().async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), + call_type="not_a_call_type", + ) + + assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "call type 'not_a_call_type' is not one litellm can scan" in message + and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message + and "skipping during-call scanning" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + @pytest.mark.asyncio + async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail = RecordingGuardrail() + + with self._capturing(caplog): + for _ in range(5): + await UnifiedLLMGuardrails().async_moderation_hook( + data=self._request(guardrail), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses/resp_123"), + call_type="aget_responses", + ) + await UnifiedLLMGuardrails().async_moderation_hook( + data=self._request(guardrail), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), + call_type="aget_responses", + ) + + unscanned = [ + message for message in self._distinct_warnings(caplog) if "skipping during-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] + + @pytest.mark.asyncio + async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail = RecordingGuardrail() + + with self._capturing(caplog): + await UnifiedLLMGuardrails().async_post_call_success_hook( + data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=litellm.ModelResponse(), + ) + + assert guardrail.apply_calls == [] + unscanned = [message for message in self._warnings(caplog) if "skipping post-call scanning" in message] + assert unscanned, self._warnings(caplog) + assert "call type 'aresponses'" in unscanned[0], unscanned[0] + assert "CallTypes." not in unscanned[0], unscanned[0] From 90561367f9387756ff9091223c96f0b16fa0efd6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:47:00 -0700 Subject: [PATCH 2/9] fix(guardrails): retry guardrail handler packages that failed to import and log unscanned routes --- litellm/llms/__init__.py | 213 ++++++++++-------- .../unified_guardrail/unified_guardrail.py | 122 +++++++--- 2 files changed, 219 insertions(+), 116 deletions(-) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 7452a4a20ef..5bc4754835a 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -1,5 +1,10 @@ import importlib +import importlib.util import os +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger @@ -89,93 +94,133 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo return None -def discover_guardrail_translation_mappings() -> dict[CallTypes, type["BaseTranslation"]]: - """ - Discover guardrail translation mappings by scanning the llms directory structure. +_GUARDRAIL_TRANSLATION_PACKAGE: Final = "guardrail_translation" +_MCP_GUARDRAIL_TRANSLATION_MODULE: Final = "litellm.proxy._experimental.mcp_server.guardrail_translation" +_NO_MAPPINGS: Final[Mapping[CallTypes, type["BaseTranslation"]]] = MappingProxyType({}) - Scans for modules with guardrail_translation_mappings dictionaries and aggregates them. - Returns: - Dict[CallTypes, Type[BaseTranslation]]: A dictionary mapping call types to their translation handler classes - """ - discovered_mappings: Final[dict[CallTypes, type[BaseTranslation]]] = {} +@dataclass(frozen=True, slots=True) +class GuardrailTranslationDiscovery: + mappings: Mapping[CallTypes, type["BaseTranslation"]] + unavailable: Mapping[str, str] + +def _bundled_guardrail_translation_modules() -> Iterator[str]: + llms_dir: Final = os.path.dirname(__file__) + for root, dirs, files in os.walk(llms_dir): + dirs[:] = tuple(d for d in dirs if not d.startswith("__") and d != "base_llm") + if os.path.basename(root) == _GUARDRAIL_TRANSLATION_PACKAGE and "__init__.py" in files: + yield "litellm." + os.path.relpath(root, os.path.dirname(llms_dir)).replace(os.sep, ".") + + +@dataclass(frozen=True, slots=True) +class _UnavailablePackage: + reason: str + missing_dependency: bool + + +def _is_absent(module_name: str) -> bool: + root: Final = module_name.partition(".")[0] try: - # Get the path to the llms directory - current_dir: Final = os.path.dirname(__file__) - llms_dir: Final = current_dir + return importlib.util.find_spec(root) is None + except (ImportError, ValueError): + return False - if not os.path.exists(llms_dir): - verbose_logger.debug("llms directory not found") - return discovered_mappings - # Recursively scan for guardrail_translation directories - for root, dirs, files in os.walk(llms_dir): - # Skip __pycache__ and base_llm directories - dirs[:] = [d for d in dirs if not d.startswith("__") and d != "base_llm"] - - # Check if this is a guardrail_translation directory with __init__.py - if os.path.basename(root) == "guardrail_translation" and "__init__.py" in files: - # Build the module path relative to litellm - rel_path = os.path.relpath(root, os.path.dirname(llms_dir)) - module_path = "litellm." + rel_path.replace(os.sep, ".") - - try: - # Import the module - verbose_logger.debug("Discovering guardrail translations in: %s", module_path) - - module = importlib.import_module(module_path) - - # Check for guardrail_translation_mappings dictionary - if hasattr(module, "guardrail_translation_mappings"): - mappings = getattr(module, "guardrail_translation_mappings") - if isinstance(mappings, dict): - discovered_mappings.update(mappings) - verbose_logger.debug( - "Found guardrail_translation_mappings in %s: %s", module_path, list(mappings.keys()) - ) - - except ImportError as e: - verbose_logger.error("Could not import %s: %s", module_path, e) - continue - except Exception as e: - verbose_logger.error("Error processing %s: %s", module_path, e) - continue - - try: - from litellm.proxy._experimental.mcp_server.guardrail_translation import ( - guardrail_translation_mappings as mcp_guardrail_translation_mappings, - ) - - discovered_mappings.update(mcp_guardrail_translation_mappings) - verbose_logger.debug( - "Loaded MCP guardrail translation mappings: %s", - list(mcp_guardrail_translation_mappings.keys()), - ) - except ImportError: - verbose_logger.debug("MCP guardrail translation mappings not available; skipping") - - verbose_logger.debug( - "Discovered %s guardrail translation mappings: %s", - len(discovered_mappings), - list(discovered_mappings.keys()), +def _import_guardrail_translations( + module_path: str, +) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage: + try: + module: Final = importlib.import_module(module_path) + except Exception as e: # noqa: BLE001 # a package failing at import time for any reason is unavailable, not fatal + return _UnavailablePackage( + reason=f"{type(e).__name__}: {e}", + missing_dependency=isinstance(e, ModuleNotFoundError) and _is_absent(e.name or ""), ) - - except Exception as e: - verbose_logger.error("Error discovering guardrail translation mappings: %s", e) - - return discovered_mappings + declared: Final[Mapping[CallTypes, type[BaseTranslation]] | None] = getattr( + module, "guardrail_translation_mappings", None + ) + return declared if isinstance(declared, dict) else _NO_MAPPINGS -# Cache the discovered mappings -endpoint_guardrail_translation_mappings: dict[CallTypes, type["BaseTranslation"]] | None = None +def _guardrail_translations_from( + module_path: str, +) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage: + result: Final = _import_guardrail_translations(module_path) + if ( + module_path == _MCP_GUARDRAIL_TRANSLATION_MODULE + and isinstance(result, _UnavailablePackage) + and result.missing_dependency + ): + verbose_logger.debug("%s is not installed: %s", module_path, result.reason) + return _NO_MAPPINGS + return result -def load_guardrail_translation_mappings(): - global endpoint_guardrail_translation_mappings - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = discover_guardrail_translation_mappings() - return endpoint_guardrail_translation_mappings +def _guardrail_translation_modules() -> Iterator[str]: + yield from _bundled_guardrail_translation_modules() + yield _MCP_GUARDRAIL_TRANSLATION_MODULE + + +def _discover( + module_paths: Iterable[str], already_found: Mapping[CallTypes, type["BaseTranslation"]] +) -> GuardrailTranslationDiscovery: + imported: Final = tuple((module_path, _guardrail_translations_from(module_path)) for module_path in module_paths) + found: Final = ( + already_found, + *(result for _, result in imported if not isinstance(result, _UnavailablePackage)), + ) + return GuardrailTranslationDiscovery( + mappings=MappingProxyType(dict(chain.from_iterable(mappings.items() for mappings in found))), + unavailable=MappingProxyType( + {module_path: result.reason for module_path, result in imported if isinstance(result, _UnavailablePackage)} + ), + ) + + +def discover_guardrail_translations() -> GuardrailTranslationDiscovery: + return _discover(_guardrail_translation_modules(), already_found=_NO_MAPPINGS) + + +def discover_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]: + return discover_guardrail_translations().mappings + + +def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current: GuardrailTranslationDiscovery) -> None: + if previous is None and not current.unavailable: + return + if previous is None: + verbose_logger.error( + "Could not import guardrail translation handlers from %s; guardrails cannot run for their call types " + "until the import succeeds, which every lookup retries. %s", + ", ".join(current.unavailable), + "; ".join(f"{module_path}: {reason}" for module_path, reason in current.unavailable.items()), + ) + return + recovered: Final = tuple( + module_path for module_path in previous.unavailable if module_path not in current.unavailable + ) + if not recovered: + return + verbose_logger.warning("Guardrail translation handlers from %s are available again.", ", ".join(recovered)) + + +guardrail_translation_discovery: GuardrailTranslationDiscovery | None = None + + +def load_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]: + global guardrail_translation_discovery + cached: Final = guardrail_translation_discovery + if cached is not None and not cached.unavailable: + return cached.mappings + discovery: Final = ( + discover_guardrail_translations() + if cached is None + else _discover(cached.unavailable, already_found=cached.mappings) + ) + _announce_discovery(previous=cached, current=discovery) + guardrail_translation_discovery = discovery + return discovery.mappings def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTranslation"]: @@ -191,18 +236,10 @@ def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTransla Raises: ValueError: If no translation mapping exists for the given call type """ - global endpoint_guardrail_translation_mappings - - # Lazy load the mappings on first access - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = discover_guardrail_translation_mappings() - - # Get the translation handler class for the call type - if call_type not in endpoint_guardrail_translation_mappings: + mappings: Final = load_guardrail_translation_mappings() + if call_type not in mappings: raise ValueError( f"No guardrail translation mapping found for call_type: {call_type}. " - f"Available mappings: {list(endpoint_guardrail_translation_mappings.keys())}" + f"Available mappings: {list(mappings.keys())}" ) - - # Return the handler class directly - return endpoint_guardrail_translation_mappings[call_type] + return mappings[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 3ad9d1a2c42..7dbb70e99db 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 functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable import anyio @@ -224,6 +225,65 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N data["litellm_metadata"] = user_metadata +_UNSCANNED_WARNING_KEYS: Final = 4096 + + +def _resolved_call_type(call_type: str | None) -> CallTypes | None: + try: + return CallTypes(call_type) + except ValueError: + return None + + +@lru_cache(maxsize=_UNSCANNED_WARNING_KEYS) +def _warn_left_unscanned_once( + guardrail_name: str | None, + request_route: str | None, + call_type: str | None, + consequence: str, +) -> None: + 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. " + "Add a guardrail translation handler for that call type.", + guardrail_name, + request_route, + call_type, + consequence, + ) + return + unscannable, remedy = ( + ( + f"call type '{call_type}' is not one litellm can scan", + "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES", + ) + if call_type + else ("its call type could not be resolved", "Add the route to API_ROUTE_TO_CALL_TYPES") + ) + 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, + unscannable, + consequence, + remedy, + ) + + +def _warn_left_unscanned( + guardrail_to_apply: "CustomGuardrail", + user_api_key_dict: UserAPIKeyAuth, + call_type: str | None, + 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, + ) + + class UnifiedLLMGuardrails(CustomLogger): def __init__( self, @@ -274,14 +334,17 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_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, + call_type=call_type, + consequence="skipping pre-call scanning", + ) + return data - try: - if CallTypes(call_type) not in mappings: - return data - except ValueError: - return data # handle unmapped call types - - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[resolved_call_type]()) _ensure_litellm_metadata(data, user_api_key_dict) @@ -326,10 +389,17 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() - if call_type is not None and CallTypes(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, + call_type=call_type, + consequence="skipping during-call scanning", + ) return data - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[resolved_call_type]()) _ensure_litellm_metadata(data, user_api_key_dict) @@ -374,7 +444,7 @@ class UnifiedLLMGuardrails(CustomLogger): if user_api_key_dict.request_route is not None: call_types: Final = get_call_types_for_route(user_api_key_dict.request_route) if call_types is not None and len(call_types) > 0: - call_type = call_types[0] + call_type = call_types[0].value if call_type is None: call_type = infer_call_type(call_type=None, completion_response=response) @@ -391,28 +461,18 @@ class UnifiedLLMGuardrails(CustomLogger): ): call_type = logging_call_type - if call_type is None: - verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but its call type could not be resolved; " - "skipping post-call scanning. Add the route to API_ROUTE_TO_CALL_TYPES.", - guardrail_to_apply.guardrail_name, - user_api_key_dict.request_route, - ) - return response - mappings: Final = load_guardrail_translation_mappings() - - if CallTypes(call_type) not in mappings: - verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; " - "skipping post-call scanning.", - guardrail_to_apply.guardrail_name, - user_api_key_dict.request_route, - call_type, + 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, + call_type=call_type, + consequence="skipping post-call scanning", ) return response - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[resolved_call_type]()) try: response = await endpoint_translation.process_output_response( @@ -1260,6 +1320,12 @@ class UnifiedLLMGuardrails(CustomLogger): # If call type not supported, just pass through all chunks if call_type is None or CallTypes(call_type) not in mappings: + _warn_left_unscanned( + guardrail_to_apply=guardrail_to_apply, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + consequence="streaming this response to the client unscanned", + ) yield item async for remaining_item in response: yield remaining_item From 9c3fa811b28d4d68ffb957d49bba56f15a4feab3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:59:51 -0700 Subject: [PATCH 3/9] test(guardrails): move the discovery tests into a directory a unit shard runs --- .../test_guardrail_translation_discovery.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/unit/llms/{ => base_llm/guardrail_translation}/test_guardrail_translation_discovery.py (100%) diff --git a/tests/unit/llms/test_guardrail_translation_discovery.py b/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py similarity index 100% rename from tests/unit/llms/test_guardrail_translation_discovery.py rename to tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py From 9c6c1e2738b7737bb4f96780e42990f2de5d80ee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:02:57 -0700 Subject: [PATCH 4/9] fix(guardrails): keep rejecting during-call requests whose call type cannot be scanned and warn on every unscanned post-call response --- .../unified_guardrail/unified_guardrail.py | 19 +++---- .../test_unified_guardrail.py | 52 ++++++++++++------- 2 files changed, 42 insertions(+), 29 deletions(-) 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 7dbb70e99db..3fc9b7e9f28 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -235,8 +235,7 @@ def _resolved_call_type(call_type: str | None) -> CallTypes | None: return None -@lru_cache(maxsize=_UNSCANNED_WARNING_KEYS) -def _warn_left_unscanned_once( +def _log_left_unscanned( guardrail_name: str | None, request_route: str | None, call_type: str | None, @@ -244,8 +243,7 @@ def _warn_left_unscanned_once( ) -> None: 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. " - "Add a guardrail translation handler for that call type.", + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s.", guardrail_name, request_route, call_type, @@ -270,6 +268,9 @@ def _warn_left_unscanned_once( ) +_warn_left_unscanned_once: Final = lru_cache(maxsize=_UNSCANNED_WARNING_KEYS)(_log_left_unscanned) + + def _warn_left_unscanned( guardrail_to_apply: "CustomGuardrail", user_api_key_dict: UserAPIKeyAuth, @@ -389,8 +390,8 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() - resolved_call_type: Final = _resolved_call_type(call_type) - if resolved_call_type is None or resolved_call_type not in mappings: + resolved_call_type: Final = CallTypes(call_type) + if resolved_call_type not in mappings: _warn_left_unscanned( guardrail_to_apply=guardrail_to_apply, user_api_key_dict=user_api_key_dict, @@ -464,9 +465,9 @@ class UnifiedLLMGuardrails(CustomLogger): mappings: Final = load_guardrail_translation_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, + _log_left_unscanned( + guardrail_name=guardrail_to_apply.guardrail_name, + request_route=user_api_key_dict.request_route, call_type=call_type, consequence="skipping post-call scanning", ) 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 90c62531e75..1ad727c3982 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 @@ -3110,7 +3110,7 @@ class TestUnscannedStreamIsAnnounced: assert len(chunks) == 3 assert any( "no guardrail translation handler" in message - and "Add a guardrail translation handler for that call type." in message + and "streaming this response to the client unscanned" in message and "recording-guardrail" in message and "/chat/completions" in message for message in warnings @@ -3233,26 +3233,18 @@ class TestUnscannedRequestIsAnnounced: ), self._warnings(caplog) @pytest.mark.asyncio - async def test_during_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog, monkeypatch): + async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch): _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) guardrail = RecordingGuardrail() - data = self._request(guardrail) - with self._capturing(caplog): - returned = await UnifiedLLMGuardrails().async_moderation_hook( - data=data, - user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), - call_type="not_a_call_type", + with pytest.raises(ValueError, match="acompact_responses"): + 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 'not_a_call_type' is not one litellm can scan" in message - and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message - and "skipping during-call scanning" in message - for message in self._warnings(caplog) - ), self._warnings(caplog) @pytest.mark.asyncio async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog, monkeypatch): @@ -3261,19 +3253,21 @@ class TestUnscannedRequestIsAnnounced: with self._capturing(caplog): for _ in range(5): - await UnifiedLLMGuardrails().async_moderation_hook( - data=self._request(guardrail), + 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_moderation_hook( - data=self._request(guardrail), + 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", ) unscanned = [ - message for message in self._distinct_warnings(caplog) if "skipping during-call scanning" in message + 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] @@ -3297,3 +3291,21 @@ class TestUnscannedRequestIsAnnounced: assert unscanned, self._warnings(caplog) assert "call type 'aresponses'" in unscanned[0], unscanned[0] assert "CallTypes." not in unscanned[0], unscanned[0] + + @pytest.mark.asyncio + async def test_post_call_warns_on_every_unscanned_response(self, caplog, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail = RecordingGuardrail() + + with self._capturing(caplog): + for _ in range(3): + await UnifiedLLMGuardrails().async_post_call_success_hook( + data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=litellm.ModelResponse(), + ) + + unscanned = [ + message for message in self._distinct_warnings(caplog) if "skipping post-call scanning" in message + ] + assert len(unscanned) == 3, unscanned From 5264a70d2519f907ef994cd25507f5ed3f079be2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:15:09 -0700 Subject: [PATCH 5/9] fix(guardrails): retry a failed guardrail translation import at most once a second and type the new tests --- litellm/llms/__init__.py | 35 +++- .../test_guardrail_translation_discovery.py | 161 +++++++++++------- tests/unit/proxy/guardrails/conftest.py | 4 +- .../test_openai_moderation_streaming.py | 52 +++--- .../test_unified_guardrail.py | 76 +++++---- 5 files changed, 196 insertions(+), 132 deletions(-) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 74aaff384e6..7a9eeed6ebd 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -1,7 +1,8 @@ import importlib import importlib.util import os -from collections.abc import Iterable, Iterator, Mapping +import time +from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass from itertools import chain from types import MappingProxyType @@ -98,6 +99,7 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo _GUARDRAIL_TRANSLATION_PACKAGE: Final = "guardrail_translation" _MCP_GUARDRAIL_TRANSLATION_MODULE: Final = "litellm.proxy.mcp_server.guardrail_translation" _NO_MAPPINGS: Final[Mapping[CallTypes, type["BaseTranslation"]]] = MappingProxyType({}) +GUARDRAIL_TRANSLATION_RETRY_SECONDS: Final = 1.0 @dataclass(frozen=True, slots=True) @@ -191,8 +193,9 @@ def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current: if previous is None: verbose_logger.error( "Could not import guardrail translation handlers from %s; guardrails cannot run for their call types " - "until the import succeeds, which every lookup retries. %s", + "until the import succeeds, which is retried at most once every %s seconds. %s", ", ".join(current.unavailable), + GUARDRAIL_TRANSLATION_RETRY_SECONDS, "; ".join(f"{module_path}: {reason}" for module_path, reason in current.unavailable.items()), ) return @@ -204,21 +207,35 @@ def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current: verbose_logger.warning("Guardrail translation handlers from %s are available again.", ", ".join(recovered)) -guardrail_translation_discovery: GuardrailTranslationDiscovery | None = None +@dataclass(frozen=True, slots=True) +class CachedGuardrailTranslationDiscovery: + discovery: GuardrailTranslationDiscovery + retry_at: float -def load_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]: +guardrail_translation_discovery: CachedGuardrailTranslationDiscovery | None = None # rebind-ok: per-worker cache + + +def load_guardrail_translation_mappings( + now: Callable[[], float] = time.monotonic, +) -> Mapping[CallTypes, type["BaseTranslation"]]: global guardrail_translation_discovery cached: Final = guardrail_translation_discovery - if cached is not None and not cached.unavailable: - return cached.mappings + if cached is not None and not cached.discovery.unavailable: + return cached.discovery.mappings + checked_at: Final = now() + if cached is not None and checked_at < cached.retry_at: + return cached.discovery.mappings discovery: Final = ( discover_guardrail_translations() if cached is None - else _discover(cached.unavailable, already_found=cached.mappings) + else _discover(cached.discovery.unavailable, already_found=cached.discovery.mappings) ) - _announce_discovery(previous=cached, current=discovery) - guardrail_translation_discovery = discovery + _announce_discovery(previous=cached.discovery if cached is not None else None, current=discovery) + refreshed: Final = CachedGuardrailTranslationDiscovery( + discovery=discovery, retry_at=checked_at + GUARDRAIL_TRANSLATION_RETRY_SECONDS + ) + guardrail_translation_discovery = refreshed # rebind-ok: per-worker cache, replaced whole return discovery.mappings diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py b/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py index 4694eba8759..f0b577a6222 100644 --- a/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py +++ b/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py @@ -1,10 +1,13 @@ import importlib.abc +import importlib.machinery import importlib.util import logging import sys -from collections.abc import Iterator +import time +from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager from types import ModuleType +from typing import Final import pytest @@ -15,33 +18,46 @@ from litellm.proxy.mcp_server.guardrail_translation.handler import ( ) from litellm.types.utils import CallTypes -OPENAI_CHAT_TRANSLATION_MODULE = "litellm.llms.openai.chat.guardrail_translation" -MCP_TRANSLATION_MODULE = "litellm.proxy.mcp_server.guardrail_translation" +OPENAI_CHAT_TRANSLATION_MODULE: Final = "litellm.llms.openai.chat.guardrail_translation" +MCP_TRANSLATION_MODULE: Final = "litellm.proxy.mcp_server.guardrail_translation" +RETRY_SECONDS: Final = llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS + + +def at(seconds: float) -> Callable[[], float]: + return lambda: seconds class RaisingLoader(importlib.abc.MetaPathFinder, importlib.abc.Loader): def __init__(self, module_path: str, error: BaseException) -> None: self.module_path = module_path self.error = error + self.attempts = 0 - def find_spec(self, fullname: str, path=None, target=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) -> ModuleType | None: + def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType | None: return None def exec_module(self, module: ModuleType) -> None: + self.attempts += 1 raise self.error @contextmanager -def raising_on_import(module_path: str, error: BaseException) -> Iterator[None]: +def raising_on_import(module_path: str, error: BaseException) -> Iterator[RaisingLoader]: + loader: Final = RaisingLoader(module_path, error) with pytest.MonkeyPatch.context() as mp: mp.delitem(sys.modules, module_path, raising=False) - mp.setattr(sys, "meta_path", [RaisingLoader(module_path, error), *sys.meta_path]) - yield + mp.setattr(sys, "meta_path", [loader, *sys.meta_path]) + yield loader @contextmanager @@ -59,17 +75,18 @@ def capturing(caplog: pytest.LogCaptureFixture, level: int) -> Iterator[None]: yield +def errors_in(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple(record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR) + + @pytest.fixture(autouse=True) -def reset_guardrail_translation_discovery(): - saved = llms_package.guardrail_translation_discovery - llms_package.guardrail_translation_discovery = None - yield - llms_package.guardrail_translation_discovery = saved +def reset_guardrail_translation_discovery(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(llms_package, "guardrail_translation_discovery", None) -def test_discovery_reports_the_handler_package_it_could_not_import(): +def test_discovery_reports_the_handler_package_it_could_not_import() -> None: with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): - discovery = llms_package.discover_guardrail_translations() + discovery: Final = llms_package.discover_guardrail_translations() assert tuple(discovery.unavailable) == (OPENAI_CHAT_TRANSLATION_MODULE,) assert "None in sys.modules" in discovery.unavailable[OPENAI_CHAT_TRANSLATION_MODULE] @@ -78,128 +95,148 @@ def test_discovery_reports_the_handler_package_it_could_not_import(): assert CallTypes.aembedding in discovery.mappings -def test_complete_discovery_reports_nothing_unavailable(): - discovery = llms_package.discover_guardrail_translations() +def test_complete_discovery_reports_nothing_unavailable() -> None: + discovery: Final = llms_package.discover_guardrail_translations() assert not discovery.unavailable assert CallTypes.acompletion in discovery.mappings -def test_the_next_lookup_retries_a_package_that_failed_to_import(): +def test_a_lookup_after_the_retry_interval_retries_a_package_that_failed_to_import() -> None: with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): - partial = llms_package.load_guardrail_translation_mappings() + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) assert CallTypes.acompletion not in partial assert CallTypes.aembedding in partial - recovered = llms_package.load_guardrail_translation_mappings() + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) assert CallTypes.acompletion in recovered assert CallTypes.completion in recovered assert CallTypes.aembedding in recovered - assert not llms_package.guardrail_translation_discovery.unavailable -def test_a_complete_discovery_is_not_scanned_again(): - healthy = llms_package.load_guardrail_translation_mappings() +def test_a_package_that_fails_to_import_is_retried_at_most_once_per_interval() -> None: + with raising_on_import(OPENAI_CHAT_TRANSLATION_MODULE, ImportError("simulated broken install")) as loader: + for seconds in (0.0, RETRY_SECONDS / 2, RETRY_SECONDS * 0.99): + assert CallTypes.acompletion not in llms_package.load_guardrail_translation_mappings(now=at(seconds)) + assert loader.attempts == 1 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) + assert loader.attempts == 2 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 1.5)) + assert loader.attempts == 2 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 2)) + assert loader.attempts == 3 + + +def test_a_complete_discovery_is_not_scanned_again() -> None: + healthy: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) assert CallTypes.acompletion in healthy with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): - after_the_package_breaks = llms_package.load_guardrail_translation_mappings() + after_the_package_breaks: Final = llms_package.load_guardrail_translation_mappings( + now=at(RETRY_SECONDS * 10) + ) assert CallTypes.acompletion in after_the_package_breaks - assert not llms_package.guardrail_translation_discovery.unavailable -def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announced(caplog): +def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announced( + caplog: pytest.LogCaptureFixture, +) -> None: with capturing(caplog, logging.INFO), unimportable(OPENAI_CHAT_TRANSLATION_MODULE): - for _ in range(3): - llms_package.load_guardrail_translation_mappings() + for attempt in range(3): + llms_package.load_guardrail_translation_mappings(now=at(attempt * RETRY_SECONDS)) - errors = [record for record in caplog.records if record.levelno == logging.ERROR] - assert len(errors) == 1, [record.getMessage() for record in errors] - assert OPENAI_CHAT_TRANSLATION_MODULE in errors[0].getMessage() - assert "None in sys.modules" in errors[0].getMessage() + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert OPENAI_CHAT_TRANSLATION_MODULE in errors[0] + assert "None in sys.modules" in errors[0] caplog.clear() with capturing(caplog, logging.INFO): - llms_package.load_guardrail_translation_mappings() - llms_package.load_guardrail_translation_mappings() + llms_package.load_guardrail_translation_mappings(now=at(3 * RETRY_SECONDS)) + llms_package.load_guardrail_translation_mappings(now=at(4 * RETRY_SECONDS)) - recoveries = [record for record in caplog.records if "available again" in record.getMessage()] + recoveries: Final = [record for record in caplog.records if "available again" in record.getMessage()] assert len(recoveries) == 1 assert OPENAI_CHAT_TRANSLATION_MODULE in recoveries[0].getMessage() assert recoveries[0].levelno >= logging.WARNING - assert not [record for record in caplog.records if record.levelno == logging.ERROR] + assert not errors_in(caplog) -def test_lookup_recovers_after_a_failed_discovery(): +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.load_guardrail_translation_mappings(now=at(time.monotonic() + RETRY_SECONDS)) + assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) is not None -def test_an_mcp_package_that_fails_to_import_is_reported_and_retried(): - with raising_on_import(MCP_TRANSLATION_MODULE, AttributeError("module 'mcp.types' has no attribute 'ToolCall'")): - partial = llms_package.load_guardrail_translation_mappings() +def test_an_mcp_package_that_fails_to_import_is_reported_and_retried(caplog: pytest.LogCaptureFixture) -> None: + broken: Final = AttributeError("module 'mcp.types' has no attribute 'ToolCall'") + + with capturing(caplog, logging.DEBUG), raising_on_import(MCP_TRANSLATION_MODULE, broken): + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) assert CallTypes.call_mcp_tool not in partial assert CallTypes.acompletion in partial - assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,) - assert "AttributeError" in llms_package.guardrail_translation_discovery.unavailable[MCP_TRANSLATION_MODULE] + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert MCP_TRANSLATION_MODULE in errors[0] + assert "AttributeError" in errors[0] - recovered = llms_package.load_guardrail_translation_mappings() + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) assert CallTypes.call_mcp_tool in recovered assert CallTypes.acompletion in recovered - assert not llms_package.guardrail_translation_discovery.unavailable -def test_an_install_without_the_mcp_server_is_discovered_once_and_quietly(caplog): - absent = ModuleNotFoundError("No module named 'mcp'", name="mcp") +def test_an_install_without_the_mcp_server_is_discovered_once_and_quietly(caplog: pytest.LogCaptureFixture) -> None: + absent: Final = ModuleNotFoundError("No module named 'mcp'", name="mcp") with capturing(caplog, logging.DEBUG), unimportable("mcp"), raising_on_import(MCP_TRANSLATION_MODULE, absent): - first = llms_package.load_guardrail_translation_mappings() - second = llms_package.load_guardrail_translation_mappings() + first: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + much_later: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 10)) - assert first is second + assert first is much_later assert CallTypes.call_mcp_tool not in first assert CallTypes.acompletion in first - assert not llms_package.guardrail_translation_discovery.unavailable assert not [record for record in caplog.records if record.levelno >= logging.WARNING] -def test_a_broken_mcp_dependency_is_reported_and_retried(caplog): - broken = ModuleNotFoundError("No module named 'mcp.types'", name="mcp.types") +def test_a_broken_mcp_dependency_is_reported_and_retried(caplog: pytest.LogCaptureFixture) -> None: + broken: Final = ModuleNotFoundError("No module named 'mcp.types'", name="mcp.types") with capturing(caplog, logging.DEBUG), raising_on_import(MCP_TRANSLATION_MODULE, broken): - partial = llms_package.load_guardrail_translation_mappings() + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) assert CallTypes.call_mcp_tool not in partial assert CallTypes.acompletion in partial - assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,) - errors = [record for record in caplog.records if record.levelno >= logging.ERROR] - assert len(errors) == 1, [record.getMessage() for record in caplog.records] - assert "mcp.types" in errors[0].getMessage() + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert "mcp.types" in errors[0] - recovered = llms_package.load_guardrail_translation_mappings() + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) assert CallTypes.call_mcp_tool in recovered - assert not llms_package.guardrail_translation_discovery.unavailable class _StandInMCPHandler: pass -def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch): - colliding = ModuleType("litellm.llms.colliding_stub.guardrail_translation") +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} monkeypatch.setitem(sys.modules, colliding.__name__, colliding) - discovery = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={}) + discovery: Final = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={}) assert discovery.mappings[CallTypes.call_mcp_tool] is MCPGuardrailTranslationHandler diff --git a/tests/unit/proxy/guardrails/conftest.py b/tests/unit/proxy/guardrails/conftest.py index 56a93c515bd..ea11941df85 100644 --- a/tests/unit/proxy/guardrails/conftest.py +++ b/tests/unit/proxy/guardrails/conftest.py @@ -1,3 +1,5 @@ +from collections.abc import Iterator + import pytest from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( @@ -6,7 +8,7 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( @pytest.fixture(autouse=True) -def _forget_unscanned_warnings(): +def _forget_unscanned_warnings() -> Iterator[None]: unified_module._warn_left_unscanned_once.cache_clear() yield unified_module._warn_left_unscanned_once.cache_clear() 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 971f5e29a0f..e4a6f355c0f 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,5 +1,8 @@ import os import sys +import time +from collections.abc import AsyncIterator +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -532,34 +535,32 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab @pytest.fixture -def reset_guardrail_translation_discovery(): - saved = llms_package.guardrail_translation_discovery - llms_package.guardrail_translation_discovery = None - yield llms_package - llms_package.guardrail_translation_discovery = saved +def reset_guardrail_translation_discovery(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(llms_package, "guardrail_translation_discovery", None) @pytest.mark.asyncio async def test_moderation_still_runs_after_a_failed_translation_discovery( - reset_guardrail_translation_discovery, -): + reset_guardrail_translation_discovery: None, +) -> None: + a_retry_interval_ago: Final = time.monotonic() - llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS with pytest.MonkeyPatch.context() as poison: poison.setitem(sys.modules, "litellm.llms.openai.chat.guardrail_translation", None) - poisoned = llms_package.load_guardrail_translation_mappings() + poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: a_retry_interval_ago) assert CallTypes.acompletion not in poisoned with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): - openai_guardrail = OpenAIModerationGuardrail( + openai_guardrail: Final = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", event_hook="post_call", ) - unified_guardrail = UnifiedLLMGuardrails() + unified_guardrail: Final = UnifiedLLMGuardrails() - mock_mod_response = MagicMock() + mock_mod_response: Final = MagicMock() mock_mod_response.results = [] - async def mock_stream(): - chunks_data = ["Hello", " ", "world", "!", " Goodbye"] + async def mock_stream() -> AsyncIterator[ModelResponseStream]: + chunks_data: Final = ("Hello", " ", "world", "!", " Goodbye") for i, content in enumerate(chunks_data): yield ModelResponseStream( model="gpt-4", @@ -573,19 +574,20 @@ async def test_moderation_still_runs_after_a_failed_translation_discovery( ) with patch.object(openai_guardrail, "async_make_request", return_value=mock_mod_response) as patched_make_request: - chunks_received = 0 - async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), - response=mock_stream(), - request_data={ - "messages": [{"role": "user", "content": "hi"}], - "guardrail_to_apply": openai_guardrail, - "metadata": {"guardrails": ["test-openai-moderation"]}, - }, - ): - chunks_received += 1 + chunks_received: Final = [ + chunk + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), + response=mock_stream(), + request_data={ + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + }, + ) + ] - assert chunks_received == 5 + assert len(chunks_received) == 5 assert patched_make_request.await_count > 0, ( "Moderation never ran: the failed discovery was cached and the streaming hook " "passed every chunk through unscanned" 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 d84265e387c..ecf61bea77d 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 @@ -5,7 +5,7 @@ import contextlib import io import json import logging -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterator, Mapping, Sequence from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal @@ -3091,17 +3091,23 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject class TestUnscannedStreamIsAnnounced: @staticmethod - async def _drive(caplog, monkeypatch, request_route, mappings, response_chunks): + async def _drive( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + request_route: str, + mappings: Mapping[CallTypes, type[BaseTranslation]], + response_chunks: Sequence[object], + ) -> tuple[list[object], list[str]]: _patch_translation_mappings(monkeypatch, mappings) - async def stream(): + async def stream() -> AsyncIterator[object]: for chunk in response_chunks: yield chunk caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") unified_module.verbose_proxy_logger.addHandler(caplog.handler) try: - chunks = [ + chunks: Final = [ chunk async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth( @@ -3125,8 +3131,8 @@ class TestUnscannedStreamIsAnnounced: @pytest.mark.asyncio async def test_warns_when_the_call_type_has_no_translation_handler( - self, caplog, monkeypatch - ): + self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: chunks, warnings = await self._drive( caplog, monkeypatch, @@ -3150,7 +3156,7 @@ class TestUnscannedStreamIsAnnounced: ), warnings @pytest.mark.asyncio - async def test_warns_when_the_call_type_cannot_be_resolved(self, caplog, monkeypatch): + async def test_warns_when_the_call_type_cannot_be_resolved(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: chunks, warnings = await self._drive( caplog, monkeypatch, @@ -3171,7 +3177,7 @@ class TestUnscannedStreamIsAnnounced: class TestUnscannedRequestIsAnnounced: @staticmethod - def _request(guardrail): + def _request(guardrail: CustomGuardrail) -> dict[str, object]: return { "guardrail_to_apply": guardrail, "model": "gpt-4o", @@ -3180,7 +3186,7 @@ class TestUnscannedRequestIsAnnounced: @staticmethod @contextlib.contextmanager - def _capturing(caplog): + def _capturing(caplog: pytest.LogCaptureFixture) -> Iterator[None]: caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") unified_module.verbose_proxy_logger.addHandler(caplog.handler) try: @@ -3189,22 +3195,22 @@ class TestUnscannedRequestIsAnnounced: unified_module.verbose_proxy_logger.removeHandler(caplog.handler) @staticmethod - def _warnings(caplog): + def _warnings(caplog: pytest.LogCaptureFixture) -> list[str]: return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] @staticmethod - def _distinct_warnings(caplog): - by_record = {id(record): record for record in caplog.records if record.levelno >= logging.WARNING} + def _distinct_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + by_record: Final = {id(record): record for record in caplog.records if record.levelno >= logging.WARNING} return [record.getMessage() for record in by_record.values()] @pytest.mark.asyncio - async def test_pre_call_warns_when_the_call_type_has_no_translation_handler(self, caplog, monkeypatch): + async def test_pre_call_warns_when_the_call_type_has_no_translation_handler(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) - guardrail = RecordingGuardrail() - data = self._request(guardrail) + guardrail: Final = RecordingGuardrail() + data: Final = self._request(guardrail) with self._capturing(caplog): - returned = await UnifiedLLMGuardrails().async_pre_call_hook( + returned: Final = await UnifiedLLMGuardrails().async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), cache=DualCache(), data=data, @@ -3223,13 +3229,13 @@ class TestUnscannedRequestIsAnnounced: ), self._warnings(caplog) @pytest.mark.asyncio - async def test_during_call_warns_when_the_call_type_has_no_translation_handler(self, caplog, monkeypatch): + async def test_during_call_warns_when_the_call_type_has_no_translation_handler(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) - guardrail = RecordingGuardrail() - data = self._request(guardrail) + guardrail: Final = RecordingGuardrail() + data: Final = self._request(guardrail) with self._capturing(caplog): - returned = await UnifiedLLMGuardrails().async_moderation_hook( + returned: Final = await UnifiedLLMGuardrails().async_moderation_hook( data=data, user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), call_type=CallTypes.acompletion.value, @@ -3243,13 +3249,13 @@ class TestUnscannedRequestIsAnnounced: ), self._warnings(caplog) @pytest.mark.asyncio - async def test_pre_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog, monkeypatch): + async def test_pre_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 = RecordingGuardrail() - data = self._request(guardrail) + guardrail: Final = RecordingGuardrail() + data: Final = self._request(guardrail) with self._capturing(caplog): - returned = await UnifiedLLMGuardrails().async_pre_call_hook( + returned: Final = await UnifiedLLMGuardrails().async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"), cache=DualCache(), data=data, @@ -3266,9 +3272,9 @@ class TestUnscannedRequestIsAnnounced: ), self._warnings(caplog) @pytest.mark.asyncio - async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch): + async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) - guardrail = RecordingGuardrail() + guardrail: Final = RecordingGuardrail() with pytest.raises(ValueError, match="acompact_responses"): await UnifiedLLMGuardrails().async_moderation_hook( @@ -3280,9 +3286,9 @@ class TestUnscannedRequestIsAnnounced: assert guardrail.apply_calls == [] @pytest.mark.asyncio - async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog, monkeypatch): + async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) - guardrail = RecordingGuardrail() + guardrail: Final = RecordingGuardrail() with self._capturing(caplog): for _ in range(5): @@ -3299,7 +3305,7 @@ class TestUnscannedRequestIsAnnounced: call_type="aget_responses", ) - unscanned = [ + unscanned: Final = [ message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message ] assert len(unscanned) == 2, unscanned @@ -3308,9 +3314,9 @@ class TestUnscannedRequestIsAnnounced: assert "/v1/moderations" in unscanned[1] @pytest.mark.asyncio - async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog, monkeypatch): + async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) - guardrail = RecordingGuardrail() + guardrail: Final = RecordingGuardrail() with self._capturing(caplog): await UnifiedLLMGuardrails().async_post_call_success_hook( @@ -3320,15 +3326,15 @@ class TestUnscannedRequestIsAnnounced: ) assert guardrail.apply_calls == [] - unscanned = [message for message in self._warnings(caplog) if "skipping post-call scanning" in message] + unscanned: Final = [message for message in self._warnings(caplog) if "skipping post-call scanning" in message] assert unscanned, self._warnings(caplog) assert "call type 'aresponses'" in unscanned[0], unscanned[0] assert "CallTypes." not in unscanned[0], unscanned[0] @pytest.mark.asyncio - async def test_post_call_warns_on_every_unscanned_response(self, caplog, monkeypatch): + async def test_post_call_warns_on_every_unscanned_response(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) - guardrail = RecordingGuardrail() + guardrail: Final = RecordingGuardrail() with self._capturing(caplog): for _ in range(3): @@ -3338,7 +3344,7 @@ class TestUnscannedRequestIsAnnounced: response=litellm.ModelResponse(), ) - unscanned = [ + unscanned: Final = [ message for message in self._distinct_warnings(caplog) if "skipping post-call scanning" in message ] assert len(unscanned) == 3, unscanned From ec457a9794b6399cf3c3407baa64430f10268f84 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:48:59 -0700 Subject: [PATCH 6/9] test(guardrails): move the guardrail translation discovery tests to mirror litellm/llms and list them in the llms-providers shard --- .github/workflows/test-unit.yml | 1 + .../test_guardrail_translation_discovery.py | 0 2 files changed, 1 insertion(+) rename tests/unit/llms/{base_llm/guardrail_translation => }/test_guardrail_translation_discovery.py (100%) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2ddf21bca17..2a18a54ba3f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -223,6 +223,7 @@ jobs: tests/unit/llms/test_file_content_block.py tests/unit/llms/test_file_search_responses.py tests/unit/llms/test_get_complete_url_trailing_slash.py + tests/unit/llms/test_guardrail_translation_discovery.py tests/unit/llms/test_lifecycle_fix.py tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py b/tests/unit/llms/test_guardrail_translation_discovery.py similarity index 100% rename from tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py rename to tests/unit/llms/test_guardrail_translation_discovery.py 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 7/9] 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: From 6bab3b9fc1b5465818d46e02dc8d11f155468bda Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:57:23 -0700 Subject: [PATCH 8/9] fix(guardrails): keep rejecting during_call requests whose call type is outside CallTypes --- .../unified_guardrail/unified_guardrail.py | 4 ++-- .../unified_guardrails/test_unified_guardrail.py | 12 +++--------- 2 files changed, 5 insertions(+), 11 deletions(-) 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 769d960bfa9..d7528bf4232 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -421,8 +421,8 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() - resolved_call_type: Final = _resolved_call_type(call_type) - if resolved_call_type is None or resolved_call_type not in mappings: + resolved_call_type: Final = CallTypes(call_type) + if 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/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index bb34dc9253a..073ed200dc4 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 @@ -3299,24 +3299,18 @@ class TestUnscannedRequestIsAnnounced: 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: + async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch: pytest.MonkeyPatch) -> None: _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) guardrail: Final = RecordingGuardrail() - with self._capturing(caplog): - returned: Final = await UnifiedLLMGuardrails().async_moderation_hook( + with pytest.raises(ValueError, match="acompact_responses"): + 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( From 4613c265faf30b2926425c4a794c926f1920f3b0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:17:58 -0700 Subject: [PATCH 9/9] fix(guardrails): name failed handler packages in the skip warning and drop the wrong remedy --- litellm/llms/__init__.py | 5 ++ .../unified_guardrail/unified_guardrail.py | 34 +++++++----- .../test_unified_guardrail.py | 52 +++++++++++++++++-- 3 files changed, 74 insertions(+), 17 deletions(-) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 4bde909f3d0..9d77a576bfa 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -222,6 +222,11 @@ class CachedGuardrailTranslationDiscovery: guardrail_translation_discovery: CachedGuardrailTranslationDiscovery | None = None # rebind-ok: per-worker cache +def unavailable_guardrail_translation_packages() -> tuple[str, ...]: + cached: Final = guardrail_translation_discovery + return () if cached is None else tuple(cached.discovery.unavailable) + + def load_guardrail_translation_mappings( now: Callable[[], float] = time.monotonic, ) -> Mapping[CallTypes, type["BaseTranslation"]]: 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 d7528bf4232..b0db1504b5c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -26,7 +26,11 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket -from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings +from litellm.llms import ( + get_guardrail_translation_mapping, + load_guardrail_translation_mappings, + unavailable_guardrail_translation_packages, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -251,30 +255,32 @@ def _log_left_unscanned( ) -> 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: + if not call_type: verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s.", + "Guardrail '%s' selected for route '%s' but its call type could not be resolved; %s. " + "Add the route to API_ROUTE_TO_CALL_TYPES.", + guardrail_name, + route, + consequence, + ) + return + if _resolved_call_type(call_type) is None: + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but call type '%s' is not one litellm can scan; %s.", guardrail_name, route, shown_call_type, consequence, ) return - unscannable, remedy = ( - ( - 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 - else ("its call type could not be resolved", "Add the route to API_ROUTE_TO_CALL_TYPES") - ) + failed_packages: Final = unavailable_guardrail_translation_packages() verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but %s, so no guardrail can run on that route; %s. %s.", + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler%s; %s.", guardrail_name, route, - unscannable, + shown_call_type, + f", and handlers from {', '.join(failed_packages)} failed to import" if failed_packages else "", consequence, - remedy, ) 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 073ed200dc4..55bb969b004 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 @@ -5,8 +5,9 @@ import contextlib import io import json import logging +import math from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterator, Mapping, Sequence -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Literal import anyio @@ -20,7 +21,12 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings +from litellm.llms import ( + CachedGuardrailTranslationDiscovery, + GuardrailTranslationDiscovery, + discover_guardrail_translation_mappings, + load_guardrail_translation_mappings, +) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, @@ -3248,6 +3254,47 @@ class TestUnscannedRequestIsAnnounced: for message in self._warnings(caplog) ), self._warnings(caplog) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("unavailable", "named_failure"), + [ + ({"litellm.llms.openai.chat.guardrail_translation": "ImportError: broken"}, True), + ({}, False), + ], + ) + async def test_the_skip_warning_names_a_package_that_failed_to_import_only_when_one_did( + self, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + unavailable: dict[str, str], + named_failure: bool, + ) -> None: + mappings: Final = MappingProxyType({CallTypes.aembedding: _NoopTranslation}) + monkeypatch.setattr( + litellm.llms, + "guardrail_translation_discovery", + CachedGuardrailTranslationDiscovery( + discovery=GuardrailTranslationDiscovery(mappings=mappings, unavailable=MappingProxyType(unavailable)), + retry_at=math.inf, + ), + ) + _patch_translation_mappings(monkeypatch, mappings) + guardrail: Final = RecordingGuardrail() + + with self._capturing(caplog): + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), + cache=DualCache(), + data=self._request(guardrail), + call_type=CallTypes.acompletion.value, + ) + + skipped: Final = [message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message] + assert len(skipped) == 1, skipped + assert ( + "handlers from litellm.llms.openai.chat.guardrail_translation failed to import" in skipped[0] + ) is named_failure, skipped + @pytest.mark.asyncio async def test_pre_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()) @@ -3266,7 +3313,6 @@ class TestUnscannedRequestIsAnnounced: assert returned["messages"] == [{"role": "user", "content": "hello world"}] assert any( "call type 'not_a_call_type' is not one litellm can scan" in message - and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message and "skipping pre-call scanning" in message for message in self._warnings(caplog) ), self._warnings(caplog)