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]