mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
fix(proxy): keep the guardrails a client names in metadata on guarded pass-through keys and merge guardrail lists in linear time
This commit is contained in:
parent
c0116078b9
commit
377124588e
4 changed files with 68 additions and 12 deletions
|
|
@ -5,7 +5,7 @@ import os
|
|||
import secrets
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import chain
|
||||
from itertools import chain, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
||||
|
|
@ -219,11 +219,16 @@ def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]:
|
|||
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)
|
||||
named_with_settings: Final = frozenset[str]().union(*settings_names)
|
||||
first_plain_index: Final = _first_indexes(
|
||||
(entry, index) for index, entry in enumerate(entries) if isinstance(entry, str)
|
||||
)
|
||||
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(entries, settings_names, named_with_settings, index)
|
||||
if not _superseded_guardrail_entry(entry, index, settings_names[index], first_plain_index, first_settings_index)
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -239,16 +244,20 @@ 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(
|
||||
entries: tuple[object, ...],
|
||||
settings_names: tuple[frozenset[str], ...],
|
||||
named_with_settings: frozenset[str],
|
||||
entry: object,
|
||||
index: int,
|
||||
settings_names: frozenset[str],
|
||||
first_plain_index: Mapping[str, int],
|
||||
first_settings_index: Mapping[str, int],
|
||||
) -> bool:
|
||||
entry: Final = entries[index]
|
||||
if isinstance(entry, str):
|
||||
return entry in named_with_settings or entry in entries[:index]
|
||||
return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index])
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
guardrail_entry_names,
|
||||
guardrail_request_data_with_streaming,
|
||||
merged_guardrail_entries,
|
||||
without_server_streaming_classification,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -875,19 +876,27 @@ async def _inherit_auth_guardrails(
|
|||
|
||||
The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team,
|
||||
project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook,
|
||||
and returned only when at least one of them runs post_call, so the response is relayed untouched otherwise.
|
||||
with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one
|
||||
of them runs post_call, so the response is relayed untouched otherwise.
|
||||
"""
|
||||
if configured_guardrails or not inherit or not parsed_body:
|
||||
return configured_guardrails
|
||||
if not isinstance(parsed_body.get("metadata"), dict):
|
||||
parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata
|
||||
client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
|
||||
await move_guardrails_to_metadata(
|
||||
data=parsed_body,
|
||||
_metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
|
||||
entries: Final = _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else []
|
||||
metadata: Final = parsed_body["metadata"]
|
||||
merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
|
||||
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 (),
|
||||
)
|
||||
if entries and isinstance(metadata, dict):
|
||||
metadata["guardrails"] = entries
|
||||
return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -365,3 +365,19 @@ def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardra
|
|||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
|
||||
def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig: AuthGuardrailRig, route: str) -> 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]}}, 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() == ()
|
||||
|
|
|
|||
|
|
@ -11900,6 +11900,28 @@ class TestInheritAuthGuardrails:
|
|||
assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"]
|
||||
assert "guardrails" not in pre_call[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrails_the_client_names_in_metadata_still_run_for_a_guarded_key(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, post_call = await self._run(
|
||||
self._chat_body(metadata={"guardrails": ["client-pre", "client-post"]}),
|
||||
inherit=True,
|
||||
user_api_key_dict=auth,
|
||||
guardrails=(
|
||||
self._guardrail("key-guard", GuardrailEventHooks.pre_call),
|
||||
self._guardrail("client-post", GuardrailEventHooks.post_call),
|
||||
),
|
||||
)
|
||||
|
||||
assert pre_call[0]["metadata"] == {"guardrails": ["key-guard", "client-pre", "client-post"]}
|
||||
post_call.assert_awaited_once()
|
||||
assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == [
|
||||
"key-guard",
|
||||
"client-pre",
|
||||
"client-post",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_request_body_guardrails_keep_the_key_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue