mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(guardrails): keep a failing translation package from breaking every guarded request
Discovery caught only ImportError on the optional MCP package, so any other failure at import time propagated out of load_guardrail_translation_mappings and took down every request with a guardrail on it. Import failures are classified instead: a package that fails for any reason is reported as unavailable and retried on the next lookup, and only an mcp this install does not ship stays cached and quiet. The pre-call and during-call hooks warned nothing when they skipped a scan, and the during-call hook raised ValueError on a call type outside the enum. Both now warn with the guardrail and the route they left unscanned, and the recovery line moved to WARNING so an install running at WARNING or above sees the outage end
This commit is contained in:
parent
ba0aa4d9d1
commit
1a8d127263
4 changed files with 239 additions and 32 deletions
|
|
@ -93,7 +93,7 @@ class GuardrailTranslationDiscovery:
|
|||
"""
|
||||
The outcome of one scan for guardrail translation handlers.
|
||||
|
||||
unavailable maps each bundled package that failed to import to the reason, which is what tells a complete
|
||||
unavailable maps each package that failed to import to the reason, which is what tells a complete
|
||||
result apart from one that is missing handlers and therefore has to be retried.
|
||||
"""
|
||||
|
||||
|
|
@ -110,12 +110,30 @@ def _bundled_guardrail_translation_modules() -> Iterator[str]:
|
|||
yield "litellm." + os.path.relpath(root, os.path.dirname(llms_dir)).replace(os.sep, ".")
|
||||
|
||||
|
||||
def _import_guardrail_translations(module_path: str) -> Mapping[CallTypes, type["BaseTranslation"]] | str:
|
||||
"""Import one guardrail_translation package, returning the reason as a string when that fails."""
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _UnavailablePackage:
|
||||
"""
|
||||
Why one guardrail_translation package could not be imported.
|
||||
|
||||
missing_dependency is set when the package asked for a module outside litellm that this install does not
|
||||
have, which is the one failure that says the package is absent rather than momentarily unimportable.
|
||||
"""
|
||||
|
||||
reason: str
|
||||
missing_dependency: bool
|
||||
|
||||
|
||||
def _import_guardrail_translations(
|
||||
module_path: str,
|
||||
) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage:
|
||||
"""Import one guardrail_translation package, reporting why that failed instead of raising."""
|
||||
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 f"{type(e).__name__}: {e}"
|
||||
return _UnavailablePackage(
|
||||
reason=f"{type(e).__name__}: {e}",
|
||||
missing_dependency=isinstance(e, ModuleNotFoundError) and not (e.name or "").startswith("litellm"),
|
||||
)
|
||||
mappings: Final = getattr(module, "guardrail_translation_mappings", None)
|
||||
if not isinstance(mappings, dict):
|
||||
return _NO_MAPPINGS
|
||||
|
|
@ -123,29 +141,47 @@ def _import_guardrail_translations(module_path: str) -> Mapping[CallTypes, type[
|
|||
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)
|
||||
def _guardrail_translations_from(
|
||||
module_path: str,
|
||||
) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage:
|
||||
"""
|
||||
Import one package's handlers, tolerating an install that does not ship the optional MCP server.
|
||||
|
||||
litellm ships every package under llms, so a failure there is a gap to retry rather than a fact about the
|
||||
install. The MCP package instead arrives with the proxy extra, and a dependency it cannot import means this
|
||||
install serves no MCP endpoints for a guardrail to scan, so there is nothing to retry or report.
|
||||
"""
|
||||
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 mcp_guardrail_translation_mappings
|
||||
return result
|
||||
|
||||
|
||||
def _guardrail_translation_modules() -> Iterator[str]:
|
||||
"""Yield every module that can declare guardrail translation handlers, the optional MCP one first."""
|
||||
yield _MCP_GUARDRAIL_TRANSLATION_MODULE
|
||||
yield from _bundled_guardrail_translation_modules()
|
||||
|
||||
|
||||
def _discover(
|
||||
module_paths: Iterable[str], already_found: Mapping[CallTypes, type["BaseTranslation"]]
|
||||
) -> GuardrailTranslationDiscovery:
|
||||
imported: Final = tuple((module_path, _import_guardrail_translations(module_path)) for module_path in module_paths)
|
||||
found: Final = (already_found, *(result for _, result in imported if not isinstance(result, str)))
|
||||
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(
|
||||
{call_type: handler for mappings in found for call_type, handler in mappings.items()}
|
||||
),
|
||||
unavailable=MappingProxyType(
|
||||
{module_path: result for module_path, result in imported if isinstance(result, str)}
|
||||
{module_path: result.reason for module_path, result in imported if isinstance(result, _UnavailablePackage)}
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -155,11 +191,9 @@ 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
|
||||
GuardrailTranslationDiscovery: the handlers found, and the packages that failed to import
|
||||
"""
|
||||
return _discover(
|
||||
_bundled_guardrail_translation_modules(), already_found=_optional_mcp_guardrail_translation_mappings()
|
||||
)
|
||||
return _discover(_guardrail_translation_modules(), already_found=_NO_MAPPINGS)
|
||||
|
||||
|
||||
def discover_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]:
|
||||
|
|
@ -188,7 +222,7 @@ def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current:
|
|||
)
|
||||
if not recovered:
|
||||
return
|
||||
verbose_logger.info("Guardrail translation handlers from %s are available again.", ", ".join(recovered))
|
||||
verbose_logger.warning("Guardrail translation handlers from %s are available again.", ", ".join(recovered))
|
||||
|
||||
|
||||
guardrail_translation_discovery: GuardrailTranslationDiscovery | None = None
|
||||
|
|
|
|||
|
|
@ -133,6 +133,14 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N
|
|||
data["litellm_metadata"] = user_metadata
|
||||
|
||||
|
||||
def _resolved_call_type(call_type: str | None) -> CallTypes | None:
|
||||
"""Return the CallTypes member a route's call type names, or None when the enum has no member for it."""
|
||||
try:
|
||||
return CallTypes(call_type)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _warn_left_unscanned(
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -141,8 +149,8 @@ def _warn_left_unscanned(
|
|||
) -> None:
|
||||
if call_type is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Guardrail '%s' selected for route '%s' but its call type could not be resolved; %s. "
|
||||
"Add the route to API_ROUTE_TO_CALL_TYPES.",
|
||||
"Guardrail '%s' selected for route '%s' but its call type could not be resolved, so no guardrail "
|
||||
"can run on that route; %s.",
|
||||
guardrail_to_apply.guardrail_name,
|
||||
user_api_key_dict.request_route,
|
||||
consequence,
|
||||
|
|
@ -207,14 +215,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)
|
||||
|
||||
|
|
@ -257,10 +268,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,10 @@
|
|||
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
|
||||
|
||||
|
|
@ -10,6 +13,34 @@ from litellm._logging import verbose_logger
|
|||
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):
|
||||
"""Serves one module path, and raises the given error when Python executes it."""
|
||||
|
||||
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
|
||||
|
|
@ -93,6 +124,7 @@ def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announce
|
|||
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]
|
||||
|
||||
|
||||
|
|
@ -102,3 +134,33 @@ def test_lookup_recovers_after_a_failed_discovery():
|
|||
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), 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]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Tests for unified guardrail."""
|
||||
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -2368,3 +2369,95 @@ class TestUnscannedStreamIsAnnounced:
|
|||
and "recording-guardrail" in message
|
||||
for message in warnings
|
||||
), warnings
|
||||
|
||||
|
||||
class TestUnscannedRequestIsAnnounced:
|
||||
"""A request hook that cannot scan must say so instead of passing the request through in silence."""
|
||||
|
||||
@staticmethod
|
||||
def _request(guardrail):
|
||||
return {
|
||||
"guardrail_to_apply": guardrail,
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@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]
|
||||
|
||||
@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="moderation",
|
||||
)
|
||||
|
||||
assert guardrail.apply_calls == []
|
||||
assert returned["messages"] == [{"role": "user", "content": "hello world"}]
|
||||
assert any(
|
||||
"moderation" in message and "skipping pre-call scanning" in message
|
||||
for message in self._warnings(caplog)
|
||||
), self._warnings(caplog)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue