scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

Merge remote-tracking branch 'origin/litellm_passthrough_auth_guardrails' into litellm_lakera_v2_gemini_contents

Browse source
This commit is contained in:
yucheng 2026-11-05 01:14:51 -08:00
commit 36970355dd
6 changed files with 271 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/litellm_pre_call_utils.py
View file

@ -3986,6 +3986,11 @@ def add_guardrails_from_key_or_team_metadata(
_add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata
def _string_policy_names(policies: Sequence[object]) -> tuple[str, ...]:
"""The policy names in a key, team or project ``policies`` list, less an entry that isn't a name"""
return tuple(name for name in policies if isinstance(name, str))
def _add_guardrails_from_policies_in_metadata(
key_metadata: dict | None,
team_metadata: dict | None,
@ -4021,19 +4026,25 @@ def _add_guardrails_from_policies_in_metadata(
if key_metadata and "policies" in key_metadata:
if isinstance(key_metadata["policies"], list) and len(key_metadata["policies"]) > 0:
premium_user_check()
policy_names.update(key_metadata["policies"])
policy_names.update(
_string_policy_names(cast("Sequence[object]", key_metadata["policies"])) # cast-ok: checked list
)
# Add team-level policies
if team_metadata and "policies" in team_metadata:
if isinstance(team_metadata["policies"], list) and len(team_metadata["policies"]) > 0:
premium_user_check()
policy_names.update(team_metadata["policies"])
policy_names.update(
_string_policy_names(cast("Sequence[object]", team_metadata["policies"])) # cast-ok: checked list
)
# Add project-level policies
if project_metadata and "policies" in project_metadata:
if isinstance(project_metadata["policies"], list) and len(project_metadata["policies"]) > 0:
premium_user_check()
policy_names.update(project_metadata["policies"])
policy_names.update(
_string_policy_names(cast("Sequence[object]", project_metadata["policies"])) # cast-ok: checked list
)
if not policy_names:
return

102
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -209,6 +209,13 @@ _FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = (
("guardrail_config", (dict,)),
("proxy_server_request", (dict,)),
)
_PROXY_FILLED_FIELD_SHAPES: Final = (
("litellm_call_id", (str,)),
("litellm_trace_id", (str,)),
("custom_llm_provider", (str,)),
("api_base", (str,)),
("api_key", (str,)),
)
_POST_CALL_SELECTION_FIELDS: Final = (
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
@ -223,6 +230,19 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool:
return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping))
def _holds_string_headers(value: object) -> bool:
headers: Final = as_str_mapping(value)
return headers is not None and all(isinstance(header, str) for header in headers.values())
def _is_guardrail_readable_client_metadata(key: str, value: object) -> bool:
if key == "headers":
return _holds_string_headers(value)
if key == "requester_metadata":
return isinstance(value, Mapping)
return _is_mergeable_client_metadata(key, value)
class _RegisteredPassThroughRoute(TypedDict):
endpoint_id: ReadOnly[str]
path: ReadOnly[str | None]
@ -925,24 +945,43 @@ async def _inherit_auth_guardrails(
client_sent_metadata: Final = "metadata" in parsed_body
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")
client_metadata: Final = as_str_mapping(parsed_body["metadata"]) or {}
client_body_guardrails: Final = {"guardrails": parsed_body["guardrails"]} if "guardrails" in parsed_body else {}
client_metadata_guardrails: Final = (
{"guardrails": client_metadata["guardrails"]} if "guardrails" in client_metadata else {}
)
client_named: Final = client_metadata.get("guardrails")
await move_guardrails_to_metadata(
data=parsed_body,
_metadata_variable_name="metadata",
user_api_key_dict=user_api_key_dict,
)
metadata: Final = parsed_body["metadata"]
merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
entries: Final = (
merged_guardrail_entries(
_GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named)
)
if isinstance(merged, list)
else ()
)
runs_guardrails: Final = bool(entries) or bool((as_str_mapping(metadata) or {}).get("_guardrail_pipelines"))
if entries and isinstance(metadata, dict):
metadata["guardrails"] = entries
elif isinstance(merged, list) and isinstance(metadata, dict):
_restore_unmerged_guardrail_fields(
parsed_body,
cast("dict[str, object]", metadata), # cast-ok: the request body's JSON metadata object
client_body_guardrails,
client_metadata_guardrails,
)
if runs_guardrails:
_drop_client_values_guardrails_cannot_read(parsed_body)
if not metadata and not client_sent_metadata:
parsed_body.pop("metadata")
return None
merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
if not isinstance(merged, list):
if not entries:
return None
entries: Final = merged_guardrail_entries(
_GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named)
)
if entries and isinstance(metadata, dict):
metadata["guardrails"] = entries
admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict)
admin_scoped_body: Final = {
**parsed_body,
@ -951,11 +990,54 @@ async def _inherit_auth_guardrails(
return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None
def _drop_client_values_guardrails_cannot_read(
parsed_body: dict[str, object], # mutable-ok: strips the request body in place before its guardrails run
) -> None:
"""Drop what a client sent in a field the proxy fills on the native routes when the guardrails can't read it there:
a ``litellm_call_id``, ``litellm_trace_id``, ``custom_llm_provider``, ``api_base`` or ``api_key`` that isn't a
string, a ``proxy_server_request`` or ``metadata.headers`` whose headers aren't strings, a ``requester_metadata``
that isn't an object, and any ``user_api_key_*`` metadata field, which only the proxy sets. The provider receives
none of these fields."""
for field, shape in _PROXY_FILLED_FIELD_SHAPES:
if field in parsed_body and not isinstance(parsed_body[field], shape):
parsed_body.pop(field)
server_request: Final = as_str_mapping(parsed_body.get("proxy_server_request")) or {}
if "headers" in server_request and not _holds_string_headers(server_request["headers"]):
parsed_body.pop("proxy_server_request")
for bucket in (parsed_body.get(key) for key in sorted(_METADATA_KEYS)):
if isinstance(bucket, dict):
_drop_metadata_guardrails_cannot_read(cast("dict[str, object]", bucket)) # cast-ok: a JSON object
def _drop_metadata_guardrails_cannot_read(
metadata: dict[str, object], # mutable-ok: strips the request body's metadata in place
) -> None:
unreadable: Final = tuple(
field for field, value in metadata.items() if not _is_guardrail_readable_client_metadata(field, value)
)
for field in unreadable:
metadata.pop(field)
def _restore_unmerged_guardrail_fields(
parsed_body: dict[str, object], # mutable-ok: puts the client's own fields back on the live request body
metadata: dict[str, object], # mutable-ok: the live request body's metadata
client_body_guardrails: Mapping[str, object],
client_metadata_guardrails: Mapping[str, object],
) -> None:
"""Put back the client's ``guardrails`` and ``metadata.guardrails`` when the merge found no guardrail to run, so
the body keeps main's shape"""
metadata.pop("guardrails", None)
metadata.update(client_metadata_guardrails)
parsed_body.update(client_body_guardrails)
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"""
"""The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read that
it doesn't turn off with ``false``"""
if isinstance(client_named, list):
return _GUARDRAIL_ENTRIES.validate_python(client_named)
return tuple(as_str_mapping(client_named) or ())
return tuple(name for name, run in (as_str_mapping(client_named) or {}).items() if run is not False)
def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]:

7
litellm/proxy/utils.py
View file

@ -892,9 +892,8 @@ def _add_applied_guardrails_in_place(
data: dict[str, object], # mutable-ok: the names are recorded in the live request body's metadata
names: Sequence[str],
) -> None:
"""Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names"""
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
"""Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names. A record a
guardrail left that isn't a list stays as it is, so the request doesn't fail on it."""
if not names:
return
metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket
@ -903,8 +902,6 @@ def _add_applied_guardrails_in_place(
return
applied: Final = metadata["applied_guardrails"]
if not isinstance(applied, list):
for name in names:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=name)
return
entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names
present: Final = frozenset(name for name in entries if isinstance(name, str))

93
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -13372,6 +13372,99 @@ class TestInheritAuthGuardrails:
assert body["metadata"] == {}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"extra",
[
{"guardrails": []},
{"guardrails": [], "metadata": {"user": "u"}},
{"guardrails": [], "metadata": {"guardrails": []}},
{"metadata": {"guardrails": []}},
],
ids=["guardrails", "with-metadata", "with-metadata-guardrails", "metadata-guardrails-only"],
)
async def test_an_empty_guardrails_list_on_a_plain_key_keeps_the_body_as_sent(self, extra: dict[str, object]):
body: Final[dict[str, object]] = self._chat_body(**extra)
assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None
assert body == self._chat_body(**extra)
@pytest.mark.asyncio
async def test_a_guardrail_the_body_turns_off_with_false_stays_out_of_the_merge(self):
body: Final[dict[str, object]] = self._chat_body(
metadata={"guardrails": {"lak-legacy": False, "body-guard": {}}}
)
with patch("litellm.proxy.proxy_server.premium_user", True):
await _inherit_auth_guardrails(
body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True
)
assert body == {**self._chat_body(), "metadata": {"guardrails": ["key-guard", "body-guard"]}}
@staticmethod
def _client_values_the_guardrails_cant_read() -> dict[str, object]:
return {
"litellm_call_id": 5,
"litellm_trace_id": 5,
"custom_llm_provider": 5,
"api_base": 5,
"api_key": 5,
"proxy_server_request": {"headers": {"user-agent": 7}},
"metadata": {
"headers": "abc",
"requester_metadata": "abc",
"user_api_key_alias": "forged",
"user_api_key_user_id": 5,
"user": "kept",
},
"litellm_metadata": {"headers": {"x-aim-user-email": 5}, "user_api_key_alias": 5, "tags": ["kept"]},
}
@pytest.mark.asyncio
async def test_client_values_the_guardrails_cant_read_are_dropped_when_guardrails_run(self):
body: Final[dict[str, object]] = self._chat_body(**self._client_values_the_guardrails_cant_read())
with patch("litellm.proxy.proxy_server.premium_user", True):
await _inherit_auth_guardrails(
body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True
)
assert body == {
**self._chat_body(),
"metadata": {"user": "kept", "guardrails": ["key-guard"]},
"litellm_metadata": {"tags": ["kept"]},
}
@pytest.mark.asyncio
async def test_client_values_the_guardrails_can_read_are_kept_when_guardrails_run(self):
readable: Final[dict[str, object]] = {
"litellm_call_id": "call",
"litellm_trace_id": "trace",
"custom_llm_provider": "openai",
"api_base": "https://example.com",
"api_key": "sk-client",
"proxy_server_request": {"headers": {"user-agent": "ua"}},
}
client_metadata: Final[dict[str, object]] = {"headers": {"x-aim-user-email": "a@b.c"}, "requester_metadata": {}}
body: Final[dict[str, object]] = self._chat_body(**readable, metadata=dict(client_metadata))
with patch("litellm.proxy.proxy_server.premium_user", True):
await _inherit_auth_guardrails(
body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True
)
assert body == {**self._chat_body(**readable), "metadata": {**client_metadata, "guardrails": ["key-guard"]}}
@pytest.mark.asyncio
async def test_client_values_stay_as_sent_when_no_guardrail_runs(self):
body: Final[dict[str, object]] = self._chat_body(**self._client_values_the_guardrails_cant_read())
assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None
assert body == self._chat_body(**self._client_values_the_guardrails_cant_read())
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [5, "text", ["list"], None])
async def test_a_proxy_server_request_that_isnt_an_object_never_reaches_the_guardrails(self, value):

32
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -5622,6 +5622,38 @@ def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_onc
assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"]
@pytest.mark.parametrize("level", ["key", "team", "project"])
def test_add_guardrails_from_auth_metadata_skips_a_policy_name_that_is_not_a_string(level: str):
from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails
metadata: Final[dict[str, object]] = {}
data: Final[dict[str, object]] = {"model": "gpt-4", "metadata": metadata}
policy_registry = get_policy_registry()
policy_registry._policies = {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))}
policy_registry._initialized = True
policies = {"policies": [{"a": 1}, ["b"], "baseline"]}
try:
with patch("litellm.proxy.utils.premium_user_check"):
add_guardrails_from_auth_metadata(
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key",
metadata=policies if level == "key" else {},
team_metadata=policies if level == "team" else {},
project_metadata=policies if level == "project" else None,
),
data=data,
metadata_variable_name="metadata",
)
finally:
policy_registry._policies = {}
policy_registry._initialized = False
assert metadata["guardrails"] == ["pii_blocker"]
def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it():
from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata
from litellm.proxy.policy_engine.policy_registry import get_policy_registry

38
tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -949,6 +949,44 @@ async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardr
assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"]
class _RecordsAppliedGuardrails(CustomGuardrail):
def __init__(self, record: object) -> None:
super().__init__(guardrail_name="recorder", event_hook=GuardrailEventHooks.pre_call, default_on=False)
self._record: Final = record
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, dict[str, object]],
call_type: CallTypesLiteral,
) -> dict[str, dict[str, object]]:
data["metadata"]["applied_guardrails"] = self._record
return data
@pytest.mark.asyncio
@pytest.mark.parametrize("record", [5, "other", ("other",)], ids=["int", "str", "tuple"])
async def test_a_non_list_applied_guardrails_record_does_not_fail_the_request(
proxy_logging, make_user_api_key_auth, monkeypatch, record: object
):
monkeypatch.setattr(litellm, "callbacks", [_RecordsAppliedGuardrails(record)])
proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
data: Final = {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"guardrails": ["recorder", "declared-post-call"]},
}
out: Final = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(),
data=data,
call_type="completion",
)
assert out["metadata"]["applied_guardrails"] == record
@pytest.mark.asyncio
async def test_skip_guardrails_still_runs_non_guardrail_callbacks(proxy_logging, make_user_api_key_auth, monkeypatch):
accountant = _Accountant()