From f359ee77e05f80d5023aebd781c6d6cc3b384f49 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 11:31:32 -0700 Subject: [PATCH 1/8] fix(proxy): skip pre-call guardrails for a model the proxy cannot route --- litellm/proxy/common_request_processing.py | 47 +- .../authorization/_operator_guardrail.py | 24 + .../test_unknown_model_guardrail_skip.py | 217 ++++ .../agent_endpoints/test_a2a_endpoints.py | 8 +- .../proxy/mcp_server/test_rest_endpoints.py | 10 +- .../proxy/test_common_request_processing.py | 1084 ++++++++++++++++- .../unit/proxy/test_model_level_guardrails.py | 2 +- 7 files changed, 1364 insertions(+), 28 deletions(-) create mode 100644 tests/integration/authorization/_operator_guardrail.py create mode 100644 tests/integration/authorization/test_unknown_model_guardrail_skip.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9282ee7f703..69284fa64e8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -133,6 +133,8 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( + RouteType, + UnroutableModelProbe, endpoint_unsupported_by_model, route_request, ) @@ -143,6 +145,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ProxyLogging, _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, + without_unscanned_content, ) from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task @@ -2166,6 +2169,7 @@ class ProxyBaseLLMRequestProcessing: self.served_model = served_model self._tags_before_guardrails: frozenset[str] | None = None self.route_type: str | None = None + self._routing_probe: UnroutableModelProbe | None = None @property def litellm_call_id(self) -> str | None: @@ -2459,6 +2463,7 @@ class ProxyBaseLLMRequestProcessing: rate_limited_model: str | None = None, skip_guardrails: bool = False, logging_auth: UserAPIKeyAuth | None = None, + dispatch_route_type: RouteType | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2640,12 +2645,29 @@ class ProxyBaseLLMRequestProcessing: prefetch_model = self.data.get("model") if llm_router is not None and isinstance(prefetch_model, str): RoutingContext(llm_router).arm_read_prefetch(prefetch_model, self.data) - self.data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, - data=self.data, - call_type=route_type, - skip_guardrails=skip_guardrails, + routing_probe: Final = ( + None + if dispatch_route_type is None + else UnroutableModelProbe( + llm_router=llm_router, + user_model=user_model, + route_type=dispatch_route_type, + user_api_key_dict=user_api_key_dict, + ) ) + self._routing_probe = routing_probe + try: + self.data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=self.data, + call_type=route_type, + skip_guardrails=skip_guardrails, + routing_probe=routing_probe, + ) + except Exception: + if routing_probe is not None and routing_probe.refusal is not None: + self._redact_request_its_skipped_guardrails_never_scanned() + raise await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, @@ -2681,6 +2703,13 @@ class ProxyBaseLLMRequestProcessing: return self.data, logging_obj + def _redact_request_its_skipped_guardrails_never_scanned(self) -> None: + """Redact the request body and its logged snapshot, so failure callbacks never see content a skipped + masking guardrail would have rewritten. + """ + self.data.update(without_unscanned_content(self.data)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + refresh_proxy_server_request_body_snapshot(self.data) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + async def _pre_call_with_fallbacks( self, request: Request, @@ -2697,6 +2726,7 @@ class ProxyBaseLLMRequestProcessing: model: str | None, route_type: str, llm_router: Router | None, + dispatch_route_type: RouteType | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ( PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, @@ -2737,6 +2767,7 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, + dispatch_route_type=dispatch_route_type, ) except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data @@ -2811,6 +2842,7 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, rate_limited_model=original_model, + dispatch_route_type=dispatch_route_type, ) except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): @@ -3066,8 +3098,13 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, + dispatch_route_type=route_type, ) phase_event("litellm.request.pre_call_completed") + refusal: Final = None if self._routing_probe is None else self._routing_probe.refusal + if refusal is not None: + self._redact_request_its_skipped_guardrails_never_scanned() + raise refusal _post_call_guardrails_active: Final = self._has_post_call_guardrails() diff --git a/tests/integration/authorization/_operator_guardrail.py b/tests/integration/authorization/_operator_guardrail.py new file mode 100644 index 00000000000..dec071acd9a --- /dev/null +++ b/tests/integration/authorization/_operator_guardrail.py @@ -0,0 +1,24 @@ +"""A pre_call guardrail loaded from the operator's own module, the way a deployment's custom guardrail is.""" + +import os +from typing import Final + +import httpx +from fastapi import HTTPException + +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + + +class OperatorGuardrail(CustomGuardrail): + """Posts the requested model to ``OPERATOR_GUARDRAIL_URL`` and rejects with 403 when the reply says block.""" + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> None: + async with httpx.AsyncClient(trust_env=False) as client: + response: Final = await client.post(os.environ["OPERATOR_GUARDRAIL_URL"], json={"model": data.get("model")}) + if response.json()["action"] == "block": + raise HTTPException(status_code=403, detail="blocked by the operator guardrail") diff --git a/tests/integration/authorization/test_unknown_model_guardrail_skip.py b/tests/integration/authorization/test_unknown_model_guardrail_skip.py new file mode 100644 index 00000000000..d502eaa07ca --- /dev/null +++ b/tests/integration/authorization/test_unknown_model_guardrail_skip.py @@ -0,0 +1,217 @@ +"""LIT-8011: a model name the proxy cannot route gets the same ``400 Invalid model name`` +it gets with no guardrail, and a ``default_on`` pre_call guardrail never sends it to the +vendor. A routable model is still scanned before the provider call, and a guardrail from the +operator's own module still runs, along with every guardrail listed before it. +""" + +import json +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL_TEXT: Final = "routed reply" +_VERDICTS: Final = { + "allow": Reply(body=json.dumps({"action": "allow", "category": "benign", "response_detected": {}}).encode()), + "block": Reply( + body=json.dumps( + {"action": "block", "category": "malicious", "prompt_detected": {"dlp": True}, "response_detected": {}} + ).encode() + ), + "unreachable": Reply(status=503, body=b'{"error": "scanner down"}'), +} + + +def _scanner(verdict: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + assert request.target == "/v1/scan/sync/request" + return _VERDICTS[verdict] + + return reply + + +def _provider(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + assert request.target == "/v1/chat/completions", request.target + return Reply( + body=json.dumps( + { + "id": "chatcmpl_synthetic", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4.1-mini", + "choices": [ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": _MODEL_TEXT}} + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + } + ).encode() + ) + + +def _operator_verdict(action: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + return Reply(body=json.dumps({"action": action}).encode()) + + return reply + + +def _airs(policy_url: str) -> dict[str, JsonValue]: + return { + "guardrail_name": "airs-" + uuid.uuid4().hex[:8], + "litellm_params": { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "default_on": True, + "api_base": policy_url, + "api_key": "synthetic-panw-key", + "profile_name": "synthetic-profile", + }, + } + + +def _operator_guardrail() -> dict[str, JsonValue]: + return { + "guardrail_name": "operator-" + uuid.uuid4().hex[:8], + "litellm_params": { + "guardrail": "integration.authorization._operator_guardrail.OperatorGuardrail", + "mode": "pre_call", + "default_on": True, + }, + } + + +def _config(tmp_path: Path, *guardrails: Mapping[str, JsonValue]) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = list(guardrails) + path: Final = tmp_path / "pre_call_guardrails.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _requests(model: str) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + messages: Final = [{"role": "user", "content": "my card is 4111 1111 1111 1111"}] + return ( + ("/v1/chat/completions", {"model": model, "messages": messages}), + ("/v1/chat/completions", {"model": model, "messages": messages, "stream": True}), + ("/v1/responses", {"model": model, "input": "my card is 4111 1111 1111 1111"}), + ("/v1/messages", {"model": model, "max_tokens": 16, "messages": messages}), + ("/v1/embeddings", {"model": model, "input": "my card is 4111 1111 1111 1111"}), + ("/v1/completions", {"model": model, "prompt": "my card is 4111 1111 1111 1111"}), + ) + + +@pytest.mark.parametrize("verdict", sorted(_VERDICTS)) +def test_unknown_model_gets_routing_400_without_calling_pre_call_guardrail( + gateway: Gateway, tmp_path: Path, verdict: str +) -> None: + model: Final = "no-such-model-" + uuid.uuid4().hex[:8] + with wire_server(_scanner(verdict)) as policy: + with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate: + for target, body in _requests(model): + control = gateway.request("POST", target, body) + response = candidate.request("POST", target, body) + + assert control.status_code == 400, control.text + assert "Invalid model name" in control.text + assert (response.status_code, response.json()) == (control.status_code, control.json()), target + assert policy.drain() == () + + +def test_routable_model_is_still_scanned_before_the_provider_call(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_scanner("allow")) as policy, wire_server(_provider) as upstream: + with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key" + ) + response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1])) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == _MODEL_TEXT + scans: Final = policy.drain() + assert len(scans) == 1 + assert "4111" in json.loads(scans[0].body)["contents"][0]["prompt"] + + +def test_routable_model_block_verdict_still_blocks(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_scanner("block")) as policy, wire_server(_provider) as upstream: + with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key" + ) + response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1])) + + assert response.status_code == 400, response.text + assert "Invalid model name" not in response.text + assert len(policy.drain()) == 1 + assert [request.target for request in upstream.drain() if request.method == "POST"] == [] + + +def test_unknown_model_burst_while_the_scanner_is_down_never_reaches_it(gateway: Gateway, tmp_path: Path) -> None: + model: Final = "no-such-model-" + uuid.uuid4().hex[:8] + unknown: Final = _requests(model) + controls: Final = tuple(gateway.request("POST", target, body) for target, body in unknown) + with wire_server(_scanner("unreachable")) as policy, wire_server(_provider) as upstream: + with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url)), workers=2) as candidate: + with candidate.scenario() as scenario: + routable: Final = scenario.model( + model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key" + ) + known: Final = ( + "/v1/chat/completions", + {"model": routable, "messages": [{"role": "user", "content": "routable burst prompt"}]}, + ) + burst: Final = unknown * 6 + (known,) * 6 + with ThreadPoolExecutor(max_workers=16) as pool: + responses: Final = tuple(pool.map(lambda item: candidate.request("POST", *item), burst)) + + for index, response in enumerate(responses[: len(unknown) * 6]): + control = controls[index % len(unknown)] + assert (response.status_code, response.json()) == (control.status_code, control.json()), index + known_statuses: Final = {response.status_code for response in responses[len(unknown) * 6 :]} + assert 200 not in known_statuses, known_statuses + scanned: Final = [json.loads(scan.body)["contents"][0]["prompt"] for scan in policy.drain()] + assert scanned and all("routable burst prompt" in prompt for prompt in scanned), scanned + assert [request.target for request in upstream.drain() if request.method == "POST"] == [] + + +def test_operator_guardrail_and_the_guardrail_before_it_still_run_for_an_unknown_model( + gateway: Gateway, tmp_path: Path +) -> None: + model: Final = "no-such-model-" + uuid.uuid4().hex[:8] + with wire_server(_scanner("allow")) as policy, wire_server(_operator_verdict("block")) as operator: + config: Final = _config(tmp_path, _airs(policy.url), _operator_guardrail()) + with owned_proxy(gateway, tmp_path, {"OPERATOR_GUARDRAIL_URL": operator.url}, config=config) as candidate: + response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1])) + + assert response.status_code == 403, response.text + assert "blocked by the operator guardrail" in response.text + assert len(policy.drain()) == 1 + assert [json.loads(request.body) for request in operator.drain()] == [{"model": model}] + + +def test_built_in_guardrail_after_an_allowing_operator_guardrail_is_skipped_for_an_unknown_model( + gateway: Gateway, tmp_path: Path +) -> None: + model: Final = "no-such-model-" + uuid.uuid4().hex[:8] + body: Final = dict(_requests(model)[0][1]) + control: Final = gateway.request("POST", "/v1/chat/completions", body) + with wire_server(_scanner("block")) as policy, wire_server(_operator_verdict("allow")) as operator: + config: Final = _config(tmp_path, _operator_guardrail(), _airs(policy.url)) + with owned_proxy(gateway, tmp_path, {"OPERATOR_GUARDRAIL_URL": operator.url}, config=config) as candidate: + response: Final = candidate.request("POST", "/v1/chat/completions", body) + + assert (response.status_code, response.json()) == (control.status_code, control.json()) + assert [json.loads(request.body) for request in operator.drain()] == [{"model": model}] + assert policy.drain() == () diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 88da07c612b..6e2a67b129d 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -743,7 +743,7 @@ async def test_message_send_reports_an_unresolvable_entra_credential_as_internal mock_proxy_logging = MagicMock() mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data ) mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) downstream = AsyncMock() @@ -1097,7 +1097,7 @@ async def test_subscribe_to_task_calls_pre_call_hook(): mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data ) mock_proxy_logging.async_post_call_streaming_iterator_hook = _passthrough_iterator mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) @@ -1226,7 +1226,7 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): mock_proxy_logging = MagicMock() mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data ) mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) @@ -1293,7 +1293,7 @@ async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400() mock_proxy_logging = MagicMock() mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data ) mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) diff --git a/tests/unit/proxy/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/mcp_server/test_rest_endpoints.py index fdb1ff3cd09..ae3879379bf 100644 --- a/tests/unit/proxy/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_rest_endpoints.py @@ -3072,7 +3072,7 @@ class TestCallToolRestAPI: pre_call_finished_at = {} - async def slow_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def slow_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): await asyncio.sleep(0.05) pre_call_finished_at["value"] = datetime.now() return data @@ -3330,10 +3330,12 @@ class TestCallToolRestAPI: message="Content blocked", model="mcp-tool-call", request_data={}, guardrail_name="block-all" ) - async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def passthrough_pre_call_hook( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): return data - async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): raise guardrail_error async def fake_execute_mcp_tool(**kwargs): @@ -3534,7 +3536,7 @@ class TestCallToolRestAPI: async def fake_add_litellm_data_to_request(**kwargs): return kwargs.get("data", {}) - async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): raise guardrail_error failure_logging = AsyncMock() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 4d1b67f9212..fbcb751277f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,7 +4,8 @@ import copy import datetime import gzip import json -from collections.abc import AsyncIterable, Mapping +from dataclasses import dataclass +from collections.abc import AsyncIterable, Awaitable, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus @@ -31,14 +32,29 @@ from litellm.litellm_core_utils.bug_report import ( bug_report_notice, build_bug_report, ) -from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + RouterGeneralSettings, + RouterRateLimitError, + RouterRateLimitErrorBasic, +) +from litellm.types.utils import ( + CachingDetails, + CallTypesLiteral, + Delta, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, +) from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_KEEPALIVE_PING_BYTES, ) +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -83,7 +99,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_spend_cache_key, team_member_budget_entity_id, ) -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, hash_token from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -439,7 +455,9 @@ class TestProxyBaseLLMRequestProcessing: async def mock_add_litellm_data_to_request(*args, **kwargs): return {} - async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type, skip_guardrails=False): + async def mock_common_processing_pre_call_logic( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): data_copy = copy.deepcopy(data) return data_copy @@ -714,6 +732,7 @@ class TestProxyBaseLLMRequestProcessing: data: dict[str, object], call_type: str, skip_guardrails: bool = False, + routing_probe: object = None, ) -> dict[str, object]: logging_obj: Final = data["litellm_logging_obj"] assert isinstance(logging_obj, LiteLLMLoggingObj) @@ -774,7 +793,7 @@ class TestProxyBaseLLMRequestProcessing: async def mock_add_litellm_data_to_request(*args, **kwargs): return copy.deepcopy(request_body) - async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): data.setdefault("metadata", {}).setdefault("tags", []).extend(guardrail_tags) return data @@ -950,7 +969,9 @@ class TestProxyBaseLLMRequestProcessing: async def retry_add_litellm_data_to_request(*args, **kwargs): return first_pass_data - async def idempotent_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def idempotent_pre_call_hook( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): return data monkeypatch.setattr( @@ -1091,7 +1112,7 @@ class TestProxyBaseLLMRequestProcessing: seen_metadata: dict = {} - async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): seen_metadata.update(data.get("metadata") or {}) return data @@ -1162,7 +1183,9 @@ class TestProxyBaseLLMRequestProcessing: async def mock_add_litellm_data_to_request(*args, **kwargs): return {} - async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_type, skip_guardrails=False): + async def mock_common_processing_pre_call_logic( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): data_copy = copy.deepcopy(data) return data_copy @@ -2208,7 +2231,7 @@ class TestProxyBaseLLMRequestProcessing: data["metadata"] = data.get("metadata", {}) return data - async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): return copy.deepcopy(data) mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) @@ -8071,6 +8094,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: data: dict[str, object], call_type: str, skip_guardrails: bool = False, + routing_probe: object = None, ) -> dict[str, object]: limiter_models.append(str(data["model"])) if data["model"] in untagged_rate_limited_models: @@ -8552,6 +8576,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: data: dict[str, object], call_type: str, skip_guardrails: bool = False, + routing_probe: object = None, ) -> dict[str, object]: limited = await run_limiter(user_api_key_dict=user_api_key_dict, data=data, call_type=call_type) if guardrail not in (limited["metadata"].get("guardrails") or []): @@ -9869,7 +9894,9 @@ class TestPerRequestModelGroupAlias: async def mock_add_litellm_data_to_request(*args, **kwargs): return kwargs.get("data", {}) - async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def passthrough_pre_call_hook( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): return copy.deepcopy(data) mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) @@ -9918,7 +9945,9 @@ class TestPerRequestModelGroupAlias: async def mock_add_litellm_data_to_request(*args, **kwargs): return kwargs.get("data", {}) - async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def passthrough_pre_call_hook( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): return copy.deepcopy(data) mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) @@ -9957,7 +9986,9 @@ class TestPerRequestModelGroupAlias: async def mock_add_litellm_data_to_request(*args, **kwargs): return kwargs.get("data", {}) - async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def passthrough_pre_call_hook( + user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None + ): return copy.deepcopy(data) mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) @@ -10040,7 +10071,7 @@ class TestVideoServedModelPin: async def mock_add_litellm_data_to_request(*args, **kwargs): return kwargs.get("data", {}) - async def pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + async def pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): hooked_data = copy.deepcopy(data) return hooked_data if hook_rewrite is None else {**hooked_data, "model": hook_rewrite} @@ -12022,6 +12053,7 @@ class TestBackgroundResponseRetrievalGovernance: data: dict[str, object], call_type: str, skip_guardrails: bool = False, + routing_probe: object = None, ) -> dict[str, object]: if data.get("response_id") == client_facing_response_id: data["response_id"] = encoded_response_id @@ -12858,3 +12890,1027 @@ async def test_a_chained_proxy_redacts_secrets_and_internal_details_from_relayed assert {key: answered[key] for key in _DOWNSTREAM_GUARDRAIL_FIELDS} == dict(_DOWNSTREAM_GUARDRAIL_FIELDS) assert answered["scanner"]["attempts"] == 2 assert [value for value in _DOWNSTREAM_LEAKABLE_VALUES if value in json.dumps(answered)] == [] +_RAW_EMAIL: Final = "jane@example.com" +_ROUTE_PATHS: Final = MappingProxyType( + { + "acompletion": "/v1/chat/completions", + "atext_completion": "/v1/completions", + "aresponses": "/v1/responses", + "aembedding": "/v1/embeddings", + } +) +_DEFAULT_BODIES: Final = MappingProxyType( + { + "acompletion": MappingProxyType({"messages": ({"role": "user", "content": f"email {_RAW_EMAIL}"},)}), + "atext_completion": MappingProxyType({"prompt": f"email {_RAW_EMAIL}"}), + "aresponses": MappingProxyType({"input": f"email {_RAW_EMAIL}"}), + "aembedding": MappingProxyType({"input": f"email {_RAW_EMAIL}"}), + } +) + + +def _routing_router(fallbacks: list[dict[str, list[str]]] | None = None) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "pong"}, + }, + { + "model_name": "text-embedding-3-small", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1, 0.2], + }, + }, + { + "model_name": "blocked-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "pong"}, + "model_info": {"blocked": True}, + }, + ], + fallbacks=fallbacks or [], + ) + + +def _mask_emails(data: dict) -> None: + for message in data.get("messages") or []: + if isinstance(message.get("content"), str): + message["content"] = message["content"].replace(_RAW_EMAIL, "[EMAIL]") + + +class _VendorGuardrail(CustomGuardrail): + """Stands in for a vendor guardrail litellm ships: counts its calls and can mask, block, fail or reroute.""" + + __module__ = "litellm.proxy.guardrails.guardrail_hooks.vendor_double" + + def __init__( + self, + behavior: Literal["pass", "mask", "block", "unreachable", "reroute"] = "pass", + name: str = "vendor-guardrail", + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + run_in_parallel: bool = False, + order: list[str] | None = None, + ) -> None: + super().__init__(guardrail_name=name, event_hook=event_hook, default_on=True) + self.behavior: Final = behavior + self.run_in_parallel = run_in_parallel + self.order: Final = order if order is not None else [] + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + self.order.append(self.guardrail_name or "") + if self.behavior == "block": + raise HTTPException(status_code=400, detail="Prompt blocked by vendor policy") + if self.behavior == "unreachable": + raise HTTPException(status_code=503, detail="guardrail vendor unreachable") + if self.behavior == "reroute": + data["model"] = "gpt-4o-mini" + if self.behavior == "mask": + _mask_emails(data) + return data + + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + self.calls += 1 + + +class _CustomerRoutingGuardrail(CustomGuardrail): + """A guardrail from the customer's own code that sends a virtual model to a loaded one.""" + + def __init__(self, run_in_parallel: bool = False) -> None: + super().__init__(guardrail_name="customer-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.run_in_parallel = run_in_parallel + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + if data.get("model") == "virtual-model": + data["model"] = "gpt-4o-mini" + return data + + +class _InheritedCustomerGuardrail(_CustomerRoutingGuardrail): + """A customer guardrail that inherits its pre-call hook from the customer's own base class.""" + + +class _CustomerTextGuardrail(CustomGuardrail): + """A customer guardrail that only implements ``apply_guardrail``.""" + + def __init__(self, run_in_parallel: bool = False) -> None: + super().__init__( + guardrail_name="customer-text-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + self.run_in_parallel = run_in_parallel + self.calls = 0 + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.calls += 1 + return inputs + + +class _BuiltInContentCheck(CustomLogger): + """Stands in for a content check litellm ships that is not a guardrail: records the messages it was sent.""" + + __module__ = "litellm.proxy.hooks.content_check_double" + + def __init__(self) -> None: + super().__init__() + self.seen: Final[list[str]] = [] + + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: + self.seen.append(json.dumps(data.get("messages"))) + return data + + +class _ModelMappingHook(CustomLogger): + """A customer callback that renames a virtual model, editing the request or returning a new one.""" + + def __init__( + self, + source: str = "virtual-model", + target: str = "gpt-4o-mini", + returns_new_request: bool = False, + order: list[str] | None = None, + ) -> None: + super().__init__() + self.source: Final = source + self.target: Final = target + self.returns_new_request: Final = returns_new_request + self.order: Final = order if order is not None else [] + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.order.append("mapping-hook") + if data.get("model") != self.source: + return data + if self.returns_new_request: + return {**data, "model": self.target} + data["model"] = self.target + return data + + +class _EmailMaskingHook(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + _mask_emails(data) + return data + + +class _RejectingHook(CustomLogger): + def __init__(self, model: str | None = None) -> None: + super().__init__() + self.model: Final = model + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if self.model in (None, data.get("model")): + raise HTTPException(status_code=403, detail="rejected by hook") + return data + + +class _RateLimitingHook(CustomLogger): + def __init__(self, model: str) -> None: + super().__init__() + self.model: Final = model + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if data.get("model") == self.model: + raise ProxyRateLimitError(detail=f"rpm limit reached for {self.model}") + return data + + +class _ContentEnforcer(CustomLogger): + enforces_request_content = True + + def __init__(self, order: list[str] | None = None) -> None: + super().__init__() + self.order: Final = order if order is not None else [] + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + self.order.append("content-enforcer") + return data + + +class _ModerationHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + self.calls += 1 + + +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logs: Final[list[dict[str, object]]] = [] + self.dumps: Final[list[str]] = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + litellm_params: Final = kwargs.get("litellm_params") or {} + logged_request: Final = (litellm_params.get("proxy_server_request") or {}).get("body") or {} + standard: Final = kwargs.get("standard_logging_object") or {} + self.dumps.append(json.dumps(kwargs, default=str)) + self.logs.append( + { + "model": kwargs.get("model"), + "messages": kwargs.get("messages"), + "standard_model": standard.get("model"), + "standard_model_group": standard.get("model_group"), + "standard_messages": standard.get("messages"), + "request_model": logged_request.get("model"), + "request_messages": logged_request.get("messages"), + "request_has_call_id": "litellm_call_id" in logged_request, + } + ) + + async def async_post_call_failure_hook( + self, request_data, original_exception, user_api_key_dict, traceback_str=None + ): + self.dumps.append(json.dumps(request_data, default=str)) + + +@dataclass(frozen=True) +class _Sent: + status: int + message: str + failure_logs: tuple[Mapping[str, object], ...] + failure_dumps: tuple[str, ...] + phase_events: tuple[str, ...] + registry_lookups: tuple[str, ...] + + +async def _send( + monkeypatch, + callbacks: Sequence[CustomLogger], + model: str | None = "does-not-exist", + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"] = "acompletion", + body: Mapping[str, object] | None = None, + llm_router: Router | None = None, + use_router: bool = True, + user_api_key_dict: ProxyUserAPIKeyAuth | None = None, + proxy_logging_obj: ProxyLogging | None = None, + registry_has: Callable[[str], Awaitable[bool]] | None = None, +) -> _Sent: + """Send one request the way the proxy's endpoint does and report what the client and the failure log saw.""" + from litellm.proxy import common_request_processing, proxy_server + from litellm.proxy.common_utils import registry_read_through + + recorder: Final = _FailureLogRecorder() + monkeypatch.setattr(litellm, "callbacks", [*callbacks, recorder]) + ProxyLogging._callback_capabilities_cache.clear() + phase_events: Final = list[str]() + monkeypatch.setattr( + common_request_processing, "phase_event", lambda name, attributes=None: phase_events.append(name) + ) + lookups: Final = list[str]() + + async def no_db_row(model_name: str) -> bool: + lookups.append(model_name) + return False + + async def recorded_lookup(model_name: str) -> bool: + lookups.append(model_name) + return await registry_has(model_name) if registry_has is not None else False + + monkeypatch.setattr( + registry_read_through, + "model_registry_read_through", + registry_read_through.RegistryReadThrough( + resync=recorded_lookup if registry_has is not None else no_db_row, + is_loaded=registry_read_through._model_is_loaded, + ), + ) + router: Final = (llm_router or _routing_router()) if use_router else None + monkeypatch.setattr(proxy_server, "llm_router", router) + path: Final = _ROUTE_PATHS[route_type] + data: Final = { + **({"model": model} if model is not None else {}), + **copy.deepcopy(dict(_DEFAULT_BODIES[route_type])), + **copy.deepcopy(dict(body or {})), + } + data.update({key: list(value) for key, value in data.items() if isinstance(value, tuple)}) + processor: Final = ProxyBaseLLMRequestProcessing(data=data) + auth: Final = user_api_key_dict or ProxyUserAPIKeyAuth(request_route=path) + logging_obj: Final = proxy_logging_obj or ProxyLogging(user_api_key_cache=UserApiKeyCache()) + try: + await processor.base_process_llm_request( + request=Request({"type": "http", "method": "POST", "path": path, "headers": [], "query_string": b""}), + fastapi_response=Response(), + user_api_key_dict=auth, + route_type=route_type, + proxy_logging_obj=logging_obj, + general_settings={}, + proxy_config=proxy_server.ProxyConfig(), + select_data_generator=proxy_server.select_data_generator, + llm_router=router, + version=None, + ) + status_code, message = 200, "" + except Exception as request_error: + try: + await processor.handle_llm_api_exception( + e=request_error, user_api_key_dict=auth, proxy_logging_obj=logging_obj + ) + except ProxyException as client_error: + status_code, message = int(client_error.code or 0), client_error.message + return _Sent( + status=status_code, + message=message, + failure_logs=tuple(recorder.logs), + failure_dumps=tuple(recorder.dumps), + phase_events=tuple(phase_events), + registry_lookups=tuple(lookups), + ) + + +def _attach_policy_pipeline(monkeypatch, mode: Literal["pre_call", "post_call"], guardrail_name: str) -> None: + from litellm.proxy.policy_engine import attachment_registry, policy_registry + + policies: Final = policy_registry.PolicyRegistry() + policies.load_policies( + { + "safety": { + "guardrails": {"add": [guardrail_name]}, + "pipeline": { + "mode": mode, + "steps": [{"guardrail": guardrail_name, "on_fail": "block", "on_pass": "allow"}], + }, + } + } + ) + attachments: Final = attachment_registry.AttachmentRegistry() + attachments.load_attachments([{"policy": "safety", "scope": "*"}]) + monkeypatch.setattr(policy_registry, "_policy_registry", policies) + monkeypatch.setattr(attachment_registry, "_attachment_registry", attachments) + + +def _shows_no_raw_content(sent: _Sent) -> bool: + return bool(sent.failure_dumps) and not any(_RAW_EMAIL in dump for dump in sent.failure_dumps) + + +class TestUnroutableModelSkipsPreCallGuardrails: + """LIT-8011: a request whose model the proxy can't route gets the routing error it gets with no guardrail + configured, and no pre-call guardrail vendor is called for it.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("behavior", ["pass", "mask", "block", "unreachable"]) + async def test_unknown_model_gets_the_invalid_model_error_without_calling_the_guardrail( + self, monkeypatch, behavior + ): + guardrail: Final = _VendorGuardrail(behavior=behavior) + + guarded: Final = await _send(monkeypatch, [guardrail]) + unguarded: Final = await _send(monkeypatch, []) + + assert (guarded.status, guarded.message) == (unguarded.status, unguarded.message) + assert guarded.status == 400 + assert "Invalid model name passed in model=does-not-exist" in guarded.message + assert guardrail.calls == 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model", + [pytest.param(None, id="missing"), pytest.param("", id="empty"), pytest.param("m" * 5000, id="5000-chars")], + ) + async def test_missing_or_malformed_model_gets_the_routing_error_without_calling_the_guardrail( + self, monkeypatch, model + ): + guardrail: Final = _VendorGuardrail(behavior="block") + + guarded: Final = await _send(monkeypatch, [guardrail], model=model) + unguarded: Final = await _send(monkeypatch, [], model=model) + + assert (guarded.status, guarded.message) == (unguarded.status, unguarded.message) + assert guarded.status == 400 + assert guardrail.calls == 0 + + @pytest.mark.asyncio + async def test_unknown_model_skips_a_parallel_guardrail(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="block", run_in_parallel=True) + + sent: Final = await _send(monkeypatch, [guardrail]) + + assert sent.status == 400 + assert "Invalid model name passed in" in sent.message + assert guardrail.calls == 0 + + @pytest.mark.asyncio + async def test_endpoint_without_a_router_skips_the_guardrail_for_an_unknown_model(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="block") + + guarded: Final = await _send(monkeypatch, [guardrail], use_router=False) + unguarded: Final = await _send(monkeypatch, [], use_router=False) + + assert (guarded.status, guarded.message) == (unguarded.status, unguarded.message) + assert "Invalid model name passed in model=does-not-exist" in guarded.message + assert guardrail.calls == 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("route_type", "body"), + [ + pytest.param("acompletion", {"stream": True}, id="streaming-chat"), + pytest.param("aresponses", {}, id="responses"), + pytest.param("aembedding", {}, id="embeddings"), + pytest.param("atext_completion", {}, id="completions"), + ], + ) + async def test_unknown_model_on_other_routes_skips_the_guardrail(self, monkeypatch, route_type, body): + guardrail: Final = _VendorGuardrail(behavior="block") + + guarded: Final = await _send(monkeypatch, [guardrail], route_type=route_type, body=body) + unguarded: Final = await _send(monkeypatch, [], route_type=route_type, body=body) + + assert (guarded.status, guarded.message) == (unguarded.status, unguarded.message) + assert guarded.status == 400 + assert guardrail.calls == 0 + + @pytest.mark.asyncio + async def test_model_a_hook_swaps_for_another_unknown_model_is_rejected_under_the_new_name(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="block") + swap: Final = _ModelMappingHook(target="still-unknown", returns_new_request=True) + + sent: Final = await _send(monkeypatch, [swap, guardrail], model="virtual-model") + + assert sent.status == 400 + assert "Invalid model name passed in model=still-unknown" in sent.message + assert guardrail.calls == 0 + + @pytest.mark.asyncio + async def test_guardrail_that_would_reroute_an_unknown_model_is_not_called(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="reroute") + + sent: Final = await _send(monkeypatch, [guardrail], model="guard-reroute") + + assert sent.status == 400 + assert "Invalid model name passed in model=guard-reroute" in sent.message + assert guardrail.calls == 0 + + @pytest.mark.asyncio + async def test_rate_limit_fallback_to_an_unknown_model_skips_the_guardrail(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="mask") + + sent: Final = await _send( + monkeypatch, + [_RateLimitingHook(model="gpt-4o-mini"), guardrail], + model="gpt-4o-mini", + llm_router=_routing_router(fallbacks=[{"gpt-4o-mini": ["does-not-exist"]}]), + ) + + assert sent.status == 400 + assert "Invalid model name passed in model=does-not-exist" in sent.message + assert guardrail.calls == 0 + assert _shows_no_raw_content(sent), sent.failure_dumps + + @pytest.mark.asyncio + async def test_key_with_router_settings_gets_the_invalid_model_error_without_calling_the_guardrail( + self, monkeypatch + ): + guardrail: Final = _VendorGuardrail(behavior="block") + key: Final = ProxyUserAPIKeyAuth(request_route="/v1/chat/completions", router_settings={"num_retries": 0}) + + guarded: Final = await _send(monkeypatch, [guardrail], user_api_key_dict=key) + unguarded: Final = await _send(monkeypatch, [], user_api_key_dict=key) + + assert (guarded.status, guarded.message) == (unguarded.status, unguarded.message) + assert guarded.status == 400 + assert guardrail.calls == 0 + + +class TestUnroutableModelFailureLogHidesUnscannedContent: + """A request turned away before a guardrail that applied to it ran must not log, or hand to the hooks after it, + the content that guardrail would have scanned, since a masking guardrail never got to mask it.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "body", + [ + pytest.param({}, id="chat-messages"), + pytest.param({"system": [{"type": "text", "text": f"email {_RAW_EMAIL}"}]}, id="anthropic-system"), + pytest.param({"instructions": f"email {_RAW_EMAIL}"}, id="responses-instructions"), + pytest.param( + {"contents": [{"role": "user", "parts": [{"text": f"email {_RAW_EMAIL}"}]}]}, id="gemini-contents" + ), + pytest.param( + {"systemInstruction": {"parts": [{"text": f"email {_RAW_EMAIL}"}]}}, id="gemini-system-instruction" + ), + pytest.param( + {"config": {"system_instruction": {"parts": [{"text": f"email {_RAW_EMAIL}"}]}, "temperature": 0}}, + id="gemini-config-system-instruction", + ), + pytest.param({"prompt": f"email {_RAW_EMAIL}", "input": [f"email {_RAW_EMAIL}"]}, id="prompt-and-input"), + pytest.param( + {"tools": [{"type": "function", "function": {"name": "mail", "description": _RAW_EMAIL}}]}, + id="tool-definitions", + ), + pytest.param( + {"query": f"email {_RAW_EMAIL}", "instruction": f"email {_RAW_EMAIL}"}, + id="rerank-query-and-instruction", + ), + ], + ) + async def test_failure_log_hides_every_field_a_skipped_guardrail_would_have_scanned(self, monkeypatch, body): + sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask")], body=body) + + assert sent.status == 400 + assert _shows_no_raw_content(sent), sent.failure_dumps + + @pytest.mark.asyncio + async def test_failure_log_hides_content_a_skipped_parallel_guardrail_would_have_scanned( + self, monkeypatch: pytest.MonkeyPatch + ): + sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask", run_in_parallel=True)]) + + assert sent.status == 400 + assert _shows_no_raw_content(sent), sent.failure_dumps + + @pytest.mark.asyncio + async def test_failure_log_hides_content_when_a_built_in_hook_rejects_after_a_skipped_guardrail( + self, monkeypatch + ): + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + key: Final = ProxyUserAPIKeyAuth( + api_key=hash_token("sk-rpm-one"), + token=hash_token("sk-rpm-one"), + rpm_limit=1, + request_route="/v1/chat/completions", + ) + guardrail: Final = _VendorGuardrail(behavior="mask") + callbacks: Final = [guardrail, proxy_logging_obj.max_parallel_request_limiter] + + first: Final = await _send(monkeypatch, callbacks, user_api_key_dict=key, proxy_logging_obj=proxy_logging_obj) + limited: Final = await _send(monkeypatch, callbacks, user_api_key_dict=key, proxy_logging_obj=proxy_logging_obj) + + assert (first.status, limited.status, guardrail.calls) == (400, 429, 0) + assert _shows_no_raw_content(limited), limited.failure_dumps + + @pytest.mark.asyncio + async def test_built_in_hook_after_a_skipped_guardrail_is_not_sent_the_content_it_would_have_masked( + self, monkeypatch: pytest.MonkeyPatch + ): + check: Final = _BuiltInContentCheck() + + sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask"), check]) + + assert sent.status == 400 + assert check.seen and not any(_RAW_EMAIL in seen for seen in check.seen), check.seen + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("stream_scope", "stream", "content_hidden"), + [ + pytest.param("streaming", True, True, id="streaming-guardrail-streaming-request"), + pytest.param("streaming", False, False, id="streaming-guardrail-plain-request"), + pytest.param("non_streaming", True, False, id="plain-guardrail-streaming-request"), + ], + ) + async def test_failure_log_hides_content_only_when_the_skipped_guardrail_applied( + self, monkeypatch, stream_scope, stream, content_hidden + ): + guardrail: Final = _VendorGuardrail(behavior="mask") + guardrail.apply_stream_scope(stream_scope) + + sent: Final = await _send(monkeypatch, [guardrail], body={"stream": stream}) + + assert sent.status == 400 + assert _shows_no_raw_content(sent) is content_hidden, sent.failure_dumps + + +class TestUnroutableModelLeavesEverythingElseAsBefore: + """Everything except the skipped pre-call guardrails behaves as it did before LIT-8011, including the cases + earlier attempts at this change broke.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("route_type", "body", "virtual", "target"), + [ + pytest.param("acompletion", {}, "virtual-model", "gpt-4o-mini", id="chat"), + pytest.param("acompletion", {"stream": True}, "virtual-model", "gpt-4o-mini", id="streaming-chat"), + pytest.param("aresponses", {}, "virtual-model", "gpt-4o-mini", id="responses"), + pytest.param("aembedding", {}, "virtual-embedding", "text-embedding-3-small", id="embeddings"), + ], + ) + @pytest.mark.parametrize("check", ["guardrail", "parallel-guardrail", "content-enforcer"]) + async def test_hook_listed_after_a_pre_call_check_can_still_rename_a_virtual_model( + self, monkeypatch, route_type, body, virtual, target, check + ): + guardrail: Final = _VendorGuardrail(run_in_parallel=check == "parallel-guardrail") + enforcer: Final = _ContentEnforcer() + mapping: Final = _ModelMappingHook(source=virtual, target=target) + + sent: Final = await _send( + monkeypatch, + [enforcer if check == "content-enforcer" else guardrail, mapping], + model=virtual, + route_type=route_type, + body=body, + ) + + assert (sent.status, sent.message) == (200, "") + assert (enforcer.calls if check == "content-enforcer" else guardrail.calls) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("customer_guardrail", [_CustomerRoutingGuardrail, _InheritedCustomerGuardrail]) + async def test_customer_guardrail_can_still_rename_a_virtual_model( + self, monkeypatch: pytest.MonkeyPatch, customer_guardrail: type[_CustomerRoutingGuardrail] + ): + guardrail: Final = _VendorGuardrail() + customer: Final = customer_guardrail() + + sent: Final = await _send(monkeypatch, [guardrail, customer], model="virtual-model") + + assert (sent.status, sent.message) == (200, "") + assert (guardrail.calls, customer.calls) == (1, 1) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "customer_guardrail", [_CustomerRoutingGuardrail, _InheritedCustomerGuardrail, _CustomerTextGuardrail] + ) + @pytest.mark.parametrize("run_in_parallel", [False, True]) + async def test_customer_guardrail_still_runs_for_an_unknown_model( + self, + monkeypatch: pytest.MonkeyPatch, + customer_guardrail: type[_CustomerRoutingGuardrail] | type[_CustomerTextGuardrail], + run_in_parallel: bool, + ): + customer: Final = customer_guardrail(run_in_parallel=run_in_parallel) + + sent: Final = await _send(monkeypatch, [customer]) + + assert sent.status == 400 + assert "Invalid model name passed in model=does-not-exist" in sent.message + assert customer.calls == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("mapping_first", [True, False]) + async def test_pre_call_pipeline_does_not_stop_a_hook_from_renaming_a_virtual_model( + self, monkeypatch, mapping_first + ): + guardrail: Final = _VendorGuardrail(name="pipeline-guardrail") + _attach_policy_pipeline(monkeypatch, "pre_call", "pipeline-guardrail") + mapping: Final = _ModelMappingHook() + + sent: Final = await _send( + monkeypatch, [mapping, guardrail] if mapping_first else [guardrail, mapping], model="virtual-model" + ) + + assert (sent.status, sent.message) == (200, "") + assert guardrail.calls == 1 + + @pytest.mark.asyncio + async def test_rate_limiter_still_counts_requests_for_unknown_models(self, monkeypatch): + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + key: Final = ProxyUserAPIKeyAuth( + api_key=hash_token("sk-rpm-one"), + token=hash_token("sk-rpm-one"), + rpm_limit=1, + request_route="/v1/chat/completions", + ) + callbacks: Final = [_VendorGuardrail(), proxy_logging_obj.max_parallel_request_limiter] + + statuses: Final = [ + ( + await _send(monkeypatch, callbacks, user_api_key_dict=key, proxy_logging_obj=proxy_logging_obj) + ).status + for _ in range(3) + ] + + assert statuses == [400, 429, 429] + + @pytest.mark.asyncio + async def test_hook_listed_after_a_guardrail_still_runs_for_an_unknown_model(self, monkeypatch): + sent: Final = await _send(monkeypatch, [_VendorGuardrail(), _RejectingHook()]) + + assert (sent.status, sent.message) == (403, "rejected by hook") + + @pytest.mark.asyncio + async def test_during_call_hooks_still_run_for_an_unknown_model(self, monkeypatch): + moderation: Final = _ModerationHook() + during_call_guardrail: Final = _VendorGuardrail(name="during-call", event_hook=GuardrailEventHooks.during_call) + + sent: Final = await _send(monkeypatch, [moderation, during_call_guardrail]) + + assert sent.status == 400 + assert (moderation.calls, during_call_guardrail.calls) == (1, 1) + + @pytest.mark.asyncio + @pytest.mark.parametrize("with_guardrail", [False, True]) + async def test_pre_call_completed_event_is_still_emitted_for_an_unknown_model(self, monkeypatch, with_guardrail): + sent: Final = await _send(monkeypatch, [_VendorGuardrail()] if with_guardrail else []) + + assert sent.status == 400 + assert "litellm.request.pre_call_completed" in sent.phase_events + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("returns_new_request", "expected"), + [ + pytest.param( + False, + { + "model": "gpt-4o-mini", + "standard_model": "openai/gpt-4o-mini", + "standard_model_group": "gpt-4o-mini", + "request_model": "virtual-model", + }, + id="hook-edits-the-request", + ), + pytest.param( + True, + { + "model": "virtual-model", + "standard_model": "virtual-model", + "standard_model_group": "", + "request_model": "virtual-model", + }, + id="hook-returns-a-new-request", + ), + ], + ) + async def test_guardrail_block_after_a_rename_logs_the_same_models_as_before( + self, monkeypatch, returns_new_request, expected + ): + mapping: Final = _ModelMappingHook(returns_new_request=returns_new_request) + + sent: Final = await _send(monkeypatch, [mapping, _VendorGuardrail(behavior="block")], model="virtual-model") + + assert (sent.status, sent.message) == (400, "Prompt blocked by vendor policy") + assert [{field: log[field] for field in expected} for log in sent.failure_logs] == [expected] + assert [log["request_has_call_id"] for log in sent.failure_logs] == [False] + + @pytest.mark.asyncio + async def test_hook_rejection_after_the_guardrail_passed_logs_the_request_as_sent(self, monkeypatch): + sent: Final = await _send( + monkeypatch, [_ModelMappingHook(), _VendorGuardrail(), _RejectingHook()], model="virtual-model" + ) + + assert (sent.status, sent.message) == (403, "rejected by hook") + assert [ + (log["request_model"], log["request_messages"], log["request_has_call_id"]) for log in sent.failure_logs + ] == [("virtual-model", [{"role": "user", "content": f"email {_RAW_EMAIL}"}], False)] + + @pytest.mark.asyncio + async def test_unknown_model_without_guardrails_logs_the_request_its_hooks_produced(self, monkeypatch): + sent: Final = await _send(monkeypatch, [_EmailMaskingHook()]) + + assert sent.status == 400 + assert [ + (log["standard_messages"], log["request_messages"], log["request_has_call_id"]) + for log in sent.failure_logs + ] == [([{"role": "user", "content": "email [EMAIL]"}], [{"role": "user", "content": "email [EMAIL]"}], True)] + + @pytest.mark.asyncio + async def test_post_call_only_pipeline_leaves_the_failure_log_of_an_unknown_model_unredacted(self, monkeypatch): + guardrail: Final = _VendorGuardrail(name="post-call-guardrail", event_hook=GuardrailEventHooks.post_call) + _attach_policy_pipeline(monkeypatch, "post_call", "post-call-guardrail") + + sent: Final = await _send(monkeypatch, [guardrail]) + + assert sent.status == 400 + assert [log["request_messages"] for log in sent.failure_logs] == [ + [{"role": "user", "content": f"email {_RAW_EMAIL}"}] + ] + + @pytest.mark.asyncio + async def test_rejection_by_a_hook_listed_before_the_guardrail_is_logged_as_before(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="mask") + + sent: Final = await _send(monkeypatch, [_RejectingHook(), guardrail]) + + assert (sent.status, sent.message) == (403, "rejected by hook") + assert guardrail.calls == 0 + assert [log["request_messages"] for log in sent.failure_logs] == [ + [{"role": "user", "content": f"email {_RAW_EMAIL}"}] + ] + + @pytest.mark.asyncio + async def test_rate_limited_unknown_model_whose_fallback_a_hook_rejects_is_logged_as_before(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="mask") + + sent: Final = await _send( + monkeypatch, + [_RateLimitingHook(model="does-not-exist"), _RejectingHook(model="gpt-4o-mini"), guardrail], + llm_router=_routing_router(fallbacks=[{"does-not-exist": ["gpt-4o-mini"]}]), + ) + + assert (sent.status, sent.message) == (403, "rejected by hook") + assert guardrail.calls == 0 + assert [(log["request_model"], log["request_messages"]) for log in sent.failure_logs] == [ + ("does-not-exist", [{"role": "user", "content": f"email {_RAW_EMAIL}"}]) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "callbacks", "expected_lookups"), + [ + pytest.param("does-not-exist", lambda: [_VendorGuardrail()], ("does-not-exist",), id="unknown-model"), + pytest.param( + "unknown-virtual", + lambda: [_ModelMappingHook(source="unknown-virtual", target="db-only-model"), _VendorGuardrail()], + ("db-only-model",), + id="renamed-to-a-model-only-in-the-db", + ), + pytest.param( + "virtual-model", + lambda: [_VendorGuardrail(), _ModelMappingHook()], + (), + id="renamed-to-a-loaded-model", + ), + ], + ) + async def test_model_registry_is_read_once_for_the_name_routing_ends_up_with( + self, monkeypatch, model, callbacks, expected_lookups + ): + sent: Final = await _send(monkeypatch, callbacks(), model=model) + + assert sent.registry_lookups == expected_lookups + + @pytest.mark.asyncio + async def test_model_created_on_another_replica_is_served_with_its_guardrail_on_the_first_request( + self, monkeypatch + ): + router: Final = _routing_router() + guardrail: Final = _VendorGuardrail() + + async def load_from_db(model_name: str) -> bool: + router.add_deployment( + Deployment( + model_name=model_name, + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", api_key="sk-test", mock_response="pong" + ), + ) + ) + return True + + sent: Final = await _send( + monkeypatch, [guardrail], model="created-on-a-sibling", llm_router=router, registry_has=load_from_db + ) + + assert (sent.status, sent.message, guardrail.calls) == (200, "", 1) + assert sent.registry_lookups == ("created-on-a-sibling",) + + @pytest.mark.asyncio + async def test_routable_model_runs_every_hook_and_guardrail_once_in_configured_order(self, monkeypatch): + order: Final = list[str]() + + sent: Final = await _send( + monkeypatch, + [_ModelMappingHook(order=order), _VendorGuardrail(order=order), _ContentEnforcer(order=order)], + model="gpt-4o-mini", + ) + + assert (sent.status, order) == (200, ["mapping-hook", "vendor-guardrail", "content-enforcer"]) + + @pytest.mark.asyncio + async def test_guardrail_still_blocks_a_routable_model(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="block") + + sent: Final = await _send(monkeypatch, [guardrail, _RejectingHook()], model="gpt-4o-mini") + + assert (sent.status, sent.message, guardrail.calls) == (400, "Prompt blocked by vendor policy", 1) + + @pytest.mark.asyncio + async def test_hook_that_renames_a_blocked_model_to_a_routable_one_is_served(self, monkeypatch): + guardrail: Final = _VendorGuardrail() + + sent: Final = await _send( + monkeypatch, [_ModelMappingHook(source="blocked-model"), guardrail], model="blocked-model" + ) + + assert (sent.status, sent.message, guardrail.calls) == (200, "", 1) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("model", "body", "llm_router"), + [ + pytest.param( + "openai/anything", + {}, + lambda: Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-test", "mock_response": "pong"}, + } + ] + ), + id="wildcard-route", + ), + pytest.param("does-not-exist", {"api_key": "sk-caller"}, _routing_router, id="caller-api-key"), + pytest.param( + "does-not-exist", + {}, + lambda: Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "mock_response": "pong", + }, + } + ], + router_general_settings=RouterGeneralSettings(pass_through_all_models=True), + ), + id="pass-through-all-models", + ), + ], + ) + async def test_request_the_router_resolves_itself_still_runs_its_guardrail( + self, monkeypatch, model, body, llm_router + ): + guardrail: Final = _VendorGuardrail() + + await _send(monkeypatch, [guardrail], model=model, body=body, llm_router=llm_router()) + + assert guardrail.calls == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("kind", ["content-enforcer", "pipeline-guardrail"]) + async def test_unknown_model_still_runs_pipelines_and_content_enforcers(self, monkeypatch, kind): + guardrail: Final = _VendorGuardrail(name="pipeline-guardrail") + enforcer: Final = _ContentEnforcer() + if kind == "pipeline-guardrail": + _attach_policy_pipeline(monkeypatch, "pre_call", "pipeline-guardrail") + + sent: Final = await _send(monkeypatch, [enforcer] if kind == "content-enforcer" else [guardrail]) + + assert sent.status == 400 + assert (enforcer.calls if kind == "content-enforcer" else guardrail.calls) == 1 + + @pytest.mark.asyncio + async def test_guardrails_still_run_for_an_unknown_model_when_one_reroutes_sensitive_data(self, monkeypatch): + guardrail: Final = _VendorGuardrail() + router_guardrail: Final = CustomGuardrail( + guardrail_name="sensitive-router", + event_hook=GuardrailEventHooks.pre_call, + on_sensitive_data="route", + sensitive_data_route_to_model="gpt-4o-mini", + ) + + sent: Final = await _send(monkeypatch, [guardrail, router_guardrail]) + + assert (sent.status, guardrail.calls) == (400, 1) + + @pytest.mark.asyncio + async def test_built_in_hook_that_routes_an_unknown_model_after_a_skipped_guardrail_gets_the_routing_error( + self, monkeypatch + ): + from litellm.proxy.hooks.sensitive_data_routing import PROXY_SensitiveDataRoutingHandler + from litellm.proxy.utils import InternalUsageCache + + guardrail: Final = _VendorGuardrail() + session_pins: Final = PROXY_SensitiveDataRoutingHandler(internal_usage_cache=InternalUsageCache(DualCache())) + key: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-pinned"), request_route="/v1/chat/completions") + await session_pins.set_session_routing("pinned-session", "gpt-4o-mini", key) + + sent: Final = await _send( + monkeypatch, + [guardrail, session_pins], + body={"litellm_session_id": "pinned-session"}, + user_api_key_dict=key, + ) + + assert (sent.status, guardrail.calls) == (400, 0) + assert "Invalid model name passed in model=does-not-exist" in sent.message + + @pytest.mark.asyncio + async def test_caller_that_skips_guardrails_still_runs_content_enforcers_for_a_renamed_model(self, monkeypatch): + from litellm.proxy import proxy_server + + enforcer: Final = _ContentEnforcer() + guardrail: Final = _VendorGuardrail() + monkeypatch.setattr(litellm, "callbacks", [_ModelMappingHook(), enforcer, guardrail]) + ProxyLogging._callback_capabilities_cache.clear() + router: Final = _routing_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": "virtual-model", "messages": [{"role": "user", "content": "hi"}]} + ) + + data, _ = await processor.common_processing_pre_call_logic( + request=Request( + {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [], "query_string": b""} + ), + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(request_route="/v1/responses"), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=proxy_server.ProxyConfig(), + route_type="aresponses", + llm_router=router, + skip_guardrails=True, + ) + + assert (data["model"], enforcer.calls, guardrail.calls) == ("gpt-4o-mini", 1, 0) diff --git a/tests/unit/proxy/test_model_level_guardrails.py b/tests/unit/proxy/test_model_level_guardrails.py index dc076e7b0fc..90dbba6bce3 100644 --- a/tests/unit/proxy/test_model_level_guardrails.py +++ b/tests/unit/proxy/test_model_level_guardrails.py @@ -601,7 +601,7 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook(): captured_pre_call_guardrails: list = [] - async def fake_pre_call_hook(*, user_api_key_dict, data, call_type, skip_guardrails=False): + async def fake_pre_call_hook(*, user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None): # Snapshot the list rather than the dict: metadata is shared by # reference, so a merge that happens after this point would otherwise # show up here retroactively and the assertion would pass either way. From a30aeb8a71eab737b7ccc7a7bd002dcca309f400 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 17:35:29 -0700 Subject: [PATCH 2/8] fix(proxy): redact the logged messages of a request refused for its model --- litellm/proxy/common_request_processing.py | 8 +++++-- .../proxy/test_common_request_processing.py | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 69284fa64e8..4e9605cc4a3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,6 +46,7 @@ from litellm.constants import ( MAX_LITELLM_CALL_ID_LENGTH, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, NON_INFERENCE_CALL_TYPES, + REDACTED_BY_LITELLM, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, STREAM_SSE_KEEPALIVE_PING_BYTES, @@ -2704,11 +2705,14 @@ class ProxyBaseLLMRequestProcessing: return self.data, logging_obj def _redact_request_its_skipped_guardrails_never_scanned(self) -> None: - """Redact the request body and its logged snapshot, so failure callbacks never see content a skipped - masking guardrail would have rewritten. + """Redact the request body, its logged snapshot and the logged messages, so failure callbacks never see + content a skipped masking guardrail would have rewritten. """ self.data.update(without_unscanned_content(self.data)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict refresh_proxy_server_request_body_snapshot(self.data) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict + logging_obj: Final = self.data.get("litellm_logging_obj") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # self.data is the untyped request dict + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.update_messages([{"role": "user", "content": REDACTED_BY_LITELLM}]) async def _pre_call_with_fallbacks( self, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fbcb751277f..6b94319ae06 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -13425,6 +13425,29 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: assert sent.status == 400 assert _shows_no_raw_content(sent), sent.failure_dumps + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("route_type", "body"), + [ + pytest.param( + "aresponses", {"input": [{"role": "user", "content": f"email {_RAW_EMAIL}"}]}, id="responses-list" + ), + pytest.param( + "aresponses", + {"input": [{"role": "user", "content": [{"type": "input_text", "text": f"email {_RAW_EMAIL}"}]}]}, + id="responses-input-text-parts", + ), + pytest.param("aresponses", {}, id="responses-string"), + pytest.param("aembedding", {"input": [f"email {_RAW_EMAIL}"]}, id="embeddings-list"), + pytest.param("atext_completion", {}, id="completions"), + ], + ) + async def test_failure_log_hides_the_input_of_other_routes(self, monkeypatch, route_type, body): + sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask")], route_type=route_type, body=body) + + assert sent.status == 400 + assert _shows_no_raw_content(sent), sent.failure_dumps + @pytest.mark.asyncio async def test_failure_log_hides_content_a_skipped_parallel_guardrail_would_have_scanned( self, monkeypatch: pytest.MonkeyPatch From cc810eb9d7045826bcd4c279ea03ab260694eef1 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 13:52:34 -0700 Subject: [PATCH 3/8] test(proxy): pin extraction state and fields out of a refused request's failure log --- tests/unit/proxy/test_common_request_processing.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 6b94319ae06..8054b3aad05 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -13417,6 +13417,15 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: {"query": f"email {_RAW_EMAIL}", "instruction": f"email {_RAW_EMAIL}"}, id="rerank-query-and-instruction", ), + pytest.param( + { + "state": f"email {_RAW_EMAIL}", + "fields": { + "contact": {"type": "string", "instructions": f"email {_RAW_EMAIL}", "unit": _RAW_EMAIL} + }, + }, + id="extraction-state-and-fields", + ), ], ) async def test_failure_log_hides_every_field_a_skipped_guardrail_would_have_scanned(self, monkeypatch, body): From f61550d2936c9dbe4dc04963a448a9f43024839c Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 14:24:00 -0700 Subject: [PATCH 4/8] fix(proxy): redact a refused request as soon as the pre-call walk ends --- litellm/proxy/common_request_processing.py | 4 +--- .../proxy/test_common_request_processing.py | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4e9605cc4a3..066d5169e9a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2665,10 +2665,9 @@ class ProxyBaseLLMRequestProcessing: skip_guardrails=skip_guardrails, routing_probe=routing_probe, ) - except Exception: + finally: if routing_probe is not None and routing_probe.refusal is not None: self._redact_request_its_skipped_guardrails_never_scanned() - raise await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, @@ -3107,7 +3106,6 @@ class ProxyBaseLLMRequestProcessing: phase_event("litellm.request.pre_call_completed") refusal: Final = None if self._routing_probe is None else self._routing_probe.refusal if refusal is not None: - self._redact_request_its_skipped_guardrails_never_scanned() raise refusal _post_call_guardrails_active: Final = self._has_post_call_guardrails() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8054b3aad05..70170be7ef5 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -13058,6 +13058,13 @@ class _ModelMappingHook(CustomLogger): return data +class _TaggingHook(CustomLogger): + """A customer callback that adds a tag to the request.""" + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return {**data, "metadata": {**(data.get("metadata") or {}), "tags": ["over-budget-tag"]}} + + class _EmailMaskingHook(CustomLogger): async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): _mask_emails(data) @@ -13486,6 +13493,21 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: assert (first.status, limited.status, guardrail.calls) == (400, 429, 0) assert _shows_no_raw_content(limited), limited.failure_dumps + @pytest.mark.asyncio + async def test_failure_log_hides_content_when_a_tag_budget_rejects_after_a_skipped_guardrail( + self, monkeypatch: pytest.MonkeyPatch + ): + async def over_budget(**_: object) -> None: + raise litellm.BudgetExceededError(current_cost=10.0, max_budget=5.0) + + monkeypatch.setattr(litellm.proxy.common_request_processing, "tag_max_budget_check_for_tags", over_budget) + guardrail: Final = _VendorGuardrail(behavior="mask") + + sent: Final = await _send(monkeypatch, [_TaggingHook(), guardrail]) + + assert (sent.status, guardrail.calls) == (422, 0), sent.message + assert _shows_no_raw_content(sent), sent.failure_dumps + @pytest.mark.asyncio async def test_built_in_hook_after_a_skipped_guardrail_is_not_sent_the_content_it_would_have_masked( self, monkeypatch: pytest.MonkeyPatch From 78aedf990deac1a396402d71c3b01afc23063d91 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 14:41:22 -0800 Subject: [PATCH 5/8] fix(proxy): keep a skipped guardrail out of the logged applied guardrails The logged applied_guardrails also lists every default_on guardrail and the names in the request's metadata.guardrails, so a skipped guardrail is taken off that list when the walk ends. The new integration test pins the header on every route. --- litellm/proxy/common_request_processing.py | 12 +++++-- .../test_unknown_model_guardrail_skip.py | 30 +++++++++++++++-- .../proxy/test_common_request_processing.py | 33 +++++++++++++++++++ 3 files changed, 70 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 066d5169e9a..8dafde140b0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2667,7 +2667,7 @@ class ProxyBaseLLMRequestProcessing: ) finally: if routing_probe is not None and routing_probe.refusal is not None: - self._redact_request_its_skipped_guardrails_never_scanned() + self._redact_request_its_skipped_guardrails_never_scanned(routing_probe.skipped_guardrails) await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, @@ -2703,15 +2703,21 @@ class ProxyBaseLLMRequestProcessing: return self.data, logging_obj - def _redact_request_its_skipped_guardrails_never_scanned(self) -> None: + def _redact_request_its_skipped_guardrails_never_scanned(self, skipped_guardrails: frozenset[str]) -> None: """Redact the request body, its logged snapshot and the logged messages, so failure callbacks never see - content a skipped masking guardrail would have rewritten. + content a skipped masking guardrail would have rewritten, and leave the skipped guardrails out of the + logged ``applied_guardrails``. """ self.data.update(without_unscanned_content(self.data)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict refresh_proxy_server_request_body_snapshot(self.data) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # self.data is the untyped request dict logging_obj: Final = self.data.get("litellm_logging_obj") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # self.data is the untyped request dict if isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.update_messages([{"role": "user", "content": REDACTED_BY_LITELLM}]) + logging_obj.model_call_details["applied_guardrails"] = [ + name + for name in logging_obj.model_call_details.get("applied_guardrails") or () + if name not in skipped_guardrails + ] async def _pre_call_with_fallbacks( self, diff --git a/tests/integration/authorization/test_unknown_model_guardrail_skip.py b/tests/integration/authorization/test_unknown_model_guardrail_skip.py index d502eaa07ca..b2fbc326112 100644 --- a/tests/integration/authorization/test_unknown_model_guardrail_skip.py +++ b/tests/integration/authorization/test_unknown_model_guardrail_skip.py @@ -65,13 +65,13 @@ def _operator_verdict(action: str) -> Callable[[Request], Reply]: return reply -def _airs(policy_url: str) -> dict[str, JsonValue]: +def _airs(policy_url: str, default_on: bool = True) -> dict[str, JsonValue]: return { "guardrail_name": "airs-" + uuid.uuid4().hex[:8], "litellm_params": { "guardrail": "panw_prisma_airs", "mode": "pre_call", - "default_on": True, + "default_on": default_on, "api_base": policy_url, "api_key": "synthetic-panw-key", "profile_name": "synthetic-profile", @@ -186,6 +186,32 @@ def test_unknown_model_burst_while_the_scanner_is_down_never_reaches_it(gateway: assert [request.target for request in upstream.drain() if request.method == "POST"] == [] +def test_opt_in_guardrail_is_named_in_the_header_only_when_it_ran(gateway: Gateway, tmp_path: Path) -> None: + model: Final = "no-such-model-" + uuid.uuid4().hex[:8] + with wire_server(_scanner("allow")) as policy, wire_server(_provider) as upstream: + guardrail: Final = _airs(policy.url, default_on=False) + opt_in: Final = {"guardrails": [guardrail["guardrail_name"]]} + with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail)) as candidate: + with candidate.scenario() as scenario: + routable: Final = scenario.model( + model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key" + ) + ran: Final = candidate.request( + "POST", "/v1/chat/completions", {**_requests(routable)[0][1], "metadata": opt_in} + ) + assert ran.status_code == 200, ran.text + assert ran.headers["x-litellm-applied-guardrails"] == guardrail["guardrail_name"] + assert len(policy.drain()) == 1 + + for target, body in _requests(model): + refused = candidate.request("POST", target, {**body, "metadata": opt_in}) + + assert refused.status_code == 400, refused.text + assert "Invalid model name" in refused.text + assert "x-litellm-applied-guardrails" not in refused.headers, target + assert policy.drain() == () + + def test_operator_guardrail_and_the_guardrail_before_it_still_run_for_an_unknown_model( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 70170be7ef5..8b2c4a2adfd 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -13137,6 +13137,7 @@ class _FailureLogRecorder(CustomLogger): "request_model": logged_request.get("model"), "request_messages": logged_request.get("messages"), "request_has_call_id": "litellm_call_id" in logged_request, + "applied_guardrails": (standard.get("metadata") or {}).get("applied_guardrails"), } ) @@ -13540,6 +13541,38 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: assert _shows_no_raw_content(sent) is content_hidden, sent.failure_dumps +class TestUnroutableModelFailureLogNamesOnlyGuardrailsThatRan: + @pytest.mark.asyncio + @pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"]) + async def test_failure_log_leaves_out_a_skipped_default_on_guardrail(self, monkeypatch, run_in_parallel): + sent: Final = await _send(monkeypatch, [_VendorGuardrail(run_in_parallel=run_in_parallel)]) + + assert sent.status == 400 + assert sent.failure_logs[0]["applied_guardrails"] == [] + + @pytest.mark.asyncio + async def test_failure_log_keeps_an_operator_guardrail_that_ran(self, monkeypatch): + class OperatorGuardrail(_VendorGuardrail): + __module__ = "operator_guardrails.checks" + + operator: Final = OperatorGuardrail(name="operator") + skipped: Final = _VendorGuardrail(name="skipped") + + sent: Final = await _send(monkeypatch, [operator, skipped]) + + assert (operator.calls, skipped.calls) == (1, 0) + assert sent.failure_logs[0]["applied_guardrails"] == ["operator"] + + @pytest.mark.asyncio + async def test_failure_log_names_a_guardrail_that_blocked_a_routable_model(self, monkeypatch): + guardrail: Final = _VendorGuardrail(behavior="block") + + sent: Final = await _send(monkeypatch, [guardrail], model="gpt-4o-mini") + + assert guardrail.calls == 1 + assert sent.failure_logs[0]["applied_guardrails"] == ["vendor-guardrail"] + + class TestUnroutableModelLeavesEverythingElseAsBefore: """Everything except the skipped pre-call guardrails behaves as it did before LIT-8011, including the cases earlier attempts at this change broke.""" From a6a02ecaca26e9d5ec87ee7b191766695b8bc197 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 17:50:38 -0800 Subject: [PATCH 6/8] test(proxy): pin a refused pass-through body and a shared guardrail name in the failure log Register the doubles that stand in for litellm's own guardrails under a module file inside the litellm package, since the walk now tells shipped code from operator code by file. Type the integration guardrail double's request parameter. --- .../authorization/_operator_guardrail.py | 7 +++- .../proxy/test_common_request_processing.py | 37 ++++++++++++++++++- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/integration/authorization/_operator_guardrail.py b/tests/integration/authorization/_operator_guardrail.py index dec071acd9a..543c71db002 100644 --- a/tests/integration/authorization/_operator_guardrail.py +++ b/tests/integration/authorization/_operator_guardrail.py @@ -1,6 +1,7 @@ """A pre_call guardrail loaded from the operator's own module, the way a deployment's custom guardrail is.""" import os +from collections.abc import Mapping from typing import Final import httpx @@ -16,7 +17,11 @@ class OperatorGuardrail(CustomGuardrail): """Posts the requested model to ``OPERATOR_GUARDRAIL_URL`` and rejects with 403 when the reply says block.""" async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: Mapping[str, object], + call_type: CallTypesLiteral, ) -> None: async with httpx.AsyncClient(trust_env=False) as client: response: Final = await client.post(os.environ["OPERATOR_GUARDRAIL_URL"], json={"model": data.get("model")}) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8b2c4a2adfd..03cbdc69f08 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,9 +4,11 @@ import copy import datetime import gzip import json +import os +import sys from dataclasses import dataclass from collections.abc import AsyncIterable, Awaitable, Mapping -from types import MappingProxyType, SimpleNamespace +from types import MappingProxyType, ModuleType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -13032,6 +13034,15 @@ class _BuiltInContentCheck(CustomLogger): return data +def _register_shipped_doubles(monkeypatch: pytest.MonkeyPatch) -> None: + """Give the doubles that stand in for litellm's own guardrails and hooks a module file inside the litellm package.""" + package_dir: Final = os.path.dirname(litellm.__file__) + for double in (_VendorGuardrail, _BuiltInContentCheck): + module = ModuleType(double.__module__) + module.__file__ = os.path.join(package_dir, *double.__module__.split(".")[1:]) + ".py" + monkeypatch.setitem(sys.modules, double.__module__, module) + + class _ModelMappingHook(CustomLogger): """A customer callback that renames a virtual model, editing the request or returning a new one.""" @@ -13174,6 +13185,7 @@ async def _send( from litellm.proxy.common_utils import registry_read_through recorder: Final = _FailureLogRecorder() + _register_shipped_doubles(monkeypatch) monkeypatch.setattr(litellm, "callbacks", [*callbacks, recorder]) ProxyLogging._callback_capabilities_cache.clear() phase_events: Final = list[str]() @@ -13434,6 +13446,15 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: }, id="extraction-state-and-fields", ), + pytest.param( + { + "data": { + "messages": [{"role": "user", "content": [{"text": f"email {_RAW_EMAIL}"}]}], + "system": [{"text": f"email {_RAW_EMAIL}"}], + } + }, + id="pass-through-body", + ), ], ) async def test_failure_log_hides_every_field_a_skipped_guardrail_would_have_scanned(self, monkeypatch, body): @@ -13563,6 +13584,19 @@ class TestUnroutableModelFailureLogNamesOnlyGuardrailsThatRan: assert (operator.calls, skipped.calls) == (1, 0) assert sent.failure_logs[0]["applied_guardrails"] == ["operator"] + @pytest.mark.asyncio + async def test_failure_log_keeps_a_name_an_operator_guardrail_shares_with_a_skipped_one(self, monkeypatch): + class OperatorGuardrail(_VendorGuardrail): + __module__ = "operator_guardrails.checks" + + operator: Final = OperatorGuardrail(name="shared") + skipped: Final = _VendorGuardrail(name="shared") + + sent: Final = await _send(monkeypatch, [operator, skipped]) + + assert (operator.calls, skipped.calls) == (1, 0) + assert "shared" in sent.failure_logs[0]["applied_guardrails"] + @pytest.mark.asyncio async def test_failure_log_names_a_guardrail_that_blocked_a_routable_model(self, monkeypatch): guardrail: Final = _VendorGuardrail(behavior="block") @@ -13979,6 +14013,7 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: enforcer: Final = _ContentEnforcer() guardrail: Final = _VendorGuardrail() + _register_shipped_doubles(monkeypatch) monkeypatch.setattr(litellm, "callbacks", [_ModelMappingHook(), enforcer, guardrail]) ProxyLogging._callback_capabilities_cache.clear() router: Final = _routing_router() From fa828f344c36d2608a4512ddf3aa5bd4ebd0ea9f Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 18:44:59 -0800 Subject: [PATCH 7/8] test(proxy): type the unroutable-model test doubles' hooks Gives the doubles' pre-call, moderation and failure hooks the CustomLogger parameter types, and types the email-masking helper. --- .../proxy/test_common_request_processing.py | 78 ++++++++++++++----- 1 file changed, 57 insertions(+), 21 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 03cbdc69f08..9fd923a4589 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -12936,10 +12936,16 @@ def _routing_router(fallbacks: list[dict[str, list[str]]] | None = None) -> Rout ) -def _mask_emails(data: dict) -> None: - for message in data.get("messages") or []: - if isinstance(message.get("content"), str): - message["content"] = message["content"].replace(_RAW_EMAIL, "[EMAIL]") +def _mask_emails(data: dict[str, object]) -> None: + messages: Final = data.get("messages") + for message in messages if isinstance(messages, list) else []: + content = message.get("content") if isinstance(message, dict) else None + if isinstance(content, str): + message["content"] = content.replace(_RAW_EMAIL, "[EMAIL]") + + +def _as_mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else {} class _VendorGuardrail(CustomGuardrail): @@ -12961,7 +12967,9 @@ class _VendorGuardrail(CustomGuardrail): self.order: Final = order if order is not None else [] self.calls = 0 - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 self.order.append(self.guardrail_name or "") if self.behavior == "block": @@ -12974,7 +12982,9 @@ class _VendorGuardrail(CustomGuardrail): _mask_emails(data) return data - async def async_moderation_hook(self, data, user_api_key_dict, call_type): + async def async_moderation_hook( + self, data: dict, user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: self.calls += 1 @@ -12986,7 +12996,9 @@ class _CustomerRoutingGuardrail(CustomGuardrail): self.run_in_parallel = run_in_parallel self.calls = 0 - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 if data.get("model") == "virtual-model": data["model"] = "gpt-4o-mini" @@ -13059,7 +13071,9 @@ class _ModelMappingHook(CustomLogger): self.returns_new_request: Final = returns_new_request self.order: Final = order if order is not None else [] - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.order.append("mapping-hook") if data.get("model") != self.source: return data @@ -13072,12 +13086,16 @@ class _ModelMappingHook(CustomLogger): class _TaggingHook(CustomLogger): """A customer callback that adds a tag to the request.""" - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: return {**data, "metadata": {**(data.get("metadata") or {}), "tags": ["over-budget-tag"]}} class _EmailMaskingHook(CustomLogger): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: _mask_emails(data) return data @@ -13087,7 +13105,9 @@ class _RejectingHook(CustomLogger): super().__init__() self.model: Final = model - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: if self.model in (None, data.get("model")): raise HTTPException(status_code=403, detail="rejected by hook") return data @@ -13098,7 +13118,9 @@ class _RateLimitingHook(CustomLogger): super().__init__() self.model: Final = model - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: if data.get("model") == self.model: raise ProxyRateLimitError(detail=f"rpm limit reached for {self.model}") return data @@ -13112,7 +13134,9 @@ class _ContentEnforcer(CustomLogger): self.order: Final = order if order is not None else [] self.calls = 0 - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 self.order.append("content-enforcer") return data @@ -13123,7 +13147,9 @@ class _ModerationHook(CustomLogger): super().__init__() self.calls = 0 - async def async_moderation_hook(self, data, user_api_key_dict, call_type): + async def async_moderation_hook( + self, data: dict, user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: self.calls += 1 @@ -13133,10 +13159,16 @@ class _FailureLogRecorder(CustomLogger): self.logs: Final[list[dict[str, object]]] = [] self.dumps: Final[list[str]] = [] - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - litellm_params: Final = kwargs.get("litellm_params") or {} - logged_request: Final = (litellm_params.get("proxy_server_request") or {}).get("body") or {} - standard: Final = kwargs.get("standard_logging_object") or {} + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + litellm_params: Final = _as_mapping(kwargs.get("litellm_params")) + logged_request: Final = _as_mapping(_as_mapping(litellm_params.get("proxy_server_request")).get("body")) + standard: Final = _as_mapping(kwargs.get("standard_logging_object")) self.dumps.append(json.dumps(kwargs, default=str)) self.logs.append( { @@ -13148,13 +13180,17 @@ class _FailureLogRecorder(CustomLogger): "request_model": logged_request.get("model"), "request_messages": logged_request.get("messages"), "request_has_call_id": "litellm_call_id" in logged_request, - "applied_guardrails": (standard.get("metadata") or {}).get("applied_guardrails"), + "applied_guardrails": _as_mapping(standard.get("metadata")).get("applied_guardrails"), } ) async def async_post_call_failure_hook( - self, request_data, original_exception, user_api_key_dict, traceback_str=None - ): + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: ProxyUserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: self.dumps.append(json.dumps(request_data, default=str)) From cb8dbda69097410147043fa99f619b5cf9652c1e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 19:11:55 -0800 Subject: [PATCH 8/8] test(proxy): type every parameter of the unroutable-model tests and doubles Gives the new tests, the send and pipeline helpers and the hook doubles concrete parameter and return types in place of bare dict. --- .../proxy/test_common_request_processing.py | 260 ++++++++++++------ 1 file changed, 179 insertions(+), 81 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 9fd923a4589..e50bbff985c 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -12968,8 +12968,12 @@ class _VendorGuardrail(CustomGuardrail): self.calls = 0 async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: self.calls += 1 self.order.append(self.guardrail_name or "") if self.behavior == "block": @@ -12983,7 +12987,7 @@ class _VendorGuardrail(CustomGuardrail): return data async def async_moderation_hook( - self, data: dict, user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + self, data: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral ) -> None: self.calls += 1 @@ -12997,8 +13001,12 @@ class _CustomerRoutingGuardrail(CustomGuardrail): self.calls = 0 async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: self.calls += 1 if data.get("model") == "virtual-model": data["model"] = "gpt-4o-mini" @@ -13022,7 +13030,7 @@ class _CustomerTextGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: @@ -13040,8 +13048,12 @@ class _BuiltInContentCheck(CustomLogger): self.seen: Final[list[str]] = [] async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: self.seen.append(json.dumps(data.get("messages"))) return data @@ -13072,8 +13084,12 @@ class _ModelMappingHook(CustomLogger): self.order: Final = order if order is not None else [] async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: self.order.append("mapping-hook") if data.get("model") != self.source: return data @@ -13087,15 +13103,23 @@ class _TaggingHook(CustomLogger): """A customer callback that adds a tag to the request.""" async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: - return {**data, "metadata": {**(data.get("metadata") or {}), "tags": ["over-budget-tag"]}} + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + return {**data, "metadata": {**_as_mapping(data.get("metadata")), "tags": ["over-budget-tag"]}} class _EmailMaskingHook(CustomLogger): async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: _mask_emails(data) return data @@ -13106,8 +13130,12 @@ class _RejectingHook(CustomLogger): self.model: Final = model async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: if self.model in (None, data.get("model")): raise HTTPException(status_code=403, detail="rejected by hook") return data @@ -13119,8 +13147,12 @@ class _RateLimitingHook(CustomLogger): self.model: Final = model async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: if data.get("model") == self.model: raise ProxyRateLimitError(detail=f"rpm limit reached for {self.model}") return data @@ -13135,8 +13167,12 @@ class _ContentEnforcer(CustomLogger): self.calls = 0 async def async_pre_call_hook( - self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: self.calls += 1 self.order.append("content-enforcer") return data @@ -13148,7 +13184,7 @@ class _ModerationHook(CustomLogger): self.calls = 0 async def async_moderation_hook( - self, data: dict, user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + self, data: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral ) -> None: self.calls += 1 @@ -13186,7 +13222,7 @@ class _FailureLogRecorder(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: ProxyUserAPIKeyAuth, traceback_str: str | None = None, @@ -13205,7 +13241,7 @@ class _Sent: async def _send( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, callbacks: Sequence[CustomLogger], model: str | None = "does-not-exist", route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"] = "acompletion", @@ -13289,7 +13325,9 @@ async def _send( ) -def _attach_policy_pipeline(monkeypatch, mode: Literal["pre_call", "post_call"], guardrail_name: str) -> None: +def _attach_policy_pipeline( + monkeypatch: pytest.MonkeyPatch, mode: Literal["pre_call", "post_call"], guardrail_name: str +) -> None: from litellm.proxy.policy_engine import attachment_registry, policy_registry policies: Final = policy_registry.PolicyRegistry() @@ -13321,8 +13359,8 @@ class TestUnroutableModelSkipsPreCallGuardrails: @pytest.mark.asyncio @pytest.mark.parametrize("behavior", ["pass", "mask", "block", "unreachable"]) async def test_unknown_model_gets_the_invalid_model_error_without_calling_the_guardrail( - self, monkeypatch, behavior - ): + self, monkeypatch: pytest.MonkeyPatch, behavior: Literal["pass", "mask", "block", "unreachable"] + ) -> None: guardrail: Final = _VendorGuardrail(behavior=behavior) guarded: Final = await _send(monkeypatch, [guardrail]) @@ -13339,8 +13377,8 @@ class TestUnroutableModelSkipsPreCallGuardrails: [pytest.param(None, id="missing"), pytest.param("", id="empty"), pytest.param("m" * 5000, id="5000-chars")], ) async def test_missing_or_malformed_model_gets_the_routing_error_without_calling_the_guardrail( - self, monkeypatch, model - ): + self, monkeypatch: pytest.MonkeyPatch, model: str | None + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") guarded: Final = await _send(monkeypatch, [guardrail], model=model) @@ -13351,7 +13389,7 @@ class TestUnroutableModelSkipsPreCallGuardrails: assert guardrail.calls == 0 @pytest.mark.asyncio - async def test_unknown_model_skips_a_parallel_guardrail(self, monkeypatch): + async def test_unknown_model_skips_a_parallel_guardrail(self, monkeypatch: pytest.MonkeyPatch) -> None: guardrail: Final = _VendorGuardrail(behavior="block", run_in_parallel=True) sent: Final = await _send(monkeypatch, [guardrail]) @@ -13361,7 +13399,9 @@ class TestUnroutableModelSkipsPreCallGuardrails: assert guardrail.calls == 0 @pytest.mark.asyncio - async def test_endpoint_without_a_router_skips_the_guardrail_for_an_unknown_model(self, monkeypatch): + async def test_endpoint_without_a_router_skips_the_guardrail_for_an_unknown_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") guarded: Final = await _send(monkeypatch, [guardrail], use_router=False) @@ -13381,7 +13421,12 @@ class TestUnroutableModelSkipsPreCallGuardrails: pytest.param("atext_completion", {}, id="completions"), ], ) - async def test_unknown_model_on_other_routes_skips_the_guardrail(self, monkeypatch, route_type, body): + async def test_unknown_model_on_other_routes_skips_the_guardrail( + self, + monkeypatch: pytest.MonkeyPatch, + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"], + body: Mapping[str, object], + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") guarded: Final = await _send(monkeypatch, [guardrail], route_type=route_type, body=body) @@ -13392,7 +13437,9 @@ class TestUnroutableModelSkipsPreCallGuardrails: assert guardrail.calls == 0 @pytest.mark.asyncio - async def test_model_a_hook_swaps_for_another_unknown_model_is_rejected_under_the_new_name(self, monkeypatch): + async def test_model_a_hook_swaps_for_another_unknown_model_is_rejected_under_the_new_name( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") swap: Final = _ModelMappingHook(target="still-unknown", returns_new_request=True) @@ -13403,7 +13450,9 @@ class TestUnroutableModelSkipsPreCallGuardrails: assert guardrail.calls == 0 @pytest.mark.asyncio - async def test_guardrail_that_would_reroute_an_unknown_model_is_not_called(self, monkeypatch): + async def test_guardrail_that_would_reroute_an_unknown_model_is_not_called( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="reroute") sent: Final = await _send(monkeypatch, [guardrail], model="guard-reroute") @@ -13413,7 +13462,9 @@ class TestUnroutableModelSkipsPreCallGuardrails: assert guardrail.calls == 0 @pytest.mark.asyncio - async def test_rate_limit_fallback_to_an_unknown_model_skips_the_guardrail(self, monkeypatch): + async def test_rate_limit_fallback_to_an_unknown_model_skips_the_guardrail( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="mask") sent: Final = await _send( @@ -13430,8 +13481,8 @@ class TestUnroutableModelSkipsPreCallGuardrails: @pytest.mark.asyncio async def test_key_with_router_settings_gets_the_invalid_model_error_without_calling_the_guardrail( - self, monkeypatch - ): + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") key: Final = ProxyUserAPIKeyAuth(request_route="/v1/chat/completions", router_settings={"num_retries": 0}) @@ -13493,7 +13544,9 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: ), ], ) - async def test_failure_log_hides_every_field_a_skipped_guardrail_would_have_scanned(self, monkeypatch, body): + async def test_failure_log_hides_every_field_a_skipped_guardrail_would_have_scanned( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object] + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask")], body=body) assert sent.status == 400 @@ -13516,7 +13569,12 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: pytest.param("atext_completion", {}, id="completions"), ], ) - async def test_failure_log_hides_the_input_of_other_routes(self, monkeypatch, route_type, body): + async def test_failure_log_hides_the_input_of_other_routes( + self, + monkeypatch: pytest.MonkeyPatch, + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"], + body: Mapping[str, object], + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask")], route_type=route_type, body=body) assert sent.status == 400 @@ -13525,7 +13583,7 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: @pytest.mark.asyncio async def test_failure_log_hides_content_a_skipped_parallel_guardrail_would_have_scanned( self, monkeypatch: pytest.MonkeyPatch - ): + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask", run_in_parallel=True)]) assert sent.status == 400 @@ -13533,8 +13591,8 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: @pytest.mark.asyncio async def test_failure_log_hides_content_when_a_built_in_hook_rejects_after_a_skipped_guardrail( - self, monkeypatch - ): + self, monkeypatch: pytest.MonkeyPatch + ) -> None: proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) key: Final = ProxyUserAPIKeyAuth( api_key=hash_token("sk-rpm-one"), @@ -13554,7 +13612,7 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: @pytest.mark.asyncio async def test_failure_log_hides_content_when_a_tag_budget_rejects_after_a_skipped_guardrail( self, monkeypatch: pytest.MonkeyPatch - ): + ) -> None: async def over_budget(**_: object) -> None: raise litellm.BudgetExceededError(current_cost=10.0, max_budget=5.0) @@ -13569,7 +13627,7 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: @pytest.mark.asyncio async def test_built_in_hook_after_a_skipped_guardrail_is_not_sent_the_content_it_would_have_masked( self, monkeypatch: pytest.MonkeyPatch - ): + ) -> None: check: Final = _BuiltInContentCheck() sent: Final = await _send(monkeypatch, [_VendorGuardrail(behavior="mask"), check]) @@ -13587,8 +13645,8 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: ], ) async def test_failure_log_hides_content_only_when_the_skipped_guardrail_applied( - self, monkeypatch, stream_scope, stream, content_hidden - ): + self, monkeypatch: pytest.MonkeyPatch, stream_scope: str, stream: bool, content_hidden: bool + ) -> None: guardrail: Final = _VendorGuardrail(behavior="mask") guardrail.apply_stream_scope(stream_scope) @@ -13601,14 +13659,16 @@ class TestUnroutableModelFailureLogHidesUnscannedContent: class TestUnroutableModelFailureLogNamesOnlyGuardrailsThatRan: @pytest.mark.asyncio @pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"]) - async def test_failure_log_leaves_out_a_skipped_default_on_guardrail(self, monkeypatch, run_in_parallel): + async def test_failure_log_leaves_out_a_skipped_default_on_guardrail( + self, monkeypatch: pytest.MonkeyPatch, run_in_parallel: bool + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail(run_in_parallel=run_in_parallel)]) assert sent.status == 400 assert sent.failure_logs[0]["applied_guardrails"] == [] @pytest.mark.asyncio - async def test_failure_log_keeps_an_operator_guardrail_that_ran(self, monkeypatch): + async def test_failure_log_keeps_an_operator_guardrail_that_ran(self, monkeypatch: pytest.MonkeyPatch) -> None: class OperatorGuardrail(_VendorGuardrail): __module__ = "operator_guardrails.checks" @@ -13621,7 +13681,9 @@ class TestUnroutableModelFailureLogNamesOnlyGuardrailsThatRan: assert sent.failure_logs[0]["applied_guardrails"] == ["operator"] @pytest.mark.asyncio - async def test_failure_log_keeps_a_name_an_operator_guardrail_shares_with_a_skipped_one(self, monkeypatch): + async def test_failure_log_keeps_a_name_an_operator_guardrail_shares_with_a_skipped_one( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: class OperatorGuardrail(_VendorGuardrail): __module__ = "operator_guardrails.checks" @@ -13634,7 +13696,9 @@ class TestUnroutableModelFailureLogNamesOnlyGuardrailsThatRan: assert "shared" in sent.failure_logs[0]["applied_guardrails"] @pytest.mark.asyncio - async def test_failure_log_names_a_guardrail_that_blocked_a_routable_model(self, monkeypatch): + async def test_failure_log_names_a_guardrail_that_blocked_a_routable_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="block") sent: Final = await _send(monkeypatch, [guardrail], model="gpt-4o-mini") @@ -13659,8 +13723,14 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ) @pytest.mark.parametrize("check", ["guardrail", "parallel-guardrail", "content-enforcer"]) async def test_hook_listed_after_a_pre_call_check_can_still_rename_a_virtual_model( - self, monkeypatch, route_type, body, virtual, target, check - ): + self, + monkeypatch: pytest.MonkeyPatch, + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"], + body: Mapping[str, object], + virtual: str, + target: str, + check: str, + ) -> None: guardrail: Final = _VendorGuardrail(run_in_parallel=check == "parallel-guardrail") enforcer: Final = _ContentEnforcer() mapping: Final = _ModelMappingHook(source=virtual, target=target) @@ -13680,7 +13750,7 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: @pytest.mark.parametrize("customer_guardrail", [_CustomerRoutingGuardrail, _InheritedCustomerGuardrail]) async def test_customer_guardrail_can_still_rename_a_virtual_model( self, monkeypatch: pytest.MonkeyPatch, customer_guardrail: type[_CustomerRoutingGuardrail] - ): + ) -> None: guardrail: Final = _VendorGuardrail() customer: Final = customer_guardrail() @@ -13699,7 +13769,7 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: monkeypatch: pytest.MonkeyPatch, customer_guardrail: type[_CustomerRoutingGuardrail] | type[_CustomerTextGuardrail], run_in_parallel: bool, - ): + ) -> None: customer: Final = customer_guardrail(run_in_parallel=run_in_parallel) sent: Final = await _send(monkeypatch, [customer]) @@ -13711,8 +13781,8 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: @pytest.mark.asyncio @pytest.mark.parametrize("mapping_first", [True, False]) async def test_pre_call_pipeline_does_not_stop_a_hook_from_renaming_a_virtual_model( - self, monkeypatch, mapping_first - ): + self, monkeypatch: pytest.MonkeyPatch, mapping_first: bool + ) -> None: guardrail: Final = _VendorGuardrail(name="pipeline-guardrail") _attach_policy_pipeline(monkeypatch, "pre_call", "pipeline-guardrail") mapping: Final = _ModelMappingHook() @@ -13725,7 +13795,7 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert guardrail.calls == 1 @pytest.mark.asyncio - async def test_rate_limiter_still_counts_requests_for_unknown_models(self, monkeypatch): + async def test_rate_limiter_still_counts_requests_for_unknown_models(self, monkeypatch: pytest.MonkeyPatch) -> None: proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) key: Final = ProxyUserAPIKeyAuth( api_key=hash_token("sk-rpm-one"), @@ -13745,13 +13815,15 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert statuses == [400, 429, 429] @pytest.mark.asyncio - async def test_hook_listed_after_a_guardrail_still_runs_for_an_unknown_model(self, monkeypatch): + async def test_hook_listed_after_a_guardrail_still_runs_for_an_unknown_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail(), _RejectingHook()]) assert (sent.status, sent.message) == (403, "rejected by hook") @pytest.mark.asyncio - async def test_during_call_hooks_still_run_for_an_unknown_model(self, monkeypatch): + async def test_during_call_hooks_still_run_for_an_unknown_model(self, monkeypatch: pytest.MonkeyPatch) -> None: moderation: Final = _ModerationHook() during_call_guardrail: Final = _VendorGuardrail(name="during-call", event_hook=GuardrailEventHooks.during_call) @@ -13762,7 +13834,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: @pytest.mark.asyncio @pytest.mark.parametrize("with_guardrail", [False, True]) - async def test_pre_call_completed_event_is_still_emitted_for_an_unknown_model(self, monkeypatch, with_guardrail): + async def test_pre_call_completed_event_is_still_emitted_for_an_unknown_model( + self, monkeypatch: pytest.MonkeyPatch, with_guardrail: bool + ) -> None: sent: Final = await _send(monkeypatch, [_VendorGuardrail()] if with_guardrail else []) assert sent.status == 400 @@ -13795,8 +13869,8 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ], ) async def test_guardrail_block_after_a_rename_logs_the_same_models_as_before( - self, monkeypatch, returns_new_request, expected - ): + self, monkeypatch: pytest.MonkeyPatch, returns_new_request: bool, expected: Mapping[str, object] + ) -> None: mapping: Final = _ModelMappingHook(returns_new_request=returns_new_request) sent: Final = await _send(monkeypatch, [mapping, _VendorGuardrail(behavior="block")], model="virtual-model") @@ -13806,7 +13880,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert [log["request_has_call_id"] for log in sent.failure_logs] == [False] @pytest.mark.asyncio - async def test_hook_rejection_after_the_guardrail_passed_logs_the_request_as_sent(self, monkeypatch): + async def test_hook_rejection_after_the_guardrail_passed_logs_the_request_as_sent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: sent: Final = await _send( monkeypatch, [_ModelMappingHook(), _VendorGuardrail(), _RejectingHook()], model="virtual-model" ) @@ -13817,7 +13893,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ] == [("virtual-model", [{"role": "user", "content": f"email {_RAW_EMAIL}"}], False)] @pytest.mark.asyncio - async def test_unknown_model_without_guardrails_logs_the_request_its_hooks_produced(self, monkeypatch): + async def test_unknown_model_without_guardrails_logs_the_request_its_hooks_produced( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: sent: Final = await _send(monkeypatch, [_EmailMaskingHook()]) assert sent.status == 400 @@ -13827,7 +13905,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ] == [([{"role": "user", "content": "email [EMAIL]"}], [{"role": "user", "content": "email [EMAIL]"}], True)] @pytest.mark.asyncio - async def test_post_call_only_pipeline_leaves_the_failure_log_of_an_unknown_model_unredacted(self, monkeypatch): + async def test_post_call_only_pipeline_leaves_the_failure_log_of_an_unknown_model_unredacted( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(name="post-call-guardrail", event_hook=GuardrailEventHooks.post_call) _attach_policy_pipeline(monkeypatch, "post_call", "post-call-guardrail") @@ -13839,7 +13919,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ] @pytest.mark.asyncio - async def test_rejection_by_a_hook_listed_before_the_guardrail_is_logged_as_before(self, monkeypatch): + async def test_rejection_by_a_hook_listed_before_the_guardrail_is_logged_as_before( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="mask") sent: Final = await _send(monkeypatch, [_RejectingHook(), guardrail]) @@ -13851,7 +13933,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ] @pytest.mark.asyncio - async def test_rate_limited_unknown_model_whose_fallback_a_hook_rejects_is_logged_as_before(self, monkeypatch): + async def test_rate_limited_unknown_model_whose_fallback_a_hook_rejects_is_logged_as_before( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail(behavior="mask") sent: Final = await _send( @@ -13886,16 +13970,20 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ], ) async def test_model_registry_is_read_once_for_the_name_routing_ends_up_with( - self, monkeypatch, model, callbacks, expected_lookups - ): + self, + monkeypatch: pytest.MonkeyPatch, + model: str, + callbacks: Callable[[], Sequence[CustomLogger]], + expected_lookups: tuple[str, ...], + ) -> None: sent: Final = await _send(monkeypatch, callbacks(), model=model) assert sent.registry_lookups == expected_lookups @pytest.mark.asyncio async def test_model_created_on_another_replica_is_served_with_its_guardrail_on_the_first_request( - self, monkeypatch - ): + self, monkeypatch: pytest.MonkeyPatch + ) -> None: router: Final = _routing_router() guardrail: Final = _VendorGuardrail() @@ -13918,7 +14006,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert sent.registry_lookups == ("created-on-a-sibling",) @pytest.mark.asyncio - async def test_routable_model_runs_every_hook_and_guardrail_once_in_configured_order(self, monkeypatch): + async def test_routable_model_runs_every_hook_and_guardrail_once_in_configured_order( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: order: Final = list[str]() sent: Final = await _send( @@ -13930,7 +14020,7 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert (sent.status, order) == (200, ["mapping-hook", "vendor-guardrail", "content-enforcer"]) @pytest.mark.asyncio - async def test_guardrail_still_blocks_a_routable_model(self, monkeypatch): + async def test_guardrail_still_blocks_a_routable_model(self, monkeypatch: pytest.MonkeyPatch) -> None: guardrail: Final = _VendorGuardrail(behavior="block") sent: Final = await _send(monkeypatch, [guardrail, _RejectingHook()], model="gpt-4o-mini") @@ -13938,7 +14028,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert (sent.status, sent.message, guardrail.calls) == (400, "Prompt blocked by vendor policy", 1) @pytest.mark.asyncio - async def test_hook_that_renames_a_blocked_model_to_a_routable_one_is_served(self, monkeypatch): + async def test_hook_that_renames_a_blocked_model_to_a_routable_one_is_served( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail() sent: Final = await _send( @@ -13986,8 +14078,8 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: ], ) async def test_request_the_router_resolves_itself_still_runs_its_guardrail( - self, monkeypatch, model, body, llm_router - ): + self, monkeypatch: pytest.MonkeyPatch, model: str, body: Mapping[str, object], llm_router: Callable[[], Router] + ) -> None: guardrail: Final = _VendorGuardrail() await _send(monkeypatch, [guardrail], model=model, body=body, llm_router=llm_router()) @@ -13996,7 +14088,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["content-enforcer", "pipeline-guardrail"]) - async def test_unknown_model_still_runs_pipelines_and_content_enforcers(self, monkeypatch, kind): + async def test_unknown_model_still_runs_pipelines_and_content_enforcers( + self, monkeypatch: pytest.MonkeyPatch, kind: str + ) -> None: guardrail: Final = _VendorGuardrail(name="pipeline-guardrail") enforcer: Final = _ContentEnforcer() if kind == "pipeline-guardrail": @@ -14008,7 +14102,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert (enforcer.calls if kind == "content-enforcer" else guardrail.calls) == 1 @pytest.mark.asyncio - async def test_guardrails_still_run_for_an_unknown_model_when_one_reroutes_sensitive_data(self, monkeypatch): + async def test_guardrails_still_run_for_an_unknown_model_when_one_reroutes_sensitive_data( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: guardrail: Final = _VendorGuardrail() router_guardrail: Final = CustomGuardrail( guardrail_name="sensitive-router", @@ -14023,8 +14119,8 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: @pytest.mark.asyncio async def test_built_in_hook_that_routes_an_unknown_model_after_a_skipped_guardrail_gets_the_routing_error( - self, monkeypatch - ): + self, monkeypatch: pytest.MonkeyPatch + ) -> None: from litellm.proxy.hooks.sensitive_data_routing import PROXY_SensitiveDataRoutingHandler from litellm.proxy.utils import InternalUsageCache @@ -14044,7 +14140,9 @@ class TestUnroutableModelLeavesEverythingElseAsBefore: assert "Invalid model name passed in model=does-not-exist" in sent.message @pytest.mark.asyncio - async def test_caller_that_skips_guardrails_still_runs_content_enforcers_for_a_renamed_model(self, monkeypatch): + async def test_caller_that_skips_guardrails_still_runs_content_enforcers_for_a_renamed_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: from litellm.proxy import proxy_server enforcer: Final = _ContentEnforcer()