test(guardrails): cover a guardrail translation package that fails to import

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

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()

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()

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()

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"
+    )

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]