scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(guardrails): retry guardrail handler packages that failed to import and log unscanned routes

Browse source
This commit is contained in:
mateo-berri 2026-10-20 20:47:00 -07:00
parent 52935639ad
commit 90561367f9
2 changed files with 219 additions and 116 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

213
litellm/llms/__init__.py
View file

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

122
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

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