diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 7a9eeed6ebd..4bde909f3d0 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -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}. " diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 1d60e79b93f..769d960bfa9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -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, diff --git a/tests/unit/llms/test_guardrail_translation_discovery.py b/tests/unit/llms/test_guardrail_translation_discovery.py index f0b577a6222..543313c6a51 100644 --- a/tests/unit/llms/test_guardrail_translation_discovery.py +++ b/tests/unit/llms/test_guardrail_translation_discovery.py @@ -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} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index e4a6f355c0f..0132e666dff 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -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"}): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index ecf61bea77d..bb34dc9253a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -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: