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] 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