From 5264a70d2519f907ef994cd25507f5ed3f079be2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:15:09 -0700 Subject: [PATCH] fix(guardrails): retry a failed guardrail translation import at most once a second and type the new tests --- litellm/llms/__init__.py | 35 +++- .../test_guardrail_translation_discovery.py | 161 +++++++++++------- tests/unit/proxy/guardrails/conftest.py | 4 +- .../test_openai_moderation_streaming.py | 52 +++--- .../test_unified_guardrail.py | 76 +++++---- 5 files changed, 196 insertions(+), 132 deletions(-) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 74aaff384e6..7a9eeed6ebd 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -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 diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py b/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py index 4694eba8759..f0b577a6222 100644 --- a/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py +++ b/tests/unit/llms/base_llm/guardrail_translation/test_guardrail_translation_discovery.py @@ -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 diff --git a/tests/unit/proxy/guardrails/conftest.py b/tests/unit/proxy/guardrails/conftest.py index 56a93c515bd..ea11941df85 100644 --- a/tests/unit/proxy/guardrails/conftest.py +++ b/tests/unit/proxy/guardrails/conftest.py @@ -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() 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 971f5e29a0f..e4a6f355c0f 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,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" 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 d84265e387c..ecf61bea77d 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,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