diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9282ee7f703..8dafde140b0 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, @@ -133,6 +134,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 +146,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 +2170,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 +2464,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 +2646,28 @@ 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, + ) + finally: + if routing_probe is not None and routing_probe.refusal is not None: + 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, @@ -2681,6 +2703,22 @@ class ProxyBaseLLMRequestProcessing: return self.data, logging_obj + 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, 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, request: Request, @@ -2697,6 +2735,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 +2776,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 +2851,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 +3107,12 @@ 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: + 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..543c71db002 --- /dev/null +++ b/tests/integration/authorization/_operator_guardrail.py @@ -0,0 +1,29 @@ +"""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 +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: 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")}) + 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..b2fbc326112 --- /dev/null +++ b/tests/integration/authorization/test_unknown_model_guardrail_skip.py @@ -0,0 +1,243 @@ +"""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, 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": default_on, + "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_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: + 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..e50bbff985c 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,8 +4,11 @@ import copy import datetime import gzip import json -from collections.abc import AsyncIterable, Mapping -from types import MappingProxyType, SimpleNamespace +import os +import sys +from dataclasses import dataclass +from collections.abc import AsyncIterable, Awaitable, Mapping +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 @@ -31,14 +34,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 +101,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 +457,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 +734,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 +795,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 +971,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 +1114,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 +1185,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 +2233,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 +8096,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 +8578,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 +9896,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 +9947,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 +9988,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 +10073,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 +12055,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 +12892,1281 @@ 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[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): + """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: 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": + 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: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + 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: 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" + 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[str, object], + 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[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.seen.append(json.dumps(data.get("messages"))) + 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.""" + + 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: 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 + if self.returns_new_request: + return {**data, "model": self.target} + data["model"] = self.target + 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: 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[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + _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: 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 + + +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: 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 + + +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: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + 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: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + 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: 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( + { + "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, + "applied_guardrails": _as_mapping(standard.get("metadata")).get("applied_guardrails"), + } + ) + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: ProxyUserAPIKeyAuth, + traceback_str: str | None = None, + ) -> 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: pytest.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() + _register_shipped_doubles(monkeypatch) + 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: 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() + 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: pytest.MonkeyPatch, behavior: Literal["pass", "mask", "block", "unreachable"] + ) -> None: + 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: pytest.MonkeyPatch, model: str | None + ) -> None: + 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: pytest.MonkeyPatch) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: 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) + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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", + ), + pytest.param( + { + "state": f"email {_RAW_EMAIL}", + "fields": { + "contact": {"type": "string", "instructions": f"email {_RAW_EMAIL}", "unit": _RAW_EMAIL} + }, + }, + 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: pytest.MonkeyPatch, body: Mapping[str, object] + ) -> None: + 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 + @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: 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 + 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 + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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_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) + + 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 + ) -> None: + 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: pytest.MonkeyPatch, stream_scope: str, stream: bool, content_hidden: bool + ) -> None: + 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 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: 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: pytest.MonkeyPatch) -> None: + 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_keeps_a_name_an_operator_guardrail_shares_with_a_skipped_one( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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.""" + + @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: 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) + + 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] + ) -> None: + 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, + ) -> None: + 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: pytest.MonkeyPatch, mapping_first: bool + ) -> None: + 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: pytest.MonkeyPatch) -> None: + 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: 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: pytest.MonkeyPatch) -> None: + 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: pytest.MonkeyPatch, with_guardrail: bool + ) -> None: + 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: 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") + + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: 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()) + + 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: pytest.MonkeyPatch, kind: str + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + 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: pytest.MonkeyPatch + ) -> None: + from litellm.proxy import proxy_server + + 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() + 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.