mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): keep pass-through post-call parity with what main and native run
- A dict in the client's metadata.guardrails names guardrails by its keys, so a guarded key keeps them instead of dropping them - The inherited post-call pass reads key and team global-guardrail opt-outs, as native does - Router model guardrails no longer run on pass-through traffic that only shares the model name - merged_guardrail_entries walks the entries once, so a large list costs a fraction of before
This commit is contained in:
yucheng
parent
622968c544
commit
ec52601580
8 changed files with 314 additions and 57 deletions
|
|
@ -5,7 +5,7 @@ import os
|
|||
import secrets
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import chain, repeat
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
||||
|
|
@ -211,55 +211,32 @@ def guardrail_request_data_with_streaming(
|
|||
}
|
||||
|
||||
|
||||
_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]:
|
||||
"""Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over
|
||||
any plain name in either order, else its first plain name stays"""
|
||||
entries: Final = tuple(chain.from_iterable(entry_lists))
|
||||
settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries)
|
||||
first_plain_index: Final = _first_indexes(
|
||||
(entry, index) for index, entry in enumerate(entries) if isinstance(entry, str)
|
||||
named_with_settings: Final = frozenset(
|
||||
chain.from_iterable(guardrail_entry_names(entry) for entry in entries if isinstance(entry, dict))
|
||||
)
|
||||
first_settings_index: Final = _first_indexes(
|
||||
chain.from_iterable(zip(names, repeat(index)) for index, names in enumerate(settings_names))
|
||||
)
|
||||
return [
|
||||
entry
|
||||
for index, entry in enumerate(entries)
|
||||
if not _superseded_guardrail_entry(entry, index, settings_names[index], first_plain_index, first_settings_index)
|
||||
]
|
||||
seen: Final[set[str]] = set() # mutable-ok: local to this walk
|
||||
merged: Final[list[object]] = [] # mutable-ok: local to this walk
|
||||
for entry in entries:
|
||||
names = guardrail_entry_names(entry)
|
||||
if (names and seen.issuperset(names)) or (isinstance(entry, str) and entry in named_with_settings):
|
||||
continue
|
||||
merged.append(entry)
|
||||
seen.update(names)
|
||||
return merged
|
||||
|
||||
|
||||
def guardrail_entry_names(entry: object) -> tuple[str, ...]:
|
||||
if isinstance(entry, str):
|
||||
return (entry,)
|
||||
if isinstance(entry, dict):
|
||||
return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry))
|
||||
return tuple(cast("Mapping[str, object]", entry)) # cast-ok: a guardrail settings entry is keyed by name
|
||||
return ()
|
||||
|
||||
|
||||
def _settings_entry_names(entry: object) -> frozenset[str]:
|
||||
return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset()
|
||||
|
||||
|
||||
def _first_indexes(pairs: Iterable[tuple[str, int]]) -> Mapping[str, int]:
|
||||
return MappingProxyType(dict(reversed(tuple(pairs))))
|
||||
|
||||
|
||||
def _superseded_guardrail_entry(
|
||||
entry: object,
|
||||
index: int,
|
||||
settings_names: frozenset[str],
|
||||
first_plain_index: Mapping[str, int],
|
||||
first_settings_index: Mapping[str, int],
|
||||
) -> bool:
|
||||
if isinstance(entry, str):
|
||||
return entry in first_settings_index or first_plain_index[entry] < index
|
||||
return bool(settings_names) and all(first_settings_index[name] < index for name in settings_names)
|
||||
|
||||
|
||||
def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool:
|
||||
if event_type in _REALTIME_STREAMING_HOOKS:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ from litellm.proxy.common_request_processing import (
|
|||
open_sse_before_first_byte,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import strip_callback_config
|
||||
from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id
|
||||
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -897,13 +898,52 @@ async def _inherit_auth_guardrails(
|
|||
)
|
||||
metadata: Final = parsed_body["metadata"]
|
||||
merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
|
||||
if not isinstance(merged, list):
|
||||
return None
|
||||
entries: Final = merged_guardrail_entries(
|
||||
_GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else (),
|
||||
_GUARDRAIL_ENTRIES.validate_python(client_named) if isinstance(client_named, list) else (),
|
||||
_GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named)
|
||||
)
|
||||
if entries and isinstance(metadata, dict):
|
||||
metadata["guardrails"] = entries
|
||||
return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None
|
||||
admin_scoped_body: Final = {
|
||||
**parsed_body,
|
||||
"metadata": {**(as_str_mapping(metadata) or {}), **_admin_guardrail_metadata(user_api_key_dict)},
|
||||
}
|
||||
return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None
|
||||
|
||||
|
||||
def _client_named_guardrail_entries(client_named: object) -> Sequence[object]:
|
||||
"""The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read"""
|
||||
if isinstance(client_named, list):
|
||||
return _GUARDRAIL_ENTRIES.validate_python(client_named)
|
||||
return tuple(as_str_mapping(client_named) or ())
|
||||
|
||||
|
||||
def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]:
|
||||
"""The key and team metadata ``should_run_guardrail`` reads a global guardrail opt-out from, as the native routes set it"""
|
||||
return MappingProxyType(
|
||||
{
|
||||
"user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata),
|
||||
"user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _post_call_hook_data(
|
||||
parsed_body: Mapping[str, object] | None,
|
||||
guardrails: Sequence[object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
admin_metadata: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload
|
||||
"""The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails,
|
||||
which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a
|
||||
pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry"""
|
||||
body: Final = parsed_body or {}
|
||||
return {
|
||||
**body,
|
||||
"metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata},
|
||||
"litellm_logging_obj": logging_obj,
|
||||
}
|
||||
|
||||
|
||||
def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool:
|
||||
|
|
@ -1987,26 +2027,18 @@ async def pass_through_request(
|
|||
)
|
||||
|
||||
if response.status_code < 400 and response_body is not None and post_call_guardrails:
|
||||
# Build an enriched data dict: _parsed_body has been stripped of
|
||||
# `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint,
|
||||
# so we re-attach the configured guardrails here so should_run_guardrail
|
||||
# sees them.
|
||||
hook_data: Final = dict(_parsed_body or {})
|
||||
existing_metadata = hook_data.get("metadata")
|
||||
if not isinstance(existing_metadata, dict):
|
||||
existing_metadata = {}
|
||||
hook_data["metadata"] = {
|
||||
**existing_metadata,
|
||||
"guardrails": post_call_guardrails,
|
||||
}
|
||||
# The kwargs split popped the logging object off _parsed_body; unified (apply_guardrail) guardrails
|
||||
# resolve a pass-through response's call type from it when the route has no API_ROUTE_TO_CALL_TYPES entry.
|
||||
hook_data["litellm_logging_obj"] = logging_obj
|
||||
hook_data: Final = _post_call_hook_data(
|
||||
cast(dict[str, object] | None, _parsed_body), # cast-ok: json
|
||||
post_call_guardrails,
|
||||
logging_obj,
|
||||
_admin_guardrail_metadata(user_api_key_dict) if inherit_auth_guardrails else MappingProxyType({}),
|
||||
)
|
||||
post_call_guardrail_data = hook_data
|
||||
response_body = await proxy_logging_obj.post_call_success_hook(
|
||||
data=hook_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response_body,
|
||||
merge_model_level_guardrails=not inherit_auth_guardrails,
|
||||
)
|
||||
if isinstance(response_body, dict):
|
||||
content = json.dumps(response_body).encode("utf-8")
|
||||
|
|
|
|||
|
|
@ -4343,6 +4343,7 @@ class ProxyLogging:
|
|||
data: dict,
|
||||
response: LLMResponseTypes,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
merge_model_level_guardrails: bool = True,
|
||||
):
|
||||
"""
|
||||
Allow user to modify outgoing data
|
||||
|
|
@ -4352,6 +4353,9 @@ class ProxyLogging:
|
|||
2. /embeddings
|
||||
3. /image/generation
|
||||
4. /files
|
||||
|
||||
``merge_model_level_guardrails`` is off for a request that never reaches a router deployment, so the guardrails
|
||||
of a model that merely shares the body's ``model`` name don't run on it.
|
||||
"""
|
||||
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
|
@ -4369,7 +4373,11 @@ class ProxyLogging:
|
|||
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
|
||||
try:
|
||||
# Merge model-level guardrails before checking which guardrails to run
|
||||
guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
|
||||
guardrail_data: Final = (
|
||||
check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
|
||||
if merge_model_level_guardrails
|
||||
else data
|
||||
)
|
||||
|
||||
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
|
||||
callback
|
||||
|
|
|
|||
|
|
@ -21,7 +21,9 @@ GUARDRAIL: Final = "passthrough-auth-guardrail"
|
|||
GLOBAL_POST_CALL: Final = "passthrough-global-post-call"
|
||||
KEY_POST_CALL: Final = "passthrough-key-post-call"
|
||||
GENERIC_POST_CALL: Final = "passthrough-generic-post-call"
|
||||
MODEL_POST_CALL: Final = "passthrough-model-post-call"
|
||||
GUARDED_MODEL: Final = "passthrough-guarded-model"
|
||||
POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model"
|
||||
BLOCK_MARKER: Final = "block-this-prompt"
|
||||
OPENAI_ROUTE: Final = "/openai/v1/chat/completions"
|
||||
ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
|
||||
|
|
@ -104,6 +106,7 @@ class AuthGuardrailRig:
|
|||
post_call_vendor: Wire
|
||||
key_post_call_vendor: Wire
|
||||
generic_post_call_vendor: Wire
|
||||
model_post_call_vendor: Wire
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
||||
|
|
@ -115,6 +118,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean))
|
||||
key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged))
|
||||
generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
|
||||
upstreams: Final = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)),
|
||||
|
|
@ -149,6 +153,14 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"guardrails": [GUARDRAIL],
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": POST_CALL_GUARDED_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.6",
|
||||
"api_key": "synthetic-openai-key",
|
||||
"guardrails": [MODEL_POST_CALL],
|
||||
},
|
||||
},
|
||||
],
|
||||
"guardrails": [
|
||||
{
|
||||
|
|
@ -181,6 +193,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": MODEL_POST_CALL,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"api_base": model_post_call_vendor.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": GLOBAL_POST_CALL,
|
||||
"litellm_params": {
|
||||
|
|
@ -205,7 +227,13 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
|
|||
}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield AuthGuardrailRig(
|
||||
owned, vendor, post_call_vendor, key_post_call_vendor, generic_post_call_vendor, upstreams
|
||||
owned,
|
||||
vendor,
|
||||
post_call_vendor,
|
||||
key_post_call_vendor,
|
||||
generic_post_call_vendor,
|
||||
model_post_call_vendor,
|
||||
upstreams,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -230,6 +258,7 @@ def _drain(rig: AuthGuardrailRig) -> None:
|
|||
rig.post_call_vendor.drain()
|
||||
rig.key_post_call_vendor.drain()
|
||||
rig.generic_post_call_vendor.drain()
|
||||
rig.model_post_call_vendor.drain()
|
||||
for upstream in rig.upstreams.values():
|
||||
upstream.drain()
|
||||
|
||||
|
|
@ -381,3 +410,77 @@ def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig
|
|||
assert response.status_code == 400, response.text
|
||||
assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True]
|
||||
assert rig.upstreams[route].drain() == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("settings", [{}, True])
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
|
||||
def test_guardrail_the_client_names_in_the_dict_form_still_runs_for_a_guarded_key(
|
||||
rig: AuthGuardrailRig, route: str, settings: JsonValue
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]})
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(
|
||||
gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": {GUARDRAIL: settings}}}, key
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True]
|
||||
assert rig.upstreams[route].drain() == ()
|
||||
|
||||
|
||||
def _key_with_opt_out(scenario: Scenario, opt_out: str) -> str:
|
||||
guarded: Final[dict[str, JsonValue]] = {"guardrails": [GENERIC_POST_CALL]}
|
||||
match opt_out:
|
||||
case "key_disable":
|
||||
return scenario.key(metadata=guarded, disable_global_guardrails=True)
|
||||
case "key_opted_out":
|
||||
return scenario.key(metadata={**guarded, "opted_out_global_guardrails": [GLOBAL_POST_CALL]})
|
||||
case "team_disable":
|
||||
return scenario.key(team_id=scenario.team(disable_global_guardrails=True), metadata=guarded)
|
||||
case "team_opted_out":
|
||||
return scenario.key(
|
||||
team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_POST_CALL]}), metadata=guarded
|
||||
)
|
||||
case _:
|
||||
return scenario.key(metadata=guarded)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("opt_out", "global_vendor_calls"),
|
||||
[("key_disable", 0), ("key_opted_out", 0), ("team_disable", 0), ("team_opted_out", 0), ("none", 1)],
|
||||
)
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, "/v1/chat/completions"])
|
||||
def test_post_call_pass_honors_a_global_guardrail_opt_out_as_the_native_route_does(
|
||||
rig: AuthGuardrailRig, route: str, opt_out: str, global_vendor_calls: int
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = _key_with_opt_out(scenario, opt_out)
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(gateway, route, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(rig.generic_post_call_vendor.drain()) == 1
|
||||
assert len(rig.post_call_vendor.drain()) == global_vendor_calls
|
||||
|
||||
|
||||
def test_model_guardrail_does_not_run_on_pass_through_traffic_that_names_the_model(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]})
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(
|
||||
gateway,
|
||||
OPENAI_ROUTE,
|
||||
{**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": POST_CALL_GUARDED_MODEL},
|
||||
key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(rig.generic_post_call_vendor.drain()) == 1
|
||||
assert rig.model_post_call_vendor.drain() == ()
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ import datetime as dt
|
|||
import json
|
||||
import pickle
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
|
|
@ -3802,6 +3804,24 @@ def test_default_on_guardrail() -> None:
|
|||
)
|
||||
|
||||
|
||||
def test_merged_guardrail_entries_stays_close_to_a_plain_dedup_on_a_large_list():
|
||||
from litellm.integrations.custom_guardrail import merged_guardrail_entries
|
||||
|
||||
names: Final = [f"g{index}" for index in range(300_000)]
|
||||
|
||||
def _best_cpu_seconds(run: Callable[[], object]) -> float:
|
||||
def _timed() -> float:
|
||||
started: Final = time.process_time()
|
||||
run()
|
||||
return time.process_time() - started
|
||||
|
||||
return min(_timed() for _ in range(3))
|
||||
|
||||
assert _best_cpu_seconds(lambda: merged_guardrail_entries(names)) < 25 * _best_cpu_seconds(
|
||||
lambda: list(dict.fromkeys(names))
|
||||
)
|
||||
|
||||
|
||||
def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings():
|
||||
from litellm.integrations.custom_guardrail import merged_guardrail_entries
|
||||
|
||||
|
|
|
|||
|
|
@ -13368,6 +13368,79 @@ class TestInheritAuthGuardrails:
|
|||
merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"])
|
||||
assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("settings", [{}, True])
|
||||
async def test_guardrail_a_client_names_in_the_dict_form_stays_after_the_keys(self, settings: object):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, _ = await self._run(
|
||||
{**self._chat_body(), "metadata": {"guardrails": {"client-guard": settings}}},
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(self._guardrail("key-guard", GuardrailEventHooks.pre_call),),
|
||||
)
|
||||
|
||||
assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self):
|
||||
pre_call, post_call = await self._run(
|
||||
{**self._chat_body(), "metadata": {"guardrails": {"client-guard": {}}}},
|
||||
inherit=True,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
guardrails=(self._guardrail("client-guard", GuardrailEventHooks.post_call),),
|
||||
)
|
||||
|
||||
assert pre_call[0]["metadata"]["guardrails"] == {"client-guard": {}}
|
||||
post_call.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pass_reads_opt_outs_from_the_key_and_team_not_the_body(self):
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
metadata={"guardrails": ["key-guard"], "disable_global_guardrails": True},
|
||||
team_metadata={"opted_out_global_guardrails": ["global-post"]},
|
||||
)
|
||||
forged: Final = {"user_api_key_metadata": {"disable_global_guardrails": False}}
|
||||
|
||||
_, post_call = await self._run(
|
||||
{**self._chat_body(), "metadata": forged},
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),),
|
||||
)
|
||||
|
||||
metadata: Final = post_call.call_args.kwargs["data"]["metadata"]
|
||||
assert metadata["user_api_key_metadata"]["disable_global_guardrails"] is True
|
||||
assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inherited_post_call_pass_skips_model_level_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
_, post_call = await self._run(
|
||||
self._chat_body(),
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),),
|
||||
)
|
||||
|
||||
assert post_call.call_args.kwargs["merge_model_level_guardrails"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_named_global_guardrail_the_key_opted_out_of_does_not_start_the_post_call_pass(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"opted_out_global_guardrails": ["global-post"]})
|
||||
|
||||
response, _, post_call = await self._run_with_response(
|
||||
{**self._chat_body(), "metadata": {"guardrails": ["global-post"]}},
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),),
|
||||
)
|
||||
|
||||
post_call.assert_not_awaited()
|
||||
assert response.body == b'{\n "id":"chatcmpl-1"\n}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_only_guardrails_relay_the_response_untouched(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
|
|
|||
|
|
@ -130,6 +130,29 @@ class TestPassthroughPostCallGuardrails:
|
|||
call_kwargs = mock_proxy_logging.post_call_success_hook.call_args
|
||||
assert call_kwargs.kwargs["response"] == _GEMINI_RESPONSE
|
||||
|
||||
@patch(_COLLECT, return_value=["rubrik"])
|
||||
async def test_configured_endpoint_keeps_model_guardrails_and_reads_no_key_opt_outs(
|
||||
self,
|
||||
mock_collect,
|
||||
):
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=_GEMINI_RESPONSE)
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
with _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)):
|
||||
await pass_through_request(
|
||||
request=_make_mock_request(),
|
||||
target="https://example.com/v1/generateContent",
|
||||
custom_headers={"Content-Type": "application/json"},
|
||||
user_api_key_dict=_make_user_api_key_dict(),
|
||||
stream=False,
|
||||
)
|
||||
|
||||
call_kwargs = mock_proxy_logging.post_call_success_hook.call_args.kwargs
|
||||
assert call_kwargs["merge_model_level_guardrails"] is True
|
||||
assert call_kwargs["data"]["metadata"].keys() == {"guardrails"}
|
||||
|
||||
@patch(_COLLECT, return_value=[])
|
||||
async def test_post_call_success_hook_skipped_when_no_guardrails(
|
||||
self,
|
||||
|
|
@ -216,7 +239,7 @@ class TestPassthroughPostCallGuardrails:
|
|||
block path dropping it."""
|
||||
mock_response = _make_httpx_response(_GEMINI_RESPONSE)
|
||||
|
||||
def _block(*, data, user_api_key_dict, response):
|
||||
def _block(*, data, user_api_key_dict, response, merge_model_level_guardrails):
|
||||
metadata = data.setdefault("metadata", {})
|
||||
metadata.setdefault("standard_logging_guardrail_information", []).append(
|
||||
{"guardrail_name": "rubrik", "guardrail_status": "guardrail_intervened"}
|
||||
|
|
|
|||
|
|
@ -59,6 +59,27 @@ async def test_post_call_success_hook_runs_other_callback_and_replaces_response(
|
|||
assert out == new_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("merge", "guardrails_seen"), [(True, ["key-guard", "model-guard"]), (False, ["key-guard"])])
|
||||
async def test_post_call_success_hook_merges_model_level_guardrails_unless_told_not_to(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, merge, guardrails_seen
|
||||
):
|
||||
router = MagicMock()
|
||||
router.registry.return_value.deployments.return_value = [{"litellm_params": {"guardrails": ["model-guard"]}}]
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
g = _make_guardrail(should_run=False)
|
||||
monkeypatch.setattr(litellm, "callbacks", [g])
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data={"model": "alias", "metadata": {"guardrails": ["key-guard"]}},
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
merge_model_level_guardrails=merge,
|
||||
)
|
||||
|
||||
assert g.should_run_guardrail.call_args.kwargs["data"]["metadata"]["guardrails"] == guardrails_seen
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_guardrail_should_not_run_skipped(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue