scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(guardrails): retry a handler package caught mid-import, skip unscannable during_call call types, and warn once per call type

Browse source 
A package found while it is still being imported is now treated as failed, so the next lookup retries it instead of caching a partial map.

during_call now skips a call type outside CallTypes with a warning, as pre_call and post_call do, instead of answering HTTP 500 on reads like GET /v1/responses/{id}.

The skip warning is keyed on the guardrail, the call type, and the hook, so resource ids in the path no longer repeat it. Route and call type values are stripped of line breaks before they are logged.
This commit is contained in:
mateo-berri 2026-10-29 00:20:51 -07:00
parent ec457a9794
commit bf52898fec
5 changed files with 161 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/llms/__init__.py
View file

@ -5,7 +5,7 @@ import time
from collections.abc import Callable, Iterable, Iterator, Mapping
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger
@ -130,6 +130,10 @@ def _is_absent(module_name: str) -> bool:
return False
def _still_importing(module: ModuleType) -> bool:
return attribute_or_none(attribute_or_none(module, "__spec__"), "_initializing") is True
def _import_guardrail_translations(
module_path: str,
) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage:
@ -140,6 +144,8 @@ def _import_guardrail_translations(
reason=f"{type(e).__name__}: {e}",
missing_dependency=isinstance(e, ModuleNotFoundError) and _is_absent(e.name or ""),
)
if _still_importing(module):
return _UnavailablePackage(reason=f"{module_path} was still being imported", missing_dependency=False)
declared: Final = attribute_or_none(module, "guardrail_translation_mappings")
return declared if isinstance(declared, dict) else _NO_MAPPINGS
@ -239,12 +245,15 @@ def load_guardrail_translation_mappings(
return discovery.mappings
def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTranslation"]:
def get_guardrail_translation_mapping(
call_type: CallTypes, now: Callable[[], float] = time.monotonic
) -> type["BaseTranslation"]:
"""
Get the guardrail translation handler for a given call type.
Args:
call_type: The type of call (e.g., completion, acompletion, anthropic_messages)
now: Clock that decides whether a package that failed to import is due for another try
Returns:
The translation handler class for the given call type
@ -252,7 +261,7 @@ def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTransla
Raises:
ValueError: If no translation mapping exists for the given call type
"""
mappings: Final = load_guardrail_translation_mappings()
mappings: Final = load_guardrail_translation_mappings(now)
if call_type not in mappings:
raise ValueError(
f"No guardrail translation mapping found for call_type: {call_type}. "

50
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 dataclasses import dataclass, field
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable
@ -238,24 +239,30 @@ def _resolved_call_type(call_type: str | None) -> CallTypes | None:
return None
def _without_line_breaks(value: str | None) -> str:
return str(value).replace("\r", "").replace("\n", "")
def _log_left_unscanned(
guardrail_name: str | None,
request_route: str | None,
call_type: str | None,
consequence: str,
) -> None:
route: Final = _without_line_breaks(request_route)
shown_call_type: Final = _without_line_breaks(call_type)
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.",
guardrail_name,
request_route,
call_type,
route,
shown_call_type,
consequence,
)
return
unscannable, remedy = (
(
f"call type '{call_type}' is not one litellm can scan",
f"call type '{shown_call_type}' is not one litellm can scan",
"Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES",
)
if call_type
@ -264,14 +271,32 @@ def _log_left_unscanned(
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,
route,
unscannable,
consequence,
remedy,
)
_warn_left_unscanned_once: Final = lru_cache(maxsize=_UNSCANNED_WARNING_KEYS)(_log_left_unscanned)
@dataclass(frozen=True, slots=True)
class _UnscannedWarning:
guardrail_name: str | None
call_type_or_route: str | None
consequence: str
request_route: str | None = field(compare=False)
call_type: str | None = field(compare=False)
def _log_unscanned_warning(warning: _UnscannedWarning) -> None:
_log_left_unscanned(
guardrail_name=warning.guardrail_name,
request_route=warning.request_route,
call_type=warning.call_type,
consequence=warning.consequence,
)
_warn_left_unscanned_once: Final = lru_cache(maxsize=_UNSCANNED_WARNING_KEYS)(_log_unscanned_warning)
def _warn_left_unscanned(
@ -281,10 +306,13 @@ def _warn_left_unscanned(
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,
_UnscannedWarning(
guardrail_name=guardrail_to_apply.guardrail_name,
call_type_or_route=call_type or user_api_key_dict.request_route,
consequence=consequence,
request_route=user_api_key_dict.request_route,
call_type=call_type,
)
)
@ -393,8 +421,8 @@ class UnifiedLLMGuardrails(CustomLogger):
return data
mappings: Final = load_guardrail_translation_mappings()
resolved_call_type: Final = CallTypes(call_type)
if resolved_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,

47
tests/unit/llms/test_guardrail_translation_discovery.py
View file

@ -3,7 +3,6 @@ import importlib.machinery
import importlib.util
import logging
import sys
import time
from collections.abc import Callable, Iterator, Sequence
from contextlib import contextmanager
from types import ModuleType
@ -60,6 +59,29 @@ def raising_on_import(module_path: str, error: BaseException) -> Iterator[Raisin
yield loader
class DiscoveringDuringItsOwnImport(importlib.abc.MetaPathFinder, importlib.abc.Loader):
def __init__(self, module_path: str) -> None:
self.module_path = module_path
self.seen_mid_import: llms_package.GuardrailTranslationDiscovery | None = None
def find_spec(
self,
fullname: str,
path: Sequence[str] | None = None,
target: ModuleType | None = None,
) -> importlib.machinery.ModuleSpec | None:
if fullname != self.module_path:
return None
return importlib.util.spec_from_loader(fullname, self)
def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType | None:
return None
def exec_module(self, module: ModuleType) -> None:
self.seen_mid_import = llms_package._discover((self.module_path,), already_found={})
module.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler}
@contextmanager
def unimportable(module_path: str) -> Iterator[None]:
with pytest.MonkeyPatch.context() as mp:
@ -172,11 +194,9 @@ def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announce
def test_lookup_recovers_after_a_failed_discovery() -> None:
with unimportable(OPENAI_CHAT_TRANSLATION_MODULE):
with pytest.raises(ValueError, match="acompletion"):
llms_package.get_guardrail_translation_mapping(CallTypes.acompletion)
llms_package.get_guardrail_translation_mapping(CallTypes.acompletion, now=at(0.0))
llms_package.load_guardrail_translation_mappings(now=at(time.monotonic() + RETRY_SECONDS))
assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) is not None
assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion, now=at(RETRY_SECONDS)) is not None
def test_an_mcp_package_that_fails_to_import_is_reported_and_retried(caplog: pytest.LogCaptureFixture) -> None:
@ -232,6 +252,23 @@ class _StandInMCPHandler:
pass
def test_a_package_caught_mid_import_is_retried_instead_of_cached_as_complete() -> None:
module_path: Final = "mid_import_guardrail_translation_stub"
loader: Final = DiscoveringDuringItsOwnImport(module_path)
with pytest.MonkeyPatch.context() as mp:
mp.setattr(sys, "meta_path", [loader, *sys.meta_path])
importlib.import_module(module_path)
mid_import: Final = loader.seen_mid_import
assert mid_import is not None
retried: Final = llms_package._discover(mid_import.unavailable, already_found=mid_import.mappings)
mp.delitem(sys.modules, module_path)
assert tuple(mid_import.unavailable) == (module_path,)
assert CallTypes.call_mcp_tool not in mid_import.mappings
assert not retried.unavailable
assert retried.mappings[CallTypes.call_mcp_tool] is _StandInMCPHandler
def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch: pytest.MonkeyPatch) -> None:
colliding: Final = ModuleType("litellm.llms.colliding_stub.guardrail_translation")
colliding.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler}

