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

fix(guardrails): stop caching a guardrail translation discovery that dropped a handler

Browse source 
An import failure inside discover_guardrail_translations() was swallowed per module and
the short mapping was then memoized process-wide, so every later request for the dropped
call type found no handler. On the streaming path that meant the selected guardrail
forwarded the whole response to the client without scanning it, and nothing was logged.

Discovery now reports which bundled handler packages it could not import, and a discovery
that lost one is returned but not cached, so the next call retries. The streaming hook
warns with the guardrail name, the route and the call type when it is about to stream a
response unscanned instead of doing it silently.
This commit is contained in:
mateo-berri 2026-09-03 06:11:44 -07:00
parent 658f50663d
commit 9e7a6e2967
5 changed files with 380 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

189
litellm/llms/__init__.py
View file

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

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

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

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

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

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

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

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

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