mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge 4613c265fa into b93d26bc85
This commit is contained in:
rj-cairn
GitHub
commit
f2446b20b5
9 changed files with 1008 additions and 127 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
279
tests/unit/llms/test_guardrail_translation_discovery.py
Normal file
279
tests/unit/llms/test_guardrail_translation_discovery.py
Normal file
|
|
@ -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
|
||||
14
tests/unit/proxy/guardrails/conftest.py
Normal file
14
tests/unit/proxy/guardrails/conftest.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue