mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(proxy): set x-litellm-applied-guardrails on pass-through responses
Pass-through responses never carried the applied-guardrails header family, so a caller on /openai, /anthropic, /gemini or /vertex_ai could not see which guardrails ran. Every pass-through exit (buffered, relayed, streamed, guardrail-modified and error) now sets the headers from what the pre-call and post-call guardrails recorded, combining both lists. The metadata fields those headers read are stripped from the client body first, so a caller cannot claim a guardrail ran.
This commit is contained in:
yucheng
parent
6eb6f10091
commit
cc57a4fbd6
3 changed files with 268 additions and 1 deletions
|
|
@ -104,6 +104,12 @@ from litellm.proxy.common_request_processing import (
|
|||
open_sse_before_first_byte,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
GUARDRAIL_SCAN_IDS_METADATA_KEY,
|
||||
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
|
||||
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY,
|
||||
get_logging_caching_headers, # pyright: ignore[reportUnknownVariableType] # untyped request_data param
|
||||
)
|
||||
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
|
||||
|
|
@ -176,14 +182,63 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging()
|
|||
_METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata"))
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object])
|
||||
_RESPONSE_HEADERS: Final = TypeAdapter(Mapping[str, str])
|
||||
_KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState))
|
||||
_PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",))
|
||||
# The metadata fields get_logging_caching_headers turns into response headers.
|
||||
_GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset(
|
||||
(
|
||||
"applied_guardrails",
|
||||
"applied_policies",
|
||||
"policy_sources",
|
||||
"semantic-similarity",
|
||||
GUARDRAIL_SCAN_IDS_METADATA_KEY,
|
||||
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
|
||||
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY,
|
||||
"pillar_response_headers",
|
||||
"pillar_flagged",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
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 _strip_client_guardrail_header_metadata(
|
||||
data: dict[str, object], # mutable-ok: strips the request body in place, like strip_client_guardrail_information
|
||||
) -> None:
|
||||
"""Drop the metadata fields that become guardrail response headers, so only the guardrails that ran set them."""
|
||||
for metadata in (as_str_mapping(data.get(key)) for key in sorted(_METADATA_KEYS)):
|
||||
if isinstance(metadata, dict):
|
||||
for field in _GUARDRAIL_HEADER_METADATA_FIELDS & metadata.keys():
|
||||
metadata.pop(field)
|
||||
|
||||
|
||||
def _merged_metadata_entry(prior: object, value: object) -> object:
|
||||
"""``value``, or both lists combined without repeats when ``prior`` and ``value`` are lists."""
|
||||
if not isinstance(prior, list) or not isinstance(value, list):
|
||||
return value
|
||||
prior_entries: Final = _GUARDRAIL_ENTRIES.validate_python(prior)
|
||||
return [*prior_entries, *(e for e in _GUARDRAIL_ENTRIES.validate_python(value) if e not in prior_entries)]
|
||||
|
||||
|
||||
def _guardrail_response_headers(*request_data: object) -> Mapping[str, str]:
|
||||
"""The ``x-litellm-applied-guardrails`` family of response headers for what guardrails recorded in these dicts.
|
||||
|
||||
Reads the metadata of each dict and of its ``litellm_params``. The pre-call and post-call hooks record on
|
||||
different dicts, so list entries such as the applied guardrail names are combined.
|
||||
"""
|
||||
merged: Final = dict[str, object]()
|
||||
for data in request_data:
|
||||
fields = as_str_mapping(data) or {}
|
||||
for scope in (fields, as_str_mapping(fields.get("litellm_params")) or {}):
|
||||
for metadata_key in sorted(_METADATA_KEYS):
|
||||
for key, value in (as_str_mapping(scope.get(metadata_key)) or {}).items():
|
||||
merged[key] = _merged_metadata_entry(merged.get(key), value)
|
||||
return _RESPONSE_HEADERS.validate_python(get_logging_caching_headers({"metadata": merged}) or {})
|
||||
|
||||
|
||||
class _RegisteredPassThroughRoute(TypedDict):
|
||||
endpoint_id: ReadOnly[str]
|
||||
path: ReadOnly[str | None]
|
||||
|
|
@ -410,7 +465,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
headers: httpx.Headers,
|
||||
litellm_call_id: str | None = None,
|
||||
custom_headers: Mapping[str, str] | None = None,
|
||||
guardrail_request_data: Sequence[object] = (),
|
||||
) -> dict:
|
||||
"""The upstream response headers to relay, plus the call id, ``custom_headers`` and the guardrail headers
|
||||
for what guardrails recorded in ``guardrail_request_data``."""
|
||||
# Exclude headers that uvicorn writes itself (server, date) and
|
||||
# encoding/length headers that don't survive re-serialization.
|
||||
# If we forward the upstream's Server header, uvicorn adds its
|
||||
|
|
@ -435,6 +493,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
key: value for key, value in custom_headers.items() if key.lower() not in excluded_headers
|
||||
}
|
||||
return_headers.update(sanitized_custom_headers)
|
||||
return_headers.update(_guardrail_response_headers(*guardrail_request_data))
|
||||
|
||||
return return_headers
|
||||
|
||||
|
|
@ -1343,6 +1402,7 @@ async def pass_through_request(
|
|||
|
||||
client_body: Final = TypeAdapter(dict[str, object]).validate_python(_parsed_body)
|
||||
strip_client_guardrail_information(client_body)
|
||||
_strip_client_guardrail_header_metadata(client_body)
|
||||
if not _key_or_team_allows_client_pricing_override(user_api_key_dict):
|
||||
_strip_client_pricing_overrides(client_body)
|
||||
eu_jev_client_body: Final = _eu_jev_client_body(custom_llm_provider, client_body)
|
||||
|
|
@ -1706,6 +1766,7 @@ async def pass_through_request(
|
|||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
|
|
@ -1808,6 +1869,7 @@ async def pass_through_request(
|
|||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=detected_relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
|
|
@ -1884,6 +1946,7 @@ async def pass_through_request(
|
|||
headers=HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
custom_headers=relay_custom_headers,
|
||||
guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -2052,6 +2115,7 @@ async def pass_through_request(
|
|||
response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
custom_headers=custom_headers,
|
||||
guardrail_request_data=(_parsed_body, kwargs, post_call_guardrail_data), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
)
|
||||
emitted_call_id: Final = (
|
||||
JSON_OBJECT.validate_python(response_headers).get(LITELLM_CALL_ID_HEADER)
|
||||
|
|
@ -2110,7 +2174,10 @@ async def pass_through_request(
|
|||
content=json.dumps(error_body),
|
||||
status_code=200,
|
||||
media_type="application/json",
|
||||
headers={"x-litellm-call-id": litellm_call_id},
|
||||
headers={
|
||||
"x-litellm-call-id": litellm_call_id,
|
||||
**_guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
|
|
@ -2120,6 +2187,7 @@ async def pass_through_request(
|
|||
cache_key=None,
|
||||
api_base=str(url._uri_reference) if url else None,
|
||||
)
|
||||
custom_headers.update(_guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data)) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # untyped request and header dicts
|
||||
if CustomGuardrail._is_guardrail_intervention(e):
|
||||
verbose_proxy_logger.warning(
|
||||
"pass_through_endpoint: request blocked by guardrail - %s",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,164 @@
|
|||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import message_json
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.openai_wire import chat_reply
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
PRE_CALL: Final = "passthrough-header-pre"
|
||||
POST_CALL: Final = "passthrough-header-post"
|
||||
HEADER: Final = "x-litellm-applied-guardrails"
|
||||
OPENAI_ROUTE: Final = "/openai/v1/chat/completions"
|
||||
ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
|
||||
GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _route_body(route: str) -> dict[str, JsonValue]:
|
||||
match route:
|
||||
case "/anthropic/v1/messages":
|
||||
return {"model": "claude-sonnet-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": "Hi"}]}
|
||||
case "/gemini/v1beta/models/gemini-2.5-flash:generateContent":
|
||||
return {"contents": [{"role": "user", "parts": [{"text": "Hi"}]}]}
|
||||
case _:
|
||||
return {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Hi"}]}
|
||||
|
||||
|
||||
def _allow(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
|
||||
def _openai_reply(request: Request) -> Reply:
|
||||
streaming: Final = _JSON.validate_json(request.body).get("stream") is True
|
||||
return chat_reply("chatcmpl-passthrough-header", "gpt-5.6", "Hello", stream=streaming)
|
||||
|
||||
|
||||
def _anthropic_reply(_request: Request) -> Reply:
|
||||
return Reply(body=message_json("msg_passthrough_header", "claude-sonnet-5-5", "Hello"))
|
||||
|
||||
|
||||
def _gemini_reply(_request: Request) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": "Hello"}]}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HeaderRig:
|
||||
owned: OwnedProxy
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[HeaderRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("passthrough-applied-guardrails-header")
|
||||
with ExitStack() as stack:
|
||||
vendor: Final = stack.enter_context(wire_server(_allow))
|
||||
upstreams: Final = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)),
|
||||
ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)),
|
||||
GEMINI_ROUTE: stack.enter_context(wire_server(_gemini_reply)),
|
||||
}
|
||||
)
|
||||
config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
|
||||
path: Final = directory / "passthrough-applied-guardrails-header.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": vendor.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
}
|
||||
for name, mode in ((PRE_CALL, "pre_call"), (POST_CALL, "post_call"))
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
environment: Final = {
|
||||
"OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url,
|
||||
"OPENAI_API_KEY": "synthetic-openai-key",
|
||||
"ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url,
|
||||
"ANTHROPIC_API_KEY": "synthetic-anthropic-key",
|
||||
"GEMINI_API_BASE": upstreams[GEMINI_ROUTE].url,
|
||||
"GEMINI_API_KEY": "synthetic-gemini-key",
|
||||
}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield HeaderRig(owned, upstreams)
|
||||
|
||||
|
||||
def _send(
|
||||
gateway: Gateway, scenario: Scenario, route: str, body: Mapping[str, JsonValue], *guardrails: str
|
||||
) -> httpx.Response:
|
||||
key: Final = scenario.key(metadata={"guardrails": list(guardrails)}) if guardrails else scenario.key()
|
||||
return gateway.request("POST", route, body, key=key, headers={"x-goog-api-key": key})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
|
||||
def test_response_names_the_guardrail_that_ran(rig: HeaderRig, route: str) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, route, _route_body(route), PRE_CALL)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
|
||||
|
||||
def test_streamed_response_names_the_guardrail_that_ran(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(
|
||||
gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), "stream": True}, PRE_CALL
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get("content-type", "").startswith("text/event-stream")
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
|
||||
|
||||
def test_response_names_pre_and_post_call_guardrails(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), PRE_CALL, POST_CALL)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == f"{PRE_CALL},{POST_CALL}"
|
||||
|
||||
|
||||
def test_client_cannot_claim_a_guardrail_ran(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
forged: Final = {
|
||||
**_route_body(OPENAI_ROUTE),
|
||||
"metadata": {"applied_guardrails": ["forged-guardrail"], "guardrail_scan_ids": ["forged-scan"]},
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, forged)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert HEADER not in response.headers
|
||||
assert "x-litellm-guardrail-scan-id" not in response.headers
|
||||
|
|
@ -51,6 +51,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
SafeRouteAdder,
|
||||
_build_passthrough_failure_request_payload,
|
||||
_registered_pass_through_routes,
|
||||
_strip_client_guardrail_header_metadata,
|
||||
_truncate_upstream_error_body,
|
||||
_update_metadata_with_tags_in_header,
|
||||
_with_trace_context,
|
||||
|
|
@ -11472,3 +11473,37 @@ class TestInheritAuthGuardrails:
|
|||
|
||||
assert "guardrails" not in pre_call[0].get("metadata", {})
|
||||
post_call.assert_not_awaited()
|
||||
|
||||
|
||||
class TestGuardrailResponseHeaders:
|
||||
def test_pre_and_post_call_guardrails_are_combined_without_repeats(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{"model": "gpt-5.6"},
|
||||
{"litellm_params": {"metadata": {"applied_guardrails": ["pre"], "guardrail_scan_ids": ["scan-1"]}}},
|
||||
{"metadata": {"applied_guardrails": ["pre", "post"]}},
|
||||
None,
|
||||
),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-applied-guardrails"] == "pre,post"
|
||||
assert headers["x-litellm-guardrail-scan-id"] == "scan-1"
|
||||
|
||||
def test_no_guardrail_data_adds_no_headers(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=({"metadata": {"tags": ["team-a"]}},),
|
||||
)
|
||||
|
||||
assert headers == {"content-type": "application/json"}
|
||||
|
||||
def test_client_sent_guardrail_header_fields_are_stripped(self):
|
||||
body: Final = {
|
||||
"metadata": {"applied_guardrails": ["forged"], "pillar_flagged": True, "tags": ["team-a"]},
|
||||
"litellm_metadata": {"guardrail_scan_ids": ["forged-scan"]},
|
||||
}
|
||||
|
||||
_strip_client_guardrail_header_metadata(body)
|
||||
|
||||
assert body == {"metadata": {"tags": ["team-a"]}, "litellm_metadata": {}}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue