mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(guardrails): retry a failed guardrail translation import at most once a second and type the new tests
This commit is contained in:
mateo-berri
parent
7846534c7b
commit
5264a70d25
5 changed files with 196 additions and 132 deletions
|
|
@ -1,7 +1,8 @@
|
|||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
import time
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
|
|
@ -98,6 +99,7 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo
|
|||
_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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -191,8 +193,9 @@ def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current:
|
|||
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 every lookup retries. %s",
|
||||
"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
|
||||
|
|
@ -204,21 +207,35 @@ def _announce_discovery(previous: GuardrailTranslationDiscovery | None, current:
|
|||
verbose_logger.warning("Guardrail translation handlers from %s are available again.", ", ".join(recovered))
|
||||
|
||||
|
||||
guardrail_translation_discovery: GuardrailTranslationDiscovery | None = None
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CachedGuardrailTranslationDiscovery:
|
||||
discovery: GuardrailTranslationDiscovery
|
||||
retry_at: float
|
||||
|
||||
|
||||
def load_guardrail_translation_mappings() -> Mapping[CallTypes, type["BaseTranslation"]]:
|
||||
guardrail_translation_discovery: CachedGuardrailTranslationDiscovery | None = None # rebind-ok: per-worker cache
|
||||
|
||||
|
||||
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.unavailable:
|
||||
return cached.mappings
|
||||
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.unavailable, already_found=cached.mappings)
|
||||
else _discover(cached.discovery.unavailable, already_found=cached.discovery.mappings)
|
||||
)
|
||||
_announce_discovery(previous=cached, current=discovery)
|
||||
guardrail_translation_discovery = discovery
|
||||
_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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,13 @@
|
|||
import importlib.abc
|
||||
import importlib.machinery
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Iterator
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -15,33 +18,46 @@ from litellm.proxy.mcp_server.guardrail_translation.handler import (
|
|||
)
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
OPENAI_CHAT_TRANSLATION_MODULE = "litellm.llms.openai.chat.guardrail_translation"
|
||||
MCP_TRANSLATION_MODULE = "litellm.proxy.mcp_server.guardrail_translation"
|
||||
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=None, target=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) -> ModuleType | None:
|
||||
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[None]:
|
||||
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", [RaisingLoader(module_path, error), *sys.meta_path])
|
||||
yield
|
||||
mp.setattr(sys, "meta_path", [loader, *sys.meta_path])
|
||||
yield loader
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
|
@ -59,17 +75,18 @@ def capturing(caplog: pytest.LogCaptureFixture, level: int) -> Iterator[None]:
|
|||
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():
|
||||
saved = llms_package.guardrail_translation_discovery
|
||||
llms_package.guardrail_translation_discovery = None
|
||||
yield
|
||||
llms_package.guardrail_translation_discovery = saved
|
||||
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():
|
||||
def test_discovery_reports_the_handler_package_it_could_not_import() -> None:
|
||||
with unimportable(OPENAI_CHAT_TRANSLATION_MODULE):
|
||||
discovery = llms_package.discover_guardrail_translations()
|
||||
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]
|
||||
|
|
@ -78,128 +95,148 @@ def test_discovery_reports_the_handler_package_it_could_not_import():
|
|||
assert CallTypes.aembedding in discovery.mappings
|
||||
|
||||
|
||||
def test_complete_discovery_reports_nothing_unavailable():
|
||||
discovery = llms_package.discover_guardrail_translations()
|
||||
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_the_next_lookup_retries_a_package_that_failed_to_import():
|
||||
def test_a_lookup_after_the_retry_interval_retries_a_package_that_failed_to_import() -> None:
|
||||
with unimportable(OPENAI_CHAT_TRANSLATION_MODULE):
|
||||
partial = llms_package.load_guardrail_translation_mappings()
|
||||
partial: Final = llms_package.load_guardrail_translation_mappings(now=at(0.0))
|
||||
|
||||
assert CallTypes.acompletion not in partial
|
||||
assert CallTypes.aembedding in partial
|
||||
|
||||
recovered = llms_package.load_guardrail_translation_mappings()
|
||||
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
|
||||
assert not llms_package.guardrail_translation_discovery.unavailable
|
||||
|
||||
|
||||
def test_a_complete_discovery_is_not_scanned_again():
|
||||
healthy = llms_package.load_guardrail_translation_mappings()
|
||||
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 = llms_package.load_guardrail_translation_mappings()
|
||||
after_the_package_breaks: Final = llms_package.load_guardrail_translation_mappings(
|
||||
now=at(RETRY_SECONDS * 10)
|
||||
)
|
||||
|
||||
assert CallTypes.acompletion in after_the_package_breaks
|
||||
assert not llms_package.guardrail_translation_discovery.unavailable
|
||||
|
||||
|
||||
def test_a_package_that_keeps_failing_is_reported_once_and_its_recovery_announced(caplog):
|
||||
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 _ in range(3):
|
||||
llms_package.load_guardrail_translation_mappings()
|
||||
for attempt in range(3):
|
||||
llms_package.load_guardrail_translation_mappings(now=at(attempt * RETRY_SECONDS))
|
||||
|
||||
errors = [record for record in caplog.records if record.levelno == logging.ERROR]
|
||||
assert len(errors) == 1, [record.getMessage() for record in errors]
|
||||
assert OPENAI_CHAT_TRANSLATION_MODULE in errors[0].getMessage()
|
||||
assert "None in sys.modules" in errors[0].getMessage()
|
||||
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()
|
||||
llms_package.load_guardrail_translation_mappings()
|
||||
llms_package.load_guardrail_translation_mappings(now=at(3 * RETRY_SECONDS))
|
||||
llms_package.load_guardrail_translation_mappings(now=at(4 * RETRY_SECONDS))
|
||||
|
||||
recoveries = [record for record in caplog.records if "available again" in record.getMessage()]
|
||||
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 [record for record in caplog.records if record.levelno == logging.ERROR]
|
||||
assert not errors_in(caplog)
|
||||
|
||||
|
||||
def test_lookup_recovers_after_a_failed_discovery():
|
||||
def test_lookup_recovers_after_a_failed_discovery() -> None:
|
||||
with unimportable(OPENAI_CHAT_TRANSLATION_MODULE):
|
||||
with pytest.raises(ValueError, match="acompletion"):
|
||||
llms_package.get_guardrail_translation_mapping(CallTypes.acompletion)
|
||||
|
||||
llms_package.load_guardrail_translation_mappings(now=at(time.monotonic() + RETRY_SECONDS))
|
||||
|
||||
assert llms_package.get_guardrail_translation_mapping(CallTypes.acompletion) is not None
|
||||
|
||||
|
||||
def test_an_mcp_package_that_fails_to_import_is_reported_and_retried():
|
||||
with raising_on_import(MCP_TRANSLATION_MODULE, AttributeError("module 'mcp.types' has no attribute 'ToolCall'")):
|
||||
partial = llms_package.load_guardrail_translation_mappings()
|
||||
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
|
||||
assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,)
|
||||
assert "AttributeError" in llms_package.guardrail_translation_discovery.unavailable[MCP_TRANSLATION_MODULE]
|
||||
errors: Final = errors_in(caplog)
|
||||
assert len(errors) == 1, errors
|
||||
assert MCP_TRANSLATION_MODULE in errors[0]
|
||||
assert "AttributeError" in errors[0]
|
||||
|
||||
recovered = llms_package.load_guardrail_translation_mappings()
|
||||
recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS))
|
||||
|
||||
assert CallTypes.call_mcp_tool in recovered
|
||||
assert CallTypes.acompletion in recovered
|
||||
assert not llms_package.guardrail_translation_discovery.unavailable
|
||||
|
||||
|
||||
def test_an_install_without_the_mcp_server_is_discovered_once_and_quietly(caplog):
|
||||
absent = ModuleNotFoundError("No module named 'mcp'", name="mcp")
|
||||
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 = llms_package.load_guardrail_translation_mappings()
|
||||
second = llms_package.load_guardrail_translation_mappings()
|
||||
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 second
|
||||
assert first is much_later
|
||||
assert CallTypes.call_mcp_tool not in first
|
||||
assert CallTypes.acompletion in first
|
||||
assert not llms_package.guardrail_translation_discovery.unavailable
|
||||
assert not [record for record in caplog.records if record.levelno >= logging.WARNING]
|
||||
|
||||
|
||||
def test_a_broken_mcp_dependency_is_reported_and_retried(caplog):
|
||||
broken = ModuleNotFoundError("No module named 'mcp.types'", name="mcp.types")
|
||||
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 = llms_package.load_guardrail_translation_mappings()
|
||||
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
|
||||
assert tuple(llms_package.guardrail_translation_discovery.unavailable) == (MCP_TRANSLATION_MODULE,)
|
||||
errors = [record for record in caplog.records if record.levelno >= logging.ERROR]
|
||||
assert len(errors) == 1, [record.getMessage() for record in caplog.records]
|
||||
assert "mcp.types" in errors[0].getMessage()
|
||||
errors: Final = errors_in(caplog)
|
||||
assert len(errors) == 1, errors
|
||||
assert "mcp.types" in errors[0]
|
||||
|
||||
recovered = llms_package.load_guardrail_translation_mappings()
|
||||
recovered: Final = llms_package.load_guardrail_translation_mappings(now=at(RETRY_SECONDS))
|
||||
|
||||
assert CallTypes.call_mcp_tool in recovered
|
||||
assert not llms_package.guardrail_translation_discovery.unavailable
|
||||
|
||||
|
||||
class _StandInMCPHandler:
|
||||
pass
|
||||
|
||||
|
||||
def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch):
|
||||
colliding = ModuleType("litellm.llms.colliding_stub.guardrail_translation")
|
||||
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 = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={})
|
||||
discovery: Final = llms_package._discover((colliding.__name__, MCP_TRANSLATION_MODULE), already_found={})
|
||||
|
||||
assert discovery.mappings[CallTypes.call_mcp_tool] is MCPGuardrailTranslationHandler
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
|
||||
|
|
@ -6,7 +8,7 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
|
|||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _forget_unscanned_warnings():
|
||||
def _forget_unscanned_warnings() -> Iterator[None]:
|
||||
unified_module._warn_left_unscanned_once.cache_clear()
|
||||
yield
|
||||
unified_module._warn_left_unscanned_once.cache_clear()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -532,34 +535,32 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def reset_guardrail_translation_discovery():
|
||||
saved = llms_package.guardrail_translation_discovery
|
||||
llms_package.guardrail_translation_discovery = None
|
||||
yield llms_package
|
||||
llms_package.guardrail_translation_discovery = saved
|
||||
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,
|
||||
):
|
||||
reset_guardrail_translation_discovery: None,
|
||||
) -> None:
|
||||
a_retry_interval_ago: Final = time.monotonic() - llms_package.GUARDRAIL_TRANSLATION_RETRY_SECONDS
|
||||
with pytest.MonkeyPatch.context() as poison:
|
||||
poison.setitem(sys.modules, "litellm.llms.openai.chat.guardrail_translation", None)
|
||||
poisoned = llms_package.load_guardrail_translation_mappings()
|
||||
poisoned: Final = llms_package.load_guardrail_translation_mappings(now=lambda: a_retry_interval_ago)
|
||||
assert CallTypes.acompletion not in poisoned
|
||||
|
||||
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
|
||||
openai_guardrail = OpenAIModerationGuardrail(
|
||||
openai_guardrail: Final = OpenAIModerationGuardrail(
|
||||
guardrail_name="test-openai-moderation",
|
||||
event_hook="post_call",
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
unified_guardrail: Final = UnifiedLLMGuardrails()
|
||||
|
||||
mock_mod_response = MagicMock()
|
||||
mock_mod_response: Final = MagicMock()
|
||||
mock_mod_response.results = []
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["Hello", " ", "world", "!", " Goodbye"]
|
||||
async def mock_stream() -> AsyncIterator[ModelResponseStream]:
|
||||
chunks_data: Final = ("Hello", " ", "world", "!", " Goodbye")
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield ModelResponseStream(
|
||||
model="gpt-4",
|
||||
|
|
@ -573,19 +574,20 @@ async def test_moderation_still_runs_after_a_failed_translation_discovery(
|
|||
)
|
||||
|
||||
with patch.object(openai_guardrail, "async_make_request", return_value=mock_mod_response) as patched_make_request:
|
||||
chunks_received = 0
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
|
||||
response=mock_stream(),
|
||||
request_data={
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": openai_guardrail,
|
||||
"metadata": {"guardrails": ["test-openai-moderation"]},
|
||||
},
|
||||
):
|
||||
chunks_received += 1
|
||||
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 chunks_received == 5
|
||||
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,7 +5,7 @@ import contextlib
|
|||
import io
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterator, Mapping, Sequence
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
|
|
@ -3091,17 +3091,23 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject
|
|||
class TestUnscannedStreamIsAnnounced:
|
||||
|
||||
@staticmethod
|
||||
async def _drive(caplog, monkeypatch, request_route, mappings, response_chunks):
|
||||
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():
|
||||
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 = [
|
||||
chunks: Final = [
|
||||
chunk
|
||||
async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
|
|
@ -3125,8 +3131,8 @@ class TestUnscannedStreamIsAnnounced:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_warns_when_the_call_type_has_no_translation_handler(
|
||||
self, caplog, monkeypatch
|
||||
):
|
||||
self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
chunks, warnings = await self._drive(
|
||||
caplog,
|
||||
monkeypatch,
|
||||
|
|
@ -3150,7 +3156,7 @@ class TestUnscannedStreamIsAnnounced:
|
|||
), warnings
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_warns_when_the_call_type_cannot_be_resolved(self, caplog, monkeypatch):
|
||||
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,
|
||||
|
|
@ -3171,7 +3177,7 @@ class TestUnscannedStreamIsAnnounced:
|
|||
class TestUnscannedRequestIsAnnounced:
|
||||
|
||||
@staticmethod
|
||||
def _request(guardrail):
|
||||
def _request(guardrail: CustomGuardrail) -> dict[str, object]:
|
||||
return {
|
||||
"guardrail_to_apply": guardrail,
|
||||
"model": "gpt-4o",
|
||||
|
|
@ -3180,7 +3186,7 @@ class TestUnscannedRequestIsAnnounced:
|
|||
|
||||
@staticmethod
|
||||
@contextlib.contextmanager
|
||||
def _capturing(caplog):
|
||||
def _capturing(caplog: pytest.LogCaptureFixture) -> Iterator[None]:
|
||||
caplog.set_level(logging.WARNING, logger="LiteLLM Proxy")
|
||||
unified_module.verbose_proxy_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
|
|
@ -3189,22 +3195,22 @@ class TestUnscannedRequestIsAnnounced:
|
|||
unified_module.verbose_proxy_logger.removeHandler(caplog.handler)
|
||||
|
||||
@staticmethod
|
||||
def _warnings(caplog):
|
||||
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):
|
||||
by_record = {id(record): record for record in caplog.records if record.levelno >= logging.WARNING}
|
||||
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, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
data = self._request(guardrail)
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
data: Final = self._request(guardrail)
|
||||
|
||||
with self._capturing(caplog):
|
||||
returned = await UnifiedLLMGuardrails().async_pre_call_hook(
|
||||
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,
|
||||
|
|
@ -3223,13 +3229,13 @@ class TestUnscannedRequestIsAnnounced:
|
|||
), self._warnings(caplog)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_warns_when_the_call_type_has_no_translation_handler(self, caplog, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
data = self._request(guardrail)
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
data: Final = self._request(guardrail)
|
||||
|
||||
with self._capturing(caplog):
|
||||
returned = await UnifiedLLMGuardrails().async_moderation_hook(
|
||||
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,
|
||||
|
|
@ -3243,13 +3249,13 @@ class TestUnscannedRequestIsAnnounced:
|
|||
), self._warnings(caplog)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_warns_instead_of_raising_on_a_call_type_outside_the_enum(self, caplog, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
data = self._request(guardrail)
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
data: Final = self._request(guardrail)
|
||||
|
||||
with self._capturing(caplog):
|
||||
returned = await UnifiedLLMGuardrails().async_pre_call_hook(
|
||||
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,
|
||||
|
|
@ -3266,9 +3272,9 @@ class TestUnscannedRequestIsAnnounced:
|
|||
), self._warnings(caplog)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_still_rejects_a_call_type_outside_the_enum(self, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
|
||||
with pytest.raises(ValueError, match="acompact_responses"):
|
||||
await UnifiedLLMGuardrails().async_moderation_hook(
|
||||
|
|
@ -3280,9 +3286,9 @@ class TestUnscannedRequestIsAnnounced:
|
|||
assert guardrail.apply_calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog, monkeypatch):
|
||||
async def test_a_route_with_no_handler_warns_once_instead_of_once_per_request(self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_translation_mappings(monkeypatch, {CallTypes.aembedding: _NoopTranslation})
|
||||
guardrail = RecordingGuardrail()
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
|
||||
with self._capturing(caplog):
|
||||
for _ in range(5):
|
||||
|
|
@ -3299,7 +3305,7 @@ class TestUnscannedRequestIsAnnounced:
|
|||
call_type="aget_responses",
|
||||
)
|
||||
|
||||
unscanned = [
|
||||
unscanned: Final = [
|
||||
message for message in self._distinct_warnings(caplog) if "skipping pre-call scanning" in message
|
||||
]
|
||||
assert len(unscanned) == 2, unscanned
|
||||
|
|
@ -3308,9 +3314,9 @@ class TestUnscannedRequestIsAnnounced:
|
|||
assert "/v1/moderations" in unscanned[1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_names_the_call_type_the_route_maps_to(self, caplog, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
|
||||
with self._capturing(caplog):
|
||||
await UnifiedLLMGuardrails().async_post_call_success_hook(
|
||||
|
|
@ -3320,15 +3326,15 @@ class TestUnscannedRequestIsAnnounced:
|
|||
)
|
||||
|
||||
assert guardrail.apply_calls == []
|
||||
unscanned = [message for message in self._warnings(caplog) if "skipping post-call scanning" in message]
|
||||
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, monkeypatch):
|
||||
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 = RecordingGuardrail()
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
|
||||
with self._capturing(caplog):
|
||||
for _ in range(3):
|
||||
|
|
@ -3338,7 +3344,7 @@ class TestUnscannedRequestIsAnnounced:
|
|||
response=litellm.ModelResponse(),
|
||||
)
|
||||
|
||||
unscanned = [
|
||||
unscanned: Final = [
|
||||
message for message in self._distinct_warnings(caplog) if "skipping post-call scanning" in message
|
||||
]
|
||||
assert len(unscanned) == 3, unscanned
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue