scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(guardrails): retry a guardrail translation package that failed to import instead of caching the gap

Browse source 
Discovery cached whichever handler map it got on the first lookup, so one package failing to import (a poisoned sys.modules entry in tests, a broken install in production) left every later guardrail lookup in that process without the handler, and the streaming hook passed responses through unscanned with nothing in the log

Discovery now records which packages failed and why, every later lookup retries only those packages until they import, the failure is logged once with its reason and the recovery once, and the streaming hook warns with the guardrail and route it left unscanned
This commit is contained in:
mateo-berri 2026-09-03 06:11:44 -07:00
parent 587311d0e4
commit ba0aa4d9d1
6 changed files with 432 additions and 110 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,8 @@
import importlib
import os
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger
@ -80,93 +83,137 @@ 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"]]:
_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({})
@dataclass(frozen=True, slots=True)
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
result apart from one that is missing handlers and therefore has to be retried.
"""
mappings: Mapping[CallTypes, type["BaseTranslation"]]
unavailable: Mapping[str, 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[:] = 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, ".")
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."""
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}"
mappings: Final = getattr(module, "guardrail_translation_mappings", None)
if not isinstance(mappings, dict):
return _NO_MAPPINGS
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 _NO_MAPPINGS
return mcp_guardrail_translation_mappings
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)))
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)}
),
)
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
"""
return _discover(
_bundled_guardrail_translation_modules(), already_found=_optional_mcp_guardrail_translation_mappings()
)
def discover_guardrail_translation_mappings() -> Mapping[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
Mapping[CallTypes, Type[BaseTranslation]]: the call types that have a translation handler class
"""
discovered_mappings: Final[dict[CallTypes, type[BaseTranslation]]] = {}
return discover_guardrail_translations().mappings
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()),
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()),
)
except Exception as e:
verbose_logger.error("Error discovering guardrail translation mappings: %s", e)
return discovered_mappings
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.info("Guardrail translation handlers from %s are available again.", ", ".join(recovered))
# Cache the discovered mappings
endpoint_guardrail_translation_mappings: dict[CallTypes, type["BaseTranslation"]] | None = None
guardrail_translation_discovery: GuardrailTranslationDiscovery | None = None
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 load_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]:
"""
Return the guardrail translation handlers, retrying any bundled package that could not be imported last time.
Serving an incomplete scan as if it were complete would silently strip the missing call types off every
guardrail for the rest of the process, so the packages that failed are imported again on each lookup, and
only the part that succeeded is kept.
"""
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"]:
@ -182,18 +229,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]

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

@ -133,6 +133,30 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N
data["litellm_metadata"] = user_metadata
def _warn_left_unscanned(
guardrail_to_apply: "CustomGuardrail",
user_api_key_dict: UserAPIKeyAuth,
call_type: str | None,
consequence: str,
) -> 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_to_apply.guardrail_name,
user_api_key_dict.request_route,
consequence,
)
return
verbose_proxy_logger.warning(
"Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s.",
guardrail_to_apply.guardrail_name,
user_api_key_dict.request_route,
call_type,
consequence,
)
class UnifiedLLMGuardrails(CustomLogger):
def __init__(
self,
@ -297,24 +321,13 @@ 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,
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="skipping post-call scanning",
)
return response
@ -1018,6 +1031,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

104
tests/test_litellm/llms/test_guardrail_translation_discovery.py Normal file
View file

@ -0,0 +1,104 @@
import logging
import sys
from collections.abc import Iterator
from contextlib import contextmanager
import pytest
import litellm.llms as llms_package
from litellm._logging import verbose_logger
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
@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_cached():
first = llms_package.load_guardrail_translation_mappings()
second = llms_package.load_guardrail_translation_mappings()
assert first is second
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 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

4
tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py
View file

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

89
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -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
@ -516,3 +526,74 @@ 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,
):
"""
A guardrail translation discovery that could not import the chat handler must not silently
disable moderation for the rest of the process.
"""
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"
)

79
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -2289,3 +2289,82 @@ class TestTranslationMappingsAreReadLive:
for name, value in vars(unified_module).items()
if isinstance(value, dict) and CallTypes.aocr in value
]
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):
_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 "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