diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 88a44f38c57..12f9a6ab715 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -1,5 +1,8 @@ import importlib import os +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger @@ -80,92 +83,116 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo return None +_GUARDRAIL_TRANSLATION_PACKAGE: Final = "guardrail_translation" +_MCP_GUARDRAIL_TRANSLATION_MODULE: Final = "litellm.proxy._experimental.mcp_server.guardrail_translation" + + +@dataclass(frozen=True, slots=True) +class GuardrailTranslationDiscovery: + """ + The outcome of one scan for guardrail translation handlers. + + unavailable_modules names the bundled packages that failed to import, which is what tells a complete + result apart from one that is missing handlers and therefore must not be cached. + """ + + mappings: Mapping[CallTypes, type["BaseTranslation"]] + unavailable_modules: tuple[str, ...] + + +def _bundled_guardrail_translation_modules() -> Iterator[str]: + """Yield the import path of every guardrail_translation package shipped under litellm/llms.""" + llms_dir: Final = os.path.dirname(__file__) + for root, dirs, files in os.walk(llms_dir): + dirs[:] = [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, ".") + + +def _guardrail_translation_mappings_of(module_path: str) -> Mapping[CallTypes, type["BaseTranslation"]] | None: + """Import one guardrail_translation package, returning None when it could not be imported at all.""" + try: + module: Final = importlib.import_module(module_path) + except Exception as e: + verbose_logger.error("Could not import guardrail translations from %s: %s", module_path, e) + return None + mappings: Final = getattr(module, "guardrail_translation_mappings", None) + if not isinstance(mappings, dict): + return {} + declared: Final[Mapping[CallTypes, type[BaseTranslation]]] = mappings + return declared + + +def _optional_mcp_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]: + """MCP call types live outside litellm/llms and are absent from installs without the MCP server.""" + try: + from litellm.proxy._experimental.mcp_server.guardrail_translation import ( + guardrail_translation_mappings as mcp_guardrail_translation_mappings, + ) + except ImportError: + verbose_logger.debug("%s not available; skipping", _MCP_GUARDRAIL_TRANSLATION_MODULE) + return {} + return mcp_guardrail_translation_mappings + + +def discover_guardrail_translations() -> GuardrailTranslationDiscovery: + """ + Scan the llms tree, plus the optional MCP package, for guardrail translation handlers. + + Returns: + GuardrailTranslationDiscovery: the handlers found, and the bundled packages that failed to import + """ + bundled: Final = tuple( + (module_path, _guardrail_translation_mappings_of(module_path)) + for module_path in _bundled_guardrail_translation_modules() + ) + found: Final = ( + *(mappings for _, mappings in bundled if mappings is not None), + _optional_mcp_guardrail_translation_mappings(), + ) + return GuardrailTranslationDiscovery( + mappings=MappingProxyType( + {call_type: handler for mappings in found for call_type, handler in mappings.items()} + ), + unavailable_modules=tuple(module_path for module_path, mappings in bundled if mappings is None), + ) + + def discover_guardrail_translation_mappings() -> dict[CallTypes, type["BaseTranslation"]]: """ Discover guardrail translation mappings by scanning the llms directory structure. - 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]]] = {} - - try: - # Get the path to the llms directory - current_dir: Final = os.path.dirname(__file__) - llms_dir: Final = current_dir - - 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()), - ) - - except Exception as e: - verbose_logger.error("Error discovering guardrail translation mappings: %s", e) - - return discovered_mappings + return dict(discover_guardrail_translations().mappings) -# Cache the discovered mappings endpoint_guardrail_translation_mappings: dict[CallTypes, type["BaseTranslation"]] | None = None -def load_guardrail_translation_mappings(): +def load_guardrail_translation_mappings() -> dict[CallTypes, type["BaseTranslation"]]: + """ + Return the guardrail translation handlers, caching only a discovery that imported every bundled package. + + An incomplete scan is served but never cached: caching one would silently strip the missing call types + off every guardrail for the rest of the process, so the next call retries the packages that failed. + """ global endpoint_guardrail_translation_mappings - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = discover_guardrail_translation_mappings() + if endpoint_guardrail_translation_mappings is not None: + return endpoint_guardrail_translation_mappings + + discovery: Final = discover_guardrail_translations() + if discovery.unavailable_modules: + verbose_logger.error( + "Found only %s guardrail translation handlers because %s could not be imported. " + "Not caching this result: guardrails for the missing call types cannot run until the import succeeds.", + len(discovery.mappings), + ", ".join(discovery.unavailable_modules), + ) + return dict(discovery.mappings) + + endpoint_guardrail_translation_mappings = dict(discovery.mappings) return endpoint_guardrail_translation_mappings @@ -182,18 +209,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 c6b8df1b493..fe16038f676 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -136,6 +136,31 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N data["litellm_metadata"] = user_metadata +def _warn_stream_left_unscanned( + guardrail_to_apply: "CustomGuardrail", + user_api_key_dict: UserAPIKeyAuth, + call_type: str | None, + mappings: Mapping[CallTypes, type["BaseTranslation"]], +) -> None: + """Say why a selected guardrail is about to forward a whole stream without scanning it.""" + if call_type is None: + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but its call type could not be resolved; streaming this " + "response to the client unscanned. Add the route to API_ROUTE_TO_CALL_TYPES.", + guardrail_to_apply.guardrail_name, + user_api_key_dict.request_route, + ) + return + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; " + "streaming this response to the client unscanned. Available call types: %s.", + guardrail_to_apply.guardrail_name, + user_api_key_dict.request_route, + call_type, + sorted(supported.value for supported in mappings), + ) + + class UnifiedLLMGuardrails(CustomLogger): def __init__( self, @@ -1038,6 +1063,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 endpoint_guardrail_translation_mappings: + _warn_stream_left_unscanned( + guardrail_to_apply=guardrail_to_apply, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + mappings=endpoint_guardrail_translation_mappings, + ) yield item async for remaining_item in response: yield remaining_item diff --git a/tests/test_litellm/llms/test_guardrail_translation_discovery.py b/tests/test_litellm/llms/test_guardrail_translation_discovery.py new file mode 100644 index 00000000000..fc93b685c78 --- /dev/null +++ b/tests/test_litellm/llms/test_guardrail_translation_discovery.py @@ -0,0 +1,71 @@ +import sys +from contextlib import contextmanager +from typing import Iterator + +import pytest + +import litellm.llms as llms_package +from litellm.types.utils import CallTypes + +OPENAI_CHAT_TRANSLATION_MODULE = "litellm.llms.openai.chat.guardrail_translation" + + +@contextmanager +def unimportable(module_path: str) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.setitem(sys.modules, module_path, None) + yield + + +@pytest.fixture(autouse=True) +def reset_guardrail_translation_cache(): + saved = llms_package.endpoint_guardrail_translation_mappings + llms_package.endpoint_guardrail_translation_mappings = None + yield + llms_package.endpoint_guardrail_translation_mappings = 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 discovery.unavailable_modules == (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 discovery.unavailable_modules == () + assert CallTypes.acompletion in discovery.mappings + + +def test_an_incomplete_discovery_is_not_cached_for_the_rest_of_the_process(): + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + partial = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion not in partial + assert llms_package.endpoint_guardrail_translation_mappings is None + + recovered = llms_package.load_guardrail_translation_mappings() + + assert CallTypes.acompletion in recovered + assert CallTypes.completion in recovered + + +def test_a_complete_discovery_is_cached(): + first = llms_package.load_guardrail_translation_mappings() + second = llms_package.load_guardrail_translation_mappings() + + assert first is second + assert llms_package.endpoint_guardrail_translation_mappings is first + + +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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index cb6772977ec..6c122cdda42 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -7,7 +7,13 @@ from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) -from litellm.types.utils import ModelResponseStream, ModelResponse +from litellm.types.utils import ( + CallTypes, + Delta, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) from litellm.proxy._types import UserAPIKeyAuth @@ -516,3 +522,84 @@ 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_caches(monkeypatch): + import litellm.llms as llms_package + import litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail as unified_module + + saved = llms_package.endpoint_guardrail_translation_mappings + llms_package.endpoint_guardrail_translation_mappings = None + monkeypatch.setattr( + unified_module, "endpoint_guardrail_translation_mappings", None, raising=False + ) + yield llms_package + llms_package.endpoint_guardrail_translation_mappings = saved + + +@pytest.mark.asyncio +async def test_moderation_still_runs_after_a_failed_translation_discovery( + reset_guardrail_translation_caches, +): + """ + A guardrail translation discovery that could not import the chat handler must not silently + disable moderation for the rest of the process. + """ + import sys + + llms_package = reset_guardrail_translation_caches + + 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/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index a28a2a71613..8294a38bdf5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -2239,3 +2239,89 @@ class TestStreamingScanDedup: assert out == chunks assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + +class TestUnscannedStreamIsAnnounced: + """The streaming hook must never forward a whole response unscanned without saying so.""" + + @staticmethod + async def _drive(caplog, monkeypatch, request_route, mappings, response_chunks): + import litellm.llms as llms_package + + monkeypatch.setattr( + unified_module, "endpoint_guardrail_translation_mappings", None, raising=False + ) + monkeypatch.setattr( + llms_package, "endpoint_guardrail_translation_mappings", 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 "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 "recording-guardrail" in message + for message in warnings + ), warnings