5
tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -1,6 +1,5 @@
import os
import sys
import time
from collections.abc import AsyncIterator
from typing import Final
from unittest.mock import MagicMock, patch
@ -543,10 +542,10 @@ def reset_guardrail_translation_discovery(monkeypatch: pytest.MonkeyPatch) -> No
async def test_moderation_still_runs_after_a_failed_translation_discovery(
reset_guardrail_translation_discovery: None,
) -> None:
a_retry_interval_ago: Final = time.monotonic() - llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS
long_before_any_clock_reading: Final = float("-inf")
with pytest.MonkeyPatch.context() as poison:
poison.setitem(sys.modules, "litellm.llms.openai.chat.guardrail_translation", None)
poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: a_retry_interval_ago)
poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: long_before_any_clock_reading)
assert CallTypes.acompletion not in poisoned
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):

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

@ -3272,46 +3272,92 @@ class TestUnscannedRequestIsAnnounced:
), self._warnings(caplog)
@pytest.mark.asyncio
async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.parametrize(
("call_type", "forged_text"),
[
pytest.param(CallTypes.acompletion.value, "/v1/responses/resp_1ERROR forged entry", id="handled-call-type"),
pytest.param("not_a_call_type\r\nERROR forged entry", "not_a_call_typeERROR forged entry", id="call-type-outside-the-enum"),
],
)
async def test_line_breaks_in_the_route_or_call_type_cannot_forge_log_lines(
self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, call_type: str, forged_text: str
) -> None:
_patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation})
with self._capturing(caplog):
await UnifiedLLMGuardrails().async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key", request_route="/v1/responses/resp_1\r\nERROR forged entry"
),
cache=DualCache(),
data=self._request(RecordingGuardrail()),
call_type=call_type,
)
warnings: Final = self._warnings(caplog)
assert any(forged_text in message for message in warnings), warnings
assert not any("\r" in message or "\n" in message for message in warnings), warnings
@pytest.mark.asyncio
async def test_during_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
guardrail: Final = RecordingGuardrail()
with pytest.raises(ValueError, match="acompact_responses"):
await UnifiedLLMGuardrails().async_moderation_hook(
with self._capturing(caplog):
returned: Final = await UnifiedLLMGuardrails().async_moderation_hook(
data=self._request(guardrail),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses/compact"),
call_type="acompact_responses",
)
assert guardrail.apply_calls == []
assert returned["messages"] == [{"role": "user", "content": "hello world"}]
assert any(
"call type 'acompact_responses' is not one litellm can scan" in message
and "skipping during-call scanning" in message
for message in self._warnings(caplog)
), self._warnings(caplog)
async def _pre_call_on(self, guardrail: CustomGuardrail, request_route: str, call_type: str | None) -> None:
await UnifiedLLMGuardrails().async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route=request_route),
cache=DualCache(),
data=self._request(guardrail),
call_type=call_type,
)
@pytest.mark.asyncio
async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
async def test_a_call_type_with_no_handler_warns_once_however_many_resource_ids_reach_it(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
_patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation})
guardrail: Final = RecordingGuardrail()
with self._capturing(caplog):
for _ in range(5):
await UnifiedLLMGuardrails().async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses/resp_123"),
cache=DualCache(),
data=self._request(guardrail),
call_type="aget_responses",
)
await UnifiedLLMGuardrails().async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/moderations"),
cache=DualCache(),
data=self._request(guardrail),
call_type="aget_responses",
)
for response_id in ("resp_1", "resp_2", "resp_3", "resp_1"):
await self._pre_call_on(guardrail, f"/v1/responses/{response_id}", "aget_responses")
await self._pre_call_on(guardrail, "/v1/responses/resp_1", "adelete_responses")
unscanned: Final = [
message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message
]
assert len(unscanned) == 2, unscanned
assert "/v1/responses/resp_123" in unscanned[0]
assert "aget_responses" in unscanned[0]
assert "/v1/moderations" in unscanned[1]
assert "/v1/responses/resp_1" in unscanned[0] and "aget_responses" in unscanned[0]
assert "adelete_responses" in unscanned[1]
@pytest.mark.asyncio
async def test_routes_without_a_call_type_each_warn_once(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
_patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation})
guardrail: Final = RecordingGuardrail()
with self._capturing(caplog):
for request_route in ("/custom/route_a", "/custom/route_b", "/custom/route_a"):
await self._pre_call_on(guardrail, request_route, None)
unscanned: Final = [
message for message in self._distinct_warnings(caplog) if "its call type could not be resolved" in message
]
assert len(unscanned) == 2, unscanned
assert "/custom/route_a" in unscanned[0]
assert "/custom/route_b" in unscanned[1]
@pytest.mark.asyncio
async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: