scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge 128ced5aab into 0b39e5faf3

Browse source
This commit is contained in:
AlbertNeuralTrust 2026-11-03 17:20:18 -08:00 • committed by GitHub
commit e1db7d8403
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 256 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -2337,6 +2337,68 @@ async def test_custom_code_guardrail(
)
_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _metadata_fields(value: object) -> Mapping[str, object]:
try:
return _GUARDRAIL_METADATA_ADAPTER.validate_python(value)
except ValidationError:
return {}
_CALLER_OWNED_KEY_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_end_user_id"})
_WITHHELD_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_request_route", "agent_id"})
_RESERVED_IDENTITY_FIELDS: Final[frozenset[str]] = frozenset(
{"requester_ip_address", "user_agent", "agent_id", "actor_agent_id"}
)
def _is_key_identity_field(key: str) -> bool:
"""Whether a metadata field attributes the authenticated caller rather than the request.
Pre-call processing owns the whole ``user_api_key_`` namespace, except the end user the call
is made for, which the caller owns here as it does on every LLM route.
Attribution also reaches outside that namespace, so the fields pre-call processing fills in
from the key and the connection are reserved by name as well.
"""
if key in _CALLER_OWNED_KEY_FIELDS:
return False
return key in _RESERVED_IDENTITY_FIELDS or key == "user_api_key" or key.startswith("user_api_key_")
def _guardrail_request_metadata(
caller: object, proxy: object, agent_id: str | None = None
) -> dict[str, object]: # mutable-ok: guardrails annotate the metadata they are handed
"""Caller metadata with the identity fields taken from the authenticated key.
The result is handed to the guardrail as ``request_data["metadata"]``, which providers
write their own detection details into, so it cannot be a read-only mapping.
The request route is withheld from both sides. The tool policy guardrail reads the request
itself when a route is present, instead of the tools replayed to it, so a route here turns
its checks off: the caller must not be able to set one, and this endpoint's own route
describes the evaluation rather than the call being evaluated.
The agent is taken from ``agent_id``, the one the key authenticated with, and omitted when the
key is bound to none. Pre-call processing keeps the request's own agent when the key has none,
so the processed metadata is not a trustworthy source for it either.
"""
caller_fields: Final = tuple(
(key, value) for key, value in _metadata_fields(caller).items() if not _is_key_identity_field(key)
)
key_identity: Final = tuple(
(key, value)
for key, value in _metadata_fields(proxy).items()
if _is_key_identity_field(key) and key not in _WITHHELD_FIELDS
)
bound_agent: Final = (("agent_id", agent_id),) if agent_id is not None else ()
return dict((*caller_fields, *key_identity, *bound_agent))
def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse:
return TestCustomCodeGuardrailResponse(
success=False,
@ -2503,6 +2565,16 @@ async def apply_guardrail(
response_handler: Final = GUARDRAIL_RESPONSE_HANDLERS[request.response_format]
request_data: Final = dict(response_handler.prepare_request_data(request))
caller_metadata: Final = request_data.get("metadata")
metadata: Final = _guardrail_request_metadata(
caller_metadata,
_GUARDRAIL_METADATA_ADAPTER.validate_python(data).get("metadata"),
agent_id=user_api_key_dict.invoked_agent_id or user_api_key_dict.agent_id,
)
if caller_metadata is not None or metadata:
request_data["metadata"] = metadata
else:
request_data.pop("metadata", None)
input_type: Final = response_handler.resolve_input_type(request, active_guardrail)
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
inputs={"texts": [request.text]},

197
tests/unit/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -4,7 +4,7 @@ from collections.abc import Callable
from datetime import datetime
from types import MappingProxyType
from typing import Dict, Final, List, Literal, Optional
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, Mock
import httpx
import pytest
@ -25,6 +25,7 @@ from litellm.proxy.guardrails.guardrail_endpoints import (
TestCustomCodeGuardrailRequest,
UpdateGuardrailRequest,
_get_guardrails_list_response,
_guardrail_request_metadata,
apply_guardrail,
approve_guardrail_submission,
create_guardrail,
@ -2169,27 +2170,29 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
}
def _patch_apply_guardrail_env(mocker, guardrail_result):
mock_guardrail = mocker.Mock()
def _patch_apply_guardrail_env(
mocker: MockerFixture, guardrail_result: object, processed_data: dict[str, object] | None = None
) -> Mock:
mock_guardrail: Final = mocker.Mock()
mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
mock_registry = mocker.Mock()
mock_registry: Final = mocker.Mock()
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
mock_logging_obj = mocker.Mock()
mock_logging_obj: Final = mocker.Mock()
mock_logging_obj.async_success_handler = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_processor = mocker.Mock()
mock_processor: Final = mocker.Mock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj)
)
mocker.patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
return_value=mock_processor,
)
mock_proxy_logging = mocker.Mock()
mock_proxy_logging: Final = mocker.Mock()
mock_proxy_logging.post_call_success_hook = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
mocker.patch("litellm.proxy.proxy_server.general_settings", {})
@ -2483,21 +2486,189 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
)
def test_the_authenticated_key_identity_replaces_the_caller_version() -> None:
metadata: Final = _guardrail_request_metadata(
{"forbidden_topics": ["tax"], "user_api_key_alias": "someone-else"},
{"user_api_key_alias": "billing-app", "user_api_key_user_id": "u-1"},
)
assert metadata == {
"forbidden_topics": ["tax"],
"user_api_key_alias": "billing-app",
"user_api_key_user_id": "u-1",
}
def test_the_bare_key_hash_cannot_come_from_the_caller() -> None:
metadata: Final = _guardrail_request_metadata({"user_api_key": "spoof"}, {"user_api_key": "hashed"})
assert metadata == {"user_api_key": "hashed"}
def test_the_caller_owns_the_end_user() -> None:
metadata: Final = _guardrail_request_metadata({"user_api_key_end_user_id": "customer-42"}, {})
assert metadata == {"user_api_key_end_user_id": "customer-42"}
def test_no_request_route_reaches_the_guardrail() -> None:
"""A route makes the tool policy guardrail read the request instead of the replayed tools."""
metadata: Final = _guardrail_request_metadata(
{"user_api_key_request_route": "/v1/chat/completions", "topic": "tax"},
{"user_api_key_request_route": "/apply_guardrail"},
)
assert metadata == {"topic": "tax"}
def test_an_unbound_agent_is_not_taken_from_the_request() -> None:
"""Pre-call processing keeps the request's own agent_id when the key has none, so the proxy's
metadata already carries a caller-chosen value and reserving the field is not enough."""
metadata: Final = _guardrail_request_metadata(
{"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"}
)
assert metadata == {}
def test_a_bound_agent_comes_from_the_authenticated_key() -> None:
metadata: Final = _guardrail_request_metadata(
{"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"}, agent_id="billing-agent"
)
assert metadata == {"agent_id": "billing-agent"}
def test_identity_fields_outside_the_key_namespace_are_reserved_too() -> None:
metadata: Final = _guardrail_request_metadata(
{
"requester_ip_address": "203.0.113.9",
"user_agent": "forged/1.0",
"agent_id": "someone-elses-agent",
"topic": "tax",
},
{"requester_ip_address": "198.51.100.4", "user_agent": "litellm-cli/2.3"},
agent_id="billing-agent",
)
assert metadata == {
"topic": "tax",
"requester_ip_address": "198.51.100.4",
"user_agent": "litellm-cli/2.3",
"agent_id": "billing-agent",
}
def test_a_guardrail_can_annotate_what_it_is_handed() -> None:
"""OpenAI moderation and Presidio write their detection details into this mapping."""
metadata: Final = _guardrail_request_metadata({}, {"user_api_key_alias": "billing-app"})
metadata["moderation_model"] = "omni-moderation-latest"
assert metadata == {"user_api_key_alias": "billing-app", "moderation_model": "omni-moderation-latest"}
@pytest.mark.asyncio
async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
"""Without metadata, request_data stays empty (backward-compatible)."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key(mocker):
"""Identity fields come from the proxy's own sanitized metadata, never from
the caller, so a request cannot impersonate another key or probe its policy."""
mock_guardrail = _patch_apply_guardrail_env(
mocker,
{"texts": ["ok"]},
processed_data={
"guardrail_name": "test-guardrail",
"metadata": {
"route": "/apply_guardrail",
"user_api_key_alias": "billing-app",
"user_api_key_user_id": "u-1",
},
},
)
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="hello",
metadata={
"forbidden_topics": ["tax"],
"user_api_key_alias": "someone-else",
"user_api_key_user_email": "victim@example.com",
},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"),
)
forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
assert forwarded == {
"forbidden_topics": ["tax"],
"user_api_key_alias": "billing-app",
"user_api_key_user_id": "u-1",
}
@pytest.mark.asyncio
async def test_apply_guardrail_ignores_a_caller_supplied_user_api_key(mocker):
"""The bare key field carries attribution as much as the prefixed ones, so a caller
cannot set it either."""
mock_guardrail = _patch_apply_guardrail_env(
mocker,
{"texts": ["ok"]},
processed_data={
"guardrail_name": "test-guardrail",
"metadata": {"route": "/apply_guardrail", "user_api_key": "hash-of-the-real-key"},
},
)
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="hello",
metadata={"forbidden_topics": ["tax"], "user_api_key": "hash-of-someone-else"},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"),
)
forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
assert forwarded == {"forbidden_topics": ["tax"], "user_api_key": "hash-of-the-real-key"}
@pytest.mark.asyncio
async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata(mocker):
"""A caller that sends no metadata still gets the authenticated identity
forwarded, the same shape every LLM route gives a guardrail."""
mock_guardrail = _patch_apply_guardrail_env(
mocker,
{"texts": ["ok"]},
processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}},
)
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"),
)
assert mock_guardrail.apply_guardrail.await_args.kwargs["request_data"] == {
"metadata": {"user_api_key_alias": "billing-app"}
}
@pytest.mark.asyncio
async def test_apply_guardrail_sends_the_key_identity_when_the_caller_sends_no_metadata(mocker: MockerFixture) -> None:
"""A caller that sends no metadata still arrives with the key's identity attached, because
pre-call processing fills it in. An empty request_data is not a shape the pipeline produces."""
mock_guardrail = _patch_apply_guardrail_env(
mocker,
{"texts": ["ok"]},
processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}},
)
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["hello"]},
request_data={},
request_data={"metadata": {"user_api_key_alias": "billing-app"}},
input_type="request",
)