diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index e7865817048..ec61d8cdbad 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -223,6 +223,7 @@ jobs: tests/unit/llms/test_file_content_block.py tests/unit/llms/test_file_search_responses.py tests/unit/llms/test_get_complete_url_trailing_slash.py + tests/unit/llms/test_guardrail_translation_discovery.py tests/unit/llms/test_lifecycle_fix.py tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index c18ce8cea68..9d77a576bfa 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -1,5 +1,11 @@ import importlib +import importlib.util import os +import time +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger @@ -90,101 +96,169 @@ 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"]]: - """ - Discover guardrail translation mappings by scanning the llms directory structure. +_GUARDRAIL_TRANSLATION_PACKAGE: Final = "guardrail_translation" +_MCP_GUARDRAIL_TRANSLATION_MODULE: Final = "litellm.proxy.mcp_server.guardrail_translation" +_NO_MAPPINGS: Final[Mapping[CallTypes, type["BaseTranslation"]]] = MappingProxyType({}) +GUARDRAIL_TRANSLATION_RETRY_SECONDS: Final = 1.0 - 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]]] = {} +@dataclass(frozen=True, slots=True) +class GuardrailTranslationDiscovery: + mappings: Mapping[CallTypes, type["BaseTranslation"]] + unavailable: Mapping[str, str] + +def _bundled_guardrail_translation_modules() -> Iterator[str]: + 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, ".") + + +@dataclass(frozen=True, slots=True) +class _UnavailablePackage: + reason: str + missing_dependency: bool + + +def _is_absent(module_name: str) -> bool: + root: Final = module_name.partition(".")[0] try: - # Get the path to the llms directory - current_dir: Final = os.path.dirname(__file__) - llms_dir: Final = current_dir + return importlib.util.find_spec(root) is None + except (ImportError, ValueError): + return False - 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"] +def _still_importing(module: ModuleType) -> bool: + return attribute_or_none(attribute_or_none(module, "__spec__"), "_initializing") is True - # 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 = attribute_or_none(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.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 _import_guardrail_translations( + module_path: str, +) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage: + 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 _UnavailablePackage( + reason=f"{type(e).__name__}: {e}", + missing_dependency=isinstance(e, ModuleNotFoundError) and _is_absent(e.name or ""), ) - - except Exception as e: - verbose_logger.error("Error discovering guardrail translation mappings: %s", e) - - return discovered_mappings + 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 -# Cache the discovered mappings -endpoint_guardrail_translation_mappings: dict[CallTypes, type["BaseTranslation"]] | None = None +def _guardrail_translations_from( + module_path: str, +) -> Mapping[CallTypes, type["BaseTranslation"]] | _UnavailablePackage: + 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 result -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 _guardrail_translation_modules() -> Iterator[str]: + yield from _bundled_guardrail_translation_modules() + yield _MCP_GUARDRAIL_TRANSLATION_MODULE -def get_guardrail_translation_mapping(call_type: CallTypes) -> type["BaseTranslation"]: +def _discover( + module_paths: Iterable[str], already_found: Mapping[CallTypes, type["BaseTranslation"]] +) -> GuardrailTranslationDiscovery: + 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(dict(chain.from_iterable(mappings.items() for mappings in found))), + unavailable=MappingProxyType( + {module_path: result.reason for module_path, result in imported if isinstance(result, _UnavailablePackage)} + ), + ) + + +def discover_guardrail_translations() -> GuardrailTranslationDiscovery: + return _discover(_guardrail_translation_modules(), already_found=_NO_MAPPINGS) + + +def discover_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]: + return discover_guardrail_translations().mappings + + +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 is retried at most once every %s seconds. %s", + ", ".join(current.unavailable), + GUARDRAIL_TRANSLATION_RETRY_SECONDS, + "; ".join(f"{module_path}: {reason}" for module_path, reason in current.unavailable.items()), + ) + 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.warning("Guardrail translation handlers from %s are available again.", ", ".join(recovered)) + + +@dataclass(frozen=True, slots=True) +class CachedGuardrailTranslationDiscovery: + discovery: GuardrailTranslationDiscovery + retry_at: float + + +guardrail_translation_discovery: CachedGuardrailTranslationDiscovery | None = None # rebind-ok: per-worker cache + + +def unavailable_guardrail_translation_packages() -> tuple[str, ...]: + cached: Final = guardrail_translation_discovery + return () if cached is None else tuple(cached.discovery.unavailable) + + +def load_guardrail_translation_mappings( + now: Callable[[], float] = time.monotonic, +) -> Mapping[CallTypes, type["BaseTranslation"]]: + global guardrail_translation_discovery + cached: Final = guardrail_translation_discovery + if cached is not None and not cached.discovery.unavailable: + return cached.discovery.mappings + checked_at: Final = now() + if cached is not None and checked_at < cached.retry_at: + return cached.discovery.mappings + discovery: Final = ( + discover_guardrail_translations() + if cached is None + else _discover(cached.discovery.unavailable, already_found=cached.discovery.mappings) + ) + _announce_discovery(previous=cached.discovery if cached is not None else None, current=discovery) + refreshed: Final = CachedGuardrailTranslationDiscovery( + discovery=discovery, retry_at=checked_at + GUARDRAIL_TRANSLATION_RETRY_SECONDS + ) + guardrail_translation_discovery = refreshed # rebind-ok: per-worker cache, replaced whole + return discovery.mappings + + +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 @@ -192,18 +266,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(now) + 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] 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 d713f1078a2..b0db1504b5c 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,8 @@ 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 import anyio @@ -24,7 +26,11 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket -from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings +from litellm.llms import ( + get_guardrail_translation_mapping, + load_guardrail_translation_mappings, + unavailable_guardrail_translation_packages, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -227,6 +233,95 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N data["litellm_metadata"] = user_metadata +_UNSCANNED_WARNING_KEYS: Final = 4096 + + +def _resolved_call_type(call_type: str | None) -> CallTypes | None: + try: + return CallTypes(call_type) + except ValueError: + 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 not call_type: + 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_name, + route, + consequence, + ) + return + if _resolved_call_type(call_type) is None: + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but call type '%s' is not one litellm can scan; %s.", + guardrail_name, + route, + shown_call_type, + consequence, + ) + return + failed_packages: Final = unavailable_guardrail_translation_packages() + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler%s; %s.", + guardrail_name, + route, + shown_call_type, + f", and handlers from {', '.join(failed_packages)} failed to import" if failed_packages else "", + consequence, + ) + + +@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( + guardrail_to_apply: "CustomGuardrail", + user_api_key_dict: UserAPIKeyAuth, + call_type: str | None, + consequence: str, +) -> None: + _warn_left_unscanned_once( + _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, + ) + ) + + class UnifiedLLMGuardrails(CustomLogger): def __init__( self, @@ -277,14 +372,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) @@ -329,10 +427,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 = CallTypes(call_type) + if 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) @@ -377,7 +482,7 @@ class UnifiedLLMGuardrails(CustomLogger): if user_api_key_dict.request_route is not None: call_types: Final = get_call_types_for_route(user_api_key_dict.request_route) if call_types is not None and len(call_types) > 0: - call_type = call_types[0] + call_type = call_types[0].value if call_type is None: call_type = infer_call_type(call_type=None, completion_response=response) @@ -394,28 +499,18 @@ 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, + resolved_call_type: Final = _resolved_call_type(call_type) + if resolved_call_type is None or resolved_call_type not in mappings: + _log_left_unscanned( + guardrail_name=guardrail_to_apply.guardrail_name, + request_route=user_api_key_dict.request_route, + call_type=call_type, + consequence="skipping post-call scanning", ) return response - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[resolved_call_type]()) try: response = await endpoint_translation.process_output_response( @@ -1266,6 +1361,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 diff --git a/tests/unit/llms/test_guardrail_translation_discovery.py b/tests/unit/llms/test_guardrail_translation_discovery.py new file mode 100644 index 00000000000..543313c6a51 --- /dev/null +++ b/tests/unit/llms/test_guardrail_translation_discovery.py @@ -0,0 +1,279 @@ +import importlib.abc +import importlib.machinery +import importlib.util +import logging +import sys +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from types import ModuleType +from typing import Final + +import pytest + +import litellm.llms as llms_package +from litellm._logging import verbose_logger +from litellm.proxy.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.types.utils import CallTypes + +OPENAI_CHAT_TRANSLATION_MODULE: Final = "litellm.llms.openai.chat.guardrail_translation" +MCP_TRANSLATION_MODULE: Final = "litellm.proxy.mcp_server.guardrail_translation" +RETRY_SECONDS: Final = llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS + + +def at(seconds: float) -> Callable[[], float]: + return lambda: seconds + + +class RaisingLoader(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def __init__(self, module_path: str, error: BaseException) -> None: + self.module_path = module_path + self.error = error + self.attempts = 0 + + 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.attempts += 1 + raise self.error + + +@contextmanager +def raising_on_import(module_path: str, error: BaseException) -> Iterator[RaisingLoader]: + loader: Final = RaisingLoader(module_path, error) + with pytest.MonkeyPatch.context() as mp: + mp.delitem(sys.modules, module_path, raising=False) + mp.setattr(sys, "meta_path", [loader, *sys.meta_path]) + 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: + 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 + + +def errors_in(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple(record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR) + + +@pytest.fixture(autouse=True) +def reset_guardrail_translation_discovery(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(llms_package, "guardrail_translation_discovery", None) + + +def test_discovery_reports_the_handler_package_it_could_not_import() -> None: + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + discovery: Final = 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() -> None: + discovery: Final = llms_package.discover_guardrail_translations() + + assert not discovery.unavailable + assert CallTypes.acompletion in discovery.mappings + + +def test_a_lookup_after_the_retry_interval_retries_a_package_that_failed_to_import() -> None: + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + + assert CallTypes.acompletion not in partial + assert CallTypes.aembedding in partial + + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) + + assert CallTypes.acompletion in recovered + assert CallTypes.completion in recovered + assert CallTypes.aembedding in recovered + + +def test_a_package_that_fails_to_import_is_retried_at_most_once_per_interval() -> None: + with raising_on_import(OPENAI_CHAT_TRANSLATION_MODULE, ImportError("simulated broken install")) as loader: + for seconds in (0.0, RETRY_SECONDS / 2, RETRY_SECONDS * 0.99): + assert CallTypes.acompletion not in llms_package.load_guardrail_translation_mappings(now=at(seconds)) + assert loader.attempts == 1 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) + assert loader.attempts == 2 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 1.5)) + assert loader.attempts == 2 + + llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 2)) + assert loader.attempts == 3 + + +def test_a_complete_discovery_is_not_scanned_again() -> None: + healthy: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + + assert CallTypes.acompletion in healthy + + with unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + after_the_package_breaks: Final = llms_package.load_guardrail_translation_mappings( + now=at(RETRY_SECONDS * 10) + ) + + assert CallTypes.acompletion in after_the_package_breaks + + +def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announced( + caplog: pytest.LogCaptureFixture, +) -> None: + with capturing(caplog, logging.INFO), unimportable(OPENAI_CHAT_TRANSLATION_MODULE): + for attempt in range(3): + llms_package.load_guardrail_translation_mappings(now=at(attempt * RETRY_SECONDS)) + + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert OPENAI_CHAT_TRANSLATION_MODULE in errors[0] + assert "None in sys.modules" in errors[0] + + caplog.clear() + with capturing(caplog, logging.INFO): + llms_package.load_guardrail_translation_mappings(now=at(3 * RETRY_SECONDS)) + llms_package.load_guardrail_translation_mappings(now=at(4 * RETRY_SECONDS)) + + recoveries: Final = [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 errors_in(caplog) + + +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, now=at(0.0)) + + 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: + broken: Final = AttributeError("module 'mcp.types' has no attribute 'ToolCall'") + + with capturing(caplog, logging.DEBUG), raising_on_import(MCP_TRANSLATION_MODULE, broken): + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + + assert CallTypes.call_mcp_tool not in partial + assert CallTypes.acompletion in partial + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert MCP_TRANSLATION_MODULE in errors[0] + assert "AttributeError" in errors[0] + + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) + + assert CallTypes.call_mcp_tool in recovered + assert CallTypes.acompletion in recovered + + +def test_an_install_without_the_mcp_server_is_discovered_once_and_quietly(caplog: pytest.LogCaptureFixture) -> None: + absent: Final = ModuleNotFoundError("No module named 'mcp'", name="mcp") + + with capturing(caplog, logging.DEBUG), unimportable("mcp"), raising_on_import(MCP_TRANSLATION_MODULE, absent): + first: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + much_later: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS * 10)) + + assert first is much_later + assert CallTypes.call_mcp_tool not in first + assert CallTypes.acompletion in first + assert not [record for record in caplog.records if record.levelno >= logging.WARNING] + + +def test_a_broken_mcp_dependency_is_reported_and_retried(caplog: pytest.LogCaptureFixture) -> None: + broken: Final = ModuleNotFoundError("No module named 'mcp.types'", name="mcp.types") + + with capturing(caplog, logging.DEBUG), raising_on_import(MCP_TRANSLATION_MODULE, broken): + partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0)) + + assert CallTypes.call_mcp_tool not in partial + assert CallTypes.acompletion in partial + errors: Final = errors_in(caplog) + assert len(errors) == 1, errors + assert "mcp.types" in errors[0] + + recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS)) + + assert CallTypes.call_mcp_tool in recovered + + +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} + monkeypatch.setitem(sys.modules, colliding.__name__, colliding) + + discovery: Final = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={}) + + assert discovery.mappings[CallTypes.call_mcp_tool] is MCPGuardrailTranslationHandler diff --git a/tests/unit/proxy/guardrails/conftest.py b/tests/unit/proxy/guardrails/conftest.py new file mode 100644 index 00000000000..ea11941df85 --- /dev/null +++ b/tests/unit/proxy/guardrails/conftest.py @@ -0,0 +1,14 @@ +from collections.abc import Iterator + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, +) + + +@pytest.fixture(autouse=True) +def _forget_unscanned_warnings() -> Iterator[None]: + unified_module._warn_left_unscanned_once.cache_clear() + yield + unified_module._warn_left_unscanned_once.cache_clear() 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 16f04073fae..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,14 +1,26 @@ -import pytest -from unittest.mock import MagicMock, patch import os +import sys +from collections.abc import AsyncIterator +from typing import Final +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 @@ -519,3 +531,63 @@ 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(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(llms_package, "guardrail_translation_discovery", None) + + +@pytest.mark.asyncio +async def test_moderation_still_runs_after_a_failed_translation_discovery( + reset_guardrail_translation_discovery: None, +) -> None: + 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: long_before_any_clock_reading) + assert CallTypes.acompletion not in poisoned + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail: Final = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail: Final = UnifiedLLMGuardrails() + + mock_mod_response: Final = MagicMock() + mock_mod_response.results = [] + + async def mock_stream() -> AsyncIterator[ModelResponseStream]: + chunks_data: Final = ("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: Final = [ + chunk + async for chunk 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"]}, + }, + ) + ] + + assert len(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" + ) 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 a6177ff9e3a..55bb969b004 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 @@ -5,8 +5,9 @@ import contextlib import io import json import logging -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator -from types import SimpleNamespace +import math +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterator, Mapping, Sequence +from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Literal import anyio @@ -20,7 +21,12 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings +from litellm.llms import ( + CachedGuardrailTranslationDiscovery, + GuardrailTranslationDiscovery, + discover_guardrail_translation_mappings, + load_guardrail_translation_mappings, +) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, @@ -3086,3 +3092,345 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] + + +class TestUnscannedStreamIsAnnounced: + + @staticmethod + async def _drive( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + request_route: str, + mappings: Mapping[CallTypes, type[BaseTranslation]], + response_chunks: Sequence[object], + ) -> tuple[list[object], list[str]]: + _patch_translation_mappings(monkeypatch, mappings) + + async def stream() -> AsyncIterator[object]: + 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: Final = [ + 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: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: + 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 "streaming this response to the client unscanned" 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: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: + 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 "Add the route to API_ROUTE_TO_CALL_TYPES." in message + and "recording-guardrail" in message + for message in warnings + ), warnings + + +class TestUnscannedRequestIsAnnounced: + + @staticmethod + def _request(guardrail: CustomGuardrail) -> dict[str, object]: + return { + "guardrail_to_apply": guardrail, + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello world"}], + } + + @staticmethod + @contextlib.contextmanager + def _capturing(caplog: pytest.LogCaptureFixture) -> Iterator[None]: + 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: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] + + @staticmethod + def _distinct_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + by_record: Final = {id(record): record for record in caplog.records if record.levelno >= logging.WARNING} + return [record.getMessage() for record in by_record.values()] + + @pytest.mark.asyncio + async def test_pre_call_warns_when_the_call_type_has_no_translation_handler(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail: Final = RecordingGuardrail() + data: Final = self._request(guardrail) + + with self._capturing(caplog): + returned: Final = 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: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail: Final = RecordingGuardrail() + data: Final = self._request(guardrail) + + with self._capturing(caplog): + returned: Final = 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 + @pytest.mark.parametrize( + ("unavailable", "named_failure"), + [ + ({"litellm.llms.openai.chat.guardrail_translation": "ImportError: broken"}, True), + ({}, False), + ], + ) + async def test_the_skip_warning_names_a_package_that_failed_to_import_only_when_one_did( + self, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + unavailable: dict[str, str], + named_failure: bool, + ) -> None: + mappings: Final = MappingProxyType({CallTypes.aembedding: _NoopTranslation}) + monkeypatch.setattr( + litellm.llms, + "guardrail_translation_discovery", + CachedGuardrailTranslationDiscovery( + discovery=GuardrailTranslationDiscovery(mappings=mappings, unavailable=MappingProxyType(unavailable)), + retry_at=math.inf, + ), + ) + _patch_translation_mappings(monkeypatch, mappings) + guardrail: Final = RecordingGuardrail() + + with self._capturing(caplog): + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), + cache=DualCache(), + data=self._request(guardrail), + call_type=CallTypes.acompletion.value, + ) + + skipped: Final = [message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message] + assert len(skipped) == 1, skipped + assert ( + "handlers from litellm.llms.openai.chat.guardrail_translation failed to import" in skipped[0] + ) is named_failure, skipped + + @pytest.mark.asyncio + async def test_pre_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() + data: Final = self._request(guardrail) + + with self._capturing(caplog): + returned: Final = 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="not_a_call_type", + ) + + assert guardrail.apply_calls == [] + assert returned["messages"] == [{"role": "user", "content": "hello world"}] + assert any( + "call type 'not_a_call_type' is not one litellm can scan" in message + and "skipping pre-call scanning" in message + for message in self._warnings(caplog) + ), self._warnings(caplog) + + @pytest.mark.asyncio + @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_still_rejects_a_call_type_outside_the_enum(self, 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( + 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 == [] + + 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_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 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_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: + _patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation}) + guardrail: Final = RecordingGuardrail() + + with self._capturing(caplog): + await UnifiedLLMGuardrails().async_post_call_success_hook( + data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=litellm.ModelResponse(), + ) + + assert guardrail.apply_calls == [] + unscanned: Final = [message for message in self._warnings(caplog) if "skipping post-call scanning" in message] + assert unscanned, self._warnings(caplog) + assert "call type 'aresponses'" in unscanned[0], unscanned[0] + assert "CallTypes." not in unscanned[0], unscanned[0] + + @pytest.mark.asyncio + async def test_post_call_warns_on_every_unscanned_response(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(3): + await UnifiedLLMGuardrails().async_post_call_success_hook( + data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=litellm.ModelResponse(), + ) + + unscanned: Final = [ + message for message in self._distinct_warnings(caplog) if "skipping post-call scanning" in message + ] + assert len(unscanned) == 3, unscanned diff --git a/tests/unit/proxy/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/unit/proxy/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index a6f94cbd961..e226215bf2c 100644 --- a/tests/unit/proxy/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/unit/proxy/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -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() diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index 13e4704bd90..1d11bb992da 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -19990,8 +19990,8 @@ def catalog_guardrail(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [guardrail]) monkeypatch.setattr( litellm_llms, - "endpoint_guardrail_translation_mappings", - litellm_llms.endpoint_guardrail_translation_mappings, + "guardrail_translation_discovery", + litellm_llms.guardrail_translation_discovery, ) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.slack_alerting_instance.send_alert = AsyncMock()