mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(guardrails): check a model-level post_call response once across the proxy and deployment layers (#50300)
#40154 made the deployment post_call hook really run apply_guardrail guardrails, so the proxy checked the same response a second time (LIT-10481). The deployment hook now records guardrail + response hash on the request and on the response, and the proxy skips a guardrail only on a matching record. The hash covers every field except ids, usage, model and other metadata set between layers, so a response rewritten between the layers (OCR included) is still checked. Guardrails with their own post_call hook keep their pre-#40154 behavior Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yujonglee
GitHub
parent
5d66f2d8ff
commit
2b4485dee5
4 changed files with 687 additions and 18 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import contextvars
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
|
|
@ -10,7 +11,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -74,6 +75,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
|
|||
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
|
||||
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
|
||||
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
|
||||
_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16)
|
||||
|
||||
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
|
||||
|
||||
|
|
@ -286,6 +288,95 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b
|
|||
return False
|
||||
|
||||
|
||||
def _record_guardrail_marker(
|
||||
data: dict[str, object], # mutable-ok: the request payload the marker is recorded on
|
||||
marker: str | None,
|
||||
) -> None:
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list):
|
||||
if marker not in executed:
|
||||
executed.append(marker)
|
||||
else:
|
||||
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
|
||||
_RESPONSE_METADATA_KEYS: Final = frozenset(
|
||||
{
|
||||
"id",
|
||||
"created",
|
||||
"created_at",
|
||||
"model",
|
||||
"usage",
|
||||
"usage_info",
|
||||
"system_fingerprint",
|
||||
"_hidden_params",
|
||||
"_response_ms",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _response_fingerprint(response: object) -> str | None:
|
||||
"""Digest of everything but metadata set between layers, so an unknown field counts as content."""
|
||||
try:
|
||||
fields: Final[Mapping[str, object]] = (
|
||||
response.model_dump(mode="json")
|
||||
if isinstance(response, BaseModel)
|
||||
else _REQUEST_FIELDS.validate_python(response)
|
||||
)
|
||||
content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS}
|
||||
if not content:
|
||||
return None
|
||||
return hashlib.sha256(
|
||||
json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
|
||||
).hexdigest()
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks"
|
||||
|
||||
|
||||
def _str_keyed_dict(value: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name
|
||||
|
||||
|
||||
def _response_hidden_params(response: object) -> Mapping[str, object] | None:
|
||||
as_dict: Final = _str_keyed_dict(response)
|
||||
return _str_keyed_dict(
|
||||
as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None)
|
||||
)
|
||||
|
||||
|
||||
def _write_response_marker(response: object, marker: str | None) -> None:
|
||||
"""Record a check on the response itself, so it survives the router copying request metadata per hop."""
|
||||
if marker is None:
|
||||
return
|
||||
hidden: Final = _response_hidden_params(response)
|
||||
if not isinstance(hidden, dict):
|
||||
return
|
||||
hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple(
|
||||
_str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker}
|
||||
)
|
||||
|
||||
|
||||
def _response_carries_marker(response: object, marker: str | None) -> bool:
|
||||
hidden: Final = _response_hidden_params(response)
|
||||
return (
|
||||
marker is not None
|
||||
and hidden is not None
|
||||
and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY))
|
||||
)
|
||||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
|
|
@ -918,20 +1009,7 @@ class CustomGuardrail(CustomLogger):
|
|||
With ``raised``, records instead that the hook raised as a policy pipeline
|
||||
step the pipeline then moved past. The deployment-level hook still runs it.
|
||||
"""
|
||||
marker: Final = self._pre_call_marker(raised=raised)
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list):
|
||||
if marker not in executed:
|
||||
executed.append(marker)
|
||||
else:
|
||||
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
_record_guardrail_marker(data, self._pre_call_marker(raised=raised))
|
||||
|
||||
def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
|
||||
"""Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
|
||||
|
|
@ -1129,10 +1207,26 @@ class CustomGuardrail(CustomLogger):
|
|||
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
|
||||
)
|
||||
|
||||
if not self._is_valid_response_type(result):
|
||||
return None
|
||||
returned: Final = result if self._is_valid_response_type(result) else None
|
||||
if target is self:
|
||||
return returned
|
||||
checked: Final = returned if returned is not None else response
|
||||
marker: Final = self._post_call_check_marker(checked)
|
||||
_record_guardrail_marker(
|
||||
cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed
|
||||
marker,
|
||||
)
|
||||
_write_response_marker(checked, marker)
|
||||
return returned
|
||||
|
||||
return result
|
||||
def _post_call_check_marker(self, response: object) -> str | None:
|
||||
fingerprint: Final = _response_fingerprint(response)
|
||||
return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None
|
||||
|
||||
def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool:
|
||||
"""Whether a deployment-level hook of this request already ran this guardrail on this exact response."""
|
||||
marker: Final = self._post_call_check_marker(response)
|
||||
return _request_carries_marker(data, marker) or _response_carries_marker(response, marker)
|
||||
|
||||
async def async_logging_hook(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -758,6 +758,7 @@ def _drop_token_ids_ahead_of_chunk_check(
|
|||
|
||||
|
||||
_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {})
|
||||
|
||||
|
||||
|
|
@ -4418,6 +4419,7 @@ class ProxyLogging:
|
|||
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
|
||||
|
||||
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
|
||||
request_view: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
|
||||
try:
|
||||
# Merge model-level guardrails before checking which guardrails to run
|
||||
|
|
@ -4448,6 +4450,9 @@ class ProxyLogging:
|
|||
):
|
||||
continue
|
||||
|
||||
if callback.post_call_already_checked(request_view, response) is True:
|
||||
continue
|
||||
|
||||
guardrail_response: Any | None = None
|
||||
|
||||
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
|
||||
|
|
@ -4521,10 +4526,13 @@ class ProxyLogging:
|
|||
immediately before awaiting, and the unified hook pops it before its first
|
||||
suspension point, so concurrent guardrails never race on that key.
|
||||
"""
|
||||
request_view: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
|
||||
async def _run_one(callback: CustomGuardrail) -> None:
|
||||
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
|
||||
return
|
||||
if callback.post_call_already_checked(request_view, response) is True:
|
||||
return
|
||||
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
|
||||
data["guardrail_to_apply"] = callback
|
||||
await self._run_guardrail_with_metrics(
|
||||
|
|
|
|||
566
tests/unit/integrations/test_custom_guardrail_post_call_once.py
Normal file
566
tests/unit/integrations/test_custom_guardrail_post_call_once.py
Normal file
|
|
@ -0,0 +1,566 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterable, Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from fastapi import HTTPException, Request, Response
|
||||
from fastapi.responses import StreamingResponse
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
||||
from litellm.router import Router
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypes, Choices, GenericGuardrailAPIInputs, LLMResponseTypes, Message, ModelResponse
|
||||
|
||||
pytestmark = pytest.mark.respx(assert_all_called=False)
|
||||
|
||||
REPLY: Final = "hi from the model"
|
||||
REWRITTEN: Final = "rewritten by the proxy"
|
||||
MASKED: Final = "masked by G"
|
||||
|
||||
Route = Literal["acompletion", "aresponses", "anthropic_messages"]
|
||||
|
||||
_OPENAI_API_BASE: Final = "https://api.openai.com/v1"
|
||||
_FAILING_API_BASE: Final = "https://primary.invalid/v1"
|
||||
_MCP_API_BASE: Final = "https://mcp-rounds.invalid/v1"
|
||||
|
||||
_CHAT_REPLY: Final[Mapping[str, object]] = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4o",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": REPLY}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
_CHAT_SSE: Final = (
|
||||
"".join(
|
||||
f"data: {json.dumps(chunk)}\n\n"
|
||||
for chunk in (
|
||||
{
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "gpt-4o",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": REPLY}, "finish_reason": None}],
|
||||
},
|
||||
{
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "gpt-4o",
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
},
|
||||
)
|
||||
)
|
||||
+ "data: [DONE]\n\n"
|
||||
)
|
||||
_RESPONSES_REPLY: Final[Mapping[str, object]] = {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 0,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": REPLY, "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
_RESPONSES_SSE: Final = "".join(
|
||||
f"data: {json.dumps(event)}\n\n"
|
||||
for event in (
|
||||
{"type": "response.created", "sequence_number": 0, "response": {**_RESPONSES_REPLY, "output": []}},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": REPLY,
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": _RESPONSES_REPLY},
|
||||
)
|
||||
)
|
||||
_TOOL_CALL_ROUND: Final[Mapping[str, object]] = {
|
||||
**_RESPONSES_REPLY,
|
||||
"id": "resp_tool_call",
|
||||
"output": [
|
||||
{
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"call_id": "call_echo",
|
||||
"name": "echo",
|
||||
"arguments": '{"value": "PING"}',
|
||||
"status": "completed",
|
||||
}
|
||||
],
|
||||
}
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
class _CountedGuardrail(CustomGuardrail):
|
||||
checked: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class ApplyGuardrailG(_CountedGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: object = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.checked = (*self.checked, " ".join(inputs.get("texts") or ()))
|
||||
return inputs
|
||||
|
||||
|
||||
def _reply_text(response: LLMResponseTypes) -> str:
|
||||
match response:
|
||||
case ModelResponse(choices=[Choices(message=Message(content=str() as text))]):
|
||||
return text
|
||||
case _:
|
||||
return str(response)
|
||||
|
||||
|
||||
class OwnHookG(_CountedGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
) -> None:
|
||||
self.checked = (*self.checked, _reply_text(response))
|
||||
|
||||
|
||||
class InheritedApplyGuardrailG(ApplyGuardrailG):
|
||||
pass
|
||||
|
||||
|
||||
class MaskingG(ApplyGuardrailG):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: object = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
checked: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
|
||||
return {**checked, "texts": [MASKED for _ in checked.get("texts") or ()]}
|
||||
|
||||
|
||||
class BlockingG(ApplyGuardrailG):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: object = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
|
||||
raise HTTPException(status_code=400, detail="blocked by G")
|
||||
|
||||
|
||||
class RewritingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: object = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return {**inputs, "texts": [REWRITTEN for _ in inputs.get("texts") or ()]}
|
||||
|
||||
|
||||
class ToolCallBlockingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: object = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if input_type == "response" and inputs.get("tool_calls"):
|
||||
raise HTTPException(status_code=400, detail="tool calls are not allowed")
|
||||
return inputs
|
||||
|
||||
|
||||
def _guardrail_g(
|
||||
style: type[_CountedGuardrail] = ApplyGuardrailG, default_on: bool = False, run_in_parallel: bool = False
|
||||
) -> _CountedGuardrail:
|
||||
return style(
|
||||
guardrail_name="G",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=default_on,
|
||||
run_in_parallel=run_in_parallel,
|
||||
)
|
||||
|
||||
|
||||
def _known_gap(action_item: str) -> pytest.MarkDecorator:
|
||||
return pytest.mark.xfail(
|
||||
raises=AssertionError,
|
||||
strict=True,
|
||||
reason=f"broken before #40154, out of scope for LIT-10481, tracked as RCA {action_item}",
|
||||
)
|
||||
|
||||
|
||||
def _deployment(model_name: str, guardrails: tuple[str, ...], api_base: str) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_key": "sk-test",
|
||||
"api_base": api_base,
|
||||
**({"guardrails": list(guardrails)} if guardrails else {}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _router(primary: tuple[str, ...], fallback: tuple[str, ...] | None = None, cached: bool = False) -> Router:
|
||||
if fallback is None:
|
||||
return Router(model_list=[_deployment("gpt", primary, _OPENAI_API_BASE)], cache_responses=cached)
|
||||
return Router(
|
||||
model_list=[
|
||||
_deployment("gpt", primary, _FAILING_API_BASE),
|
||||
_deployment("gpt-backup", fallback, _OPENAI_API_BASE),
|
||||
],
|
||||
fallbacks=[{"gpt": ["gpt-backup"]}],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
def _json_or_sse(json_reply: Mapping[str, object], sse_reply: str) -> Callable[[httpx.Request], httpx.Response]:
|
||||
def reply(request: httpx.Request) -> httpx.Response:
|
||||
if _JSON_OBJECT.validate_json(request.content).get("stream"):
|
||||
return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_reply.encode())
|
||||
return httpx.Response(200, json=dict(json_reply))
|
||||
|
||||
return reply
|
||||
|
||||
|
||||
async def _drain(stream: AsyncIterable[object]) -> str:
|
||||
return "".join([str(chunk) async for chunk in stream])
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def provider(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> Iterator[respx.MockRouter]:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
ProxyLogging._callback_capabilities_cache.clear() # pyright: ignore[reportPrivateUsage] # keyed on id(), which a collected guardrail from an earlier test can reuse
|
||||
for endpoint in ("chat/completions", "responses"):
|
||||
respx_mock.post(f"{_FAILING_API_BASE}/{endpoint}").mock(
|
||||
return_value=httpx.Response(500, json={"error": {"message": "primary down"}})
|
||||
)
|
||||
respx_mock.post(f"{_OPENAI_API_BASE}/chat/completions").mock(side_effect=_json_or_sse(_CHAT_REPLY, _CHAT_SSE))
|
||||
respx_mock.post(f"{_OPENAI_API_BASE}/responses").mock(side_effect=_json_or_sse(_RESPONSES_REPLY, _RESPONSES_SSE))
|
||||
yield respx_mock
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProxyCase:
|
||||
route: Route = "acompletion"
|
||||
stream: bool = False
|
||||
primary: tuple[str, ...] = ("G",)
|
||||
fallback: tuple[str, ...] | None = None
|
||||
request: tuple[str, ...] = ()
|
||||
default_on: bool = False
|
||||
style: type[_CountedGuardrail] = ApplyGuardrailG
|
||||
cached: bool = False
|
||||
client_sees: str = REPLY
|
||||
checks: tuple[str, ...] = (REPLY,)
|
||||
|
||||
|
||||
PROXY_CASES: Final = (
|
||||
pytest.param(ProxyCase(), id="A-chat"),
|
||||
pytest.param(ProxyCase(route="aresponses"), id="A-responses"),
|
||||
pytest.param(ProxyCase(route="anthropic_messages"), id="A-messages"),
|
||||
pytest.param(ProxyCase(request=("G",)), id="A-chat-model-and-request"),
|
||||
pytest.param(ProxyCase(style=MaskingG, client_sees=MASKED), id="A-chat-g-masks-its-own-output"),
|
||||
pytest.param(ProxyCase(stream=True), id="B-chat-stream"),
|
||||
pytest.param(ProxyCase(route="aresponses", stream=True), id="B-responses-stream"),
|
||||
pytest.param(ProxyCase(route="anthropic_messages", stream=True), id="B-messages-stream"),
|
||||
pytest.param(ProxyCase(primary=(), fallback=("G",)), id="E-only-on-fallback"),
|
||||
pytest.param(
|
||||
ProxyCase(request=("rewriter",), client_sees=REWRITTEN, checks=(REPLY, REWRITTEN)),
|
||||
id="G-rewritten-after-check",
|
||||
),
|
||||
pytest.param(ProxyCase(fallback=("G",)), id="H-on-both-deployments"),
|
||||
pytest.param(ProxyCase(primary=(), request=("G",)), id="I-request-chat"),
|
||||
pytest.param(ProxyCase(route="aresponses", primary=(), request=("G",)), id="I-request-responses"),
|
||||
pytest.param(ProxyCase(style=OwnHookG), id="J-A-chat-own-hook", marks=_known_gap("AI-5")),
|
||||
pytest.param(ProxyCase(primary=(), fallback=("G",), style=OwnHookG), id="E-only-on-fallback-own-hook"),
|
||||
pytest.param(ProxyCase(primary=(), request=("G",), style=OwnHookG), id="I-request-chat-own-hook"),
|
||||
pytest.param(
|
||||
ProxyCase(primary=(), request=("G",), style=InheritedApplyGuardrailG),
|
||||
id="K-I-request-chat-inherited-apply-guardrail",
|
||||
marks=_known_gap("AI-2"),
|
||||
),
|
||||
pytest.param(ProxyCase(primary=("other",), default_on=True), id="L-default-on"),
|
||||
)
|
||||
|
||||
_ROUTE_PATHS: Final[Mapping[Route, str]] = MappingProxyType(
|
||||
{
|
||||
"acompletion": "/v1/chat/completions",
|
||||
"aresponses": "/v1/responses",
|
||||
"anthropic_messages": "/v1/messages",
|
||||
}
|
||||
)
|
||||
_ROUTE_BODIES: Final[Mapping[Route, Mapping[str, object]]] = MappingProxyType(
|
||||
{
|
||||
"acompletion": {"messages": [{"role": "user", "content": "hello"}]},
|
||||
"aresponses": {"input": "hello"},
|
||||
"anthropic_messages": {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 32},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _proxy_reply(case: ProxyCase) -> str:
|
||||
from litellm.proxy.proxy_server import ProxyConfig, select_data_generator
|
||||
|
||||
path: Final = _ROUTE_PATHS[case.route]
|
||||
router: Final = _router(case.primary, case.fallback, case.cached)
|
||||
response: Final = await ProxyBaseLLMRequestProcessing(
|
||||
data={
|
||||
"model": "gpt",
|
||||
**_ROUTE_BODIES[case.route],
|
||||
**({"stream": True} if case.stream else {}),
|
||||
**({"guardrails": list(case.request)} if case.request else {}),
|
||||
}
|
||||
).base_process_llm_request(
|
||||
request=Request({"type": "http", "method": "POST", "path": path, "headers": [], "query_string": b""}),
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=path),
|
||||
route_type=case.route,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
|
||||
general_settings={},
|
||||
proxy_config=ProxyConfig(),
|
||||
select_data_generator=select_data_generator,
|
||||
llm_router=router,
|
||||
is_streaming_request=case.stream,
|
||||
)
|
||||
text: Final = (
|
||||
await _drain(response.body_iterator)
|
||||
if isinstance(response, StreamingResponse)
|
||||
else json.dumps(response.model_dump() if isinstance(response, BaseModel) else response)
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
return text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", PROXY_CASES)
|
||||
async def test_a_proxy_request_checks_each_distinct_response_once(
|
||||
monkeypatch: pytest.MonkeyPatch, case: ProxyCase
|
||||
) -> None:
|
||||
guardrail: Final = _guardrail_g(case.style, default_on=case.default_on)
|
||||
rewriter: Final = RewritingGuardrail(guardrail_name="rewriter", event_hook=GuardrailEventHooks.post_call)
|
||||
monkeypatch.setattr(litellm, "callbacks", [rewriter, guardrail])
|
||||
|
||||
reply: Final = await _proxy_reply(case)
|
||||
|
||||
assert case.client_sees in reply
|
||||
assert sorted(guardrail.checked) == sorted(case.checks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["acompletion", "aresponses"])
|
||||
async def test_a_cache_hit_is_checked_like_a_fresh_response(
|
||||
monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter, route: Route
|
||||
) -> None:
|
||||
guardrail: Final = _guardrail_g()
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
|
||||
case: Final = ProxyCase(route=route, cached=True)
|
||||
replies: Final = (await _proxy_reply(case), await _proxy_reply(case))
|
||||
|
||||
assert provider.calls.call_count == 1, "the second request must be served from the cache"
|
||||
assert all(REPLY in reply for reply in replies)
|
||||
assert guardrail.checked == (REPLY, REPLY)
|
||||
|
||||
|
||||
async def test_a_blocking_verdict_blocks_the_proxy_response_after_one_check(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
guardrail: Final = _guardrail_g(BlockingG)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
|
||||
with pytest.raises(HTTPException, check=lambda e: e.detail == "blocked by G"):
|
||||
await _proxy_reply(ProxyCase())
|
||||
|
||||
assert guardrail.checked == (REPLY,)
|
||||
|
||||
|
||||
SdkCaller = Literal["sdk-call", "sdk-router"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SdkCase:
|
||||
caller: SdkCaller = "sdk-call"
|
||||
stream: bool = False
|
||||
style: type[_CountedGuardrail] = ApplyGuardrailG
|
||||
|
||||
|
||||
SDK_CASES: Final = (
|
||||
pytest.param(SdkCase(), id="C-sdk-call"),
|
||||
pytest.param(SdkCase(caller="sdk-router"), id="C-sdk-router"),
|
||||
pytest.param(SdkCase(style=OwnHookG), id="C-sdk-call-own-hook"),
|
||||
pytest.param(SdkCase(stream=True), id="D-sdk-call-stream", marks=_known_gap("AI-4")),
|
||||
)
|
||||
|
||||
|
||||
async def _sdk_reply(case: SdkCase) -> str:
|
||||
messages: Final = [{"role": "user", "content": "hello"}]
|
||||
response: Final = await (
|
||||
_router(("G",)).acompletion(model="gpt", messages=messages, stream=case.stream)
|
||||
if case.caller == "sdk-router"
|
||||
else litellm.acompletion(
|
||||
model="openai/gpt-4o", api_key="sk-test", guardrails=["G"], messages=messages, stream=case.stream
|
||||
)
|
||||
)
|
||||
text: Final = await _drain(response) if isinstance(response, AsyncIterable) else str(response)
|
||||
await asyncio.sleep(0)
|
||||
return text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", SDK_CASES)
|
||||
async def test_an_sdk_call_checks_its_response_once(monkeypatch: pytest.MonkeyPatch, case: SdkCase) -> None:
|
||||
guardrail: Final = _guardrail_g(case.style)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
|
||||
reply: Final = await _sdk_reply(case)
|
||||
|
||||
assert REPLY in reply
|
||||
assert guardrail.checked == (REPLY,)
|
||||
|
||||
|
||||
async def test_mcp_auto_execute_checks_the_tool_call_round_before_running_its_tools(
|
||||
monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm, "callbacks", [ToolCallBlockingGuardrail(guardrail_name="G", event_hook=GuardrailEventHooks.post_call)]
|
||||
)
|
||||
provider.post(f"{_MCP_API_BASE}/responses").mock(
|
||||
side_effect=[httpx.Response(200, json=_TOOL_CALL_ROUND), httpx.Response(200, json=_RESPONSES_REPLY)]
|
||||
)
|
||||
echo_tool: Final = MCPTool.model_validate(
|
||||
{"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}},
|
||||
by_name=False,
|
||||
)
|
||||
run_tools: Final = AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}])
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"process_mcp_tools_without_openai_transform",
|
||||
AsyncMock(return_value=([echo_tool], {"echo": "demo"})),
|
||||
),
|
||||
patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", run_tools),
|
||||
pytest.raises(HTTPException, check=lambda e: e.detail == "tool calls are not allowed"),
|
||||
):
|
||||
await litellm.aresponses(
|
||||
model="openai/gpt-4o",
|
||||
input="Call echo with PING",
|
||||
tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}],
|
||||
api_key="sk-test",
|
||||
api_base=_MCP_API_BASE,
|
||||
guardrails=["G"],
|
||||
proxy_server_request={"url": "http://localhost:4000/v1/responses", "method": "POST"},
|
||||
litellm_metadata={"user_api_key_request_route": "/v1/responses"},
|
||||
)
|
||||
|
||||
run_tools.assert_not_awaited()
|
||||
|
||||
|
||||
def _same_response(reply: ModelResponse) -> ModelResponse:
|
||||
return reply
|
||||
|
||||
|
||||
def _model_renamed(reply: ModelResponse) -> ModelResponse:
|
||||
return reply.model_copy(update={"model": "public-alias"})
|
||||
|
||||
|
||||
def _rewritten(_: ModelResponse) -> ModelResponse:
|
||||
return ModelResponse(choices=[{"message": {"role": "assistant", "content": REWRITTEN}}])
|
||||
|
||||
|
||||
PROXY_LOOP_CASES: Final = (
|
||||
pytest.param(_same_response, (REPLY,), id="A-same-response"),
|
||||
pytest.param(_model_renamed, (REPLY,), id="A-model-renamed"),
|
||||
pytest.param(_rewritten, (REPLY, REWRITTEN), id="G-rewritten"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"])
|
||||
@pytest.mark.parametrize(("proxy_sees", "checks"), PROXY_LOOP_CASES)
|
||||
async def test_each_proxy_post_call_loop_skips_only_output_the_deployment_checked(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
parallel: bool,
|
||||
proxy_sees: Callable[[ModelResponse], ModelResponse],
|
||||
checks: tuple[str, ...],
|
||||
) -> None:
|
||||
guardrail: Final = _guardrail_g(run_in_parallel=parallel)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
data: Final = {
|
||||
"model": "gpt",
|
||||
"proxy_server_request": {"url": "http://localhost:4000/v1/chat/completions"},
|
||||
"metadata": {"guardrails": ["G"]},
|
||||
}
|
||||
reply: Final = ModelResponse(model="provider-model", choices=[{"message": {"role": "assistant", "content": REPLY}}])
|
||||
|
||||
await guardrail.async_post_call_success_deployment_hook(
|
||||
request_data={**data, "guardrails": ["G"]}, response=reply, call_type=CallTypes.acompletion
|
||||
)
|
||||
await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook(
|
||||
data=data, response=proxy_sees(reply), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert sorted(guardrail.checked) == sorted(checks)
|
||||
|
||||
|
||||
def _ocr_reply(markdown: str) -> OCRResponse:
|
||||
return OCRResponse.model_validate(
|
||||
{"pages": [{"index": 0, "markdown": markdown, "images": [], "dimensions": None}], "model": "mistral-ocr-latest"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("proxy_sees", "checks"),
|
||||
[
|
||||
pytest.param(REPLY, (REPLY,), id="A-ocr-same-pages"),
|
||||
pytest.param(REWRITTEN, (REPLY, REWRITTEN), id="G-ocr-rewritten-pages"),
|
||||
],
|
||||
)
|
||||
async def test_an_ocr_response_is_skipped_only_when_its_pages_were_checked(
|
||||
monkeypatch: pytest.MonkeyPatch, proxy_sees: str, checks: tuple[str, ...]
|
||||
) -> None:
|
||||
guardrail: Final = _guardrail_g()
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
data: Final = {
|
||||
"model": "ocr",
|
||||
"document": {"type": "document_url", "document_url": "https://example.com/doc.pdf"},
|
||||
"proxy_server_request": {"url": "http://localhost:4000/v1/ocr"},
|
||||
"metadata": {"guardrails": ["G"], "user_api_key_request_route": "/v1/ocr"},
|
||||
}
|
||||
|
||||
await guardrail.async_post_call_success_deployment_hook(
|
||||
request_data={**data, "guardrails": ["G"]}, response=_ocr_reply(REPLY), call_type=CallTypes.aocr
|
||||
)
|
||||
await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook(
|
||||
data=data, response=_ocr_reply(proxy_sees), user_api_key_dict=UserAPIKeyAuth(request_route="/v1/ocr")
|
||||
)
|
||||
|
||||
assert sorted(guardrail.checked) == sorted(checks)
|
||||
|
|
@ -699,6 +699,7 @@ def _moderation_guardrail() -> MagicMock:
|
|||
cb.should_run_guardrail = MagicMock(return_value=True)
|
||||
cb.async_moderation_hook = AsyncMock(return_value=None)
|
||||
cb.async_post_call_success_hook = AsyncMock(return_value=None)
|
||||
cb.post_call_already_checked = MagicMock(return_value=False)
|
||||
cb.run_in_parallel = False
|
||||
return cb
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue