mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 3b5f55865c into c0bb0c336f
This commit is contained in:
yucheng-berri
GitHub
commit
57f9dddee0
4 changed files with 986 additions and 26 deletions
|
|
@ -19,7 +19,7 @@ from collections.abc import (
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import partial, reduce
|
||||
from itertools import chain, count, groupby, takewhile
|
||||
from itertools import chain, count, groupby, product, takewhile
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast
|
||||
from urllib.parse import parse_qs, urlencode, urlparse
|
||||
|
|
@ -124,7 +124,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 strip_callback_config
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
GUARDRAIL_SCAN_IDS_METADATA_KEY,
|
||||
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
|
||||
get_logging_caching_headers, # pyright: ignore[reportUnknownVariableType] # untyped request_data param
|
||||
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
|
||||
|
|
@ -204,6 +209,7 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch=
|
|||
_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])
|
||||
_FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = (
|
||||
("guardrails", (list, dict)),
|
||||
("guardrail_config", (dict,)),
|
||||
|
|
@ -224,12 +230,172 @@ _POST_CALL_SELECTION_FIELDS: Final = (
|
|||
)
|
||||
_KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState))
|
||||
_PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",))
|
||||
_GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset(
|
||||
(
|
||||
"applied_guardrails",
|
||||
"applied_policies",
|
||||
"policy_sources",
|
||||
GUARDRAIL_SCAN_IDS_METADATA_KEY,
|
||||
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
|
||||
)
|
||||
)
|
||||
_JOINED_HEADER_METADATA_FIELDS: Final = frozenset(
|
||||
("applied_guardrails", "applied_policies", GUARDRAIL_SCAN_IDS_METADATA_KEY)
|
||||
)
|
||||
# RFC 9110's field-value, the one both h11 and Node's HTTP parser accept: no control characters, and only spaces and
|
||||
# tabs between visible characters.
|
||||
_SENDABLE_HEADER_VALUE: Final = re.compile(rb"(?:[\x21-\x7e\x80-\xff]+(?:[ \t]+[\x21-\x7e\x80-\xff]+)*)?")
|
||||
# Half of the 16 KB response header limit in Node.js, the smallest common client limit, leaving room for the others.
|
||||
_GUARDRAIL_HEADERS_MAX_BYTES: Final = 8192
|
||||
# Node.js's limit for a whole response, less 1 KB for the status line and what the server adds.
|
||||
_RESPONSE_HEADERS_MAX_BYTES: Final = 15360
|
||||
# The ": " and the line ending around each header on the wire.
|
||||
_HEADER_FRAMING_BYTES: Final = 4
|
||||
|
||||
|
||||
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_fields(
|
||||
data: dict[str, object], # mutable-ok: strips the request body in place, like strip_client_guardrail_information
|
||||
) -> None:
|
||||
"""Drop the guardrail logging and the metadata fields that become guardrail response headers from a client body,
|
||||
so only the guardrails that ran set them."""
|
||||
strip_client_guardrail_information(data)
|
||||
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 _is_name_list(value: object) -> bool:
|
||||
return isinstance(value, list | tuple) and all(
|
||||
isinstance(entry, str) for entry in _GUARDRAIL_ENTRIES.validate_python(value)
|
||||
)
|
||||
|
||||
|
||||
def _joinable_header_metadata(recorded: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""``recorded`` without a value ``get_logging_caching_headers`` can't join into a header, which a custom guardrail
|
||||
can record in the name lists."""
|
||||
return {
|
||||
field: value
|
||||
for field, value in recorded.items()
|
||||
if field not in _JOINED_HEADER_METADATA_FIELDS or _is_name_list(value)
|
||||
}
|
||||
|
||||
|
||||
def _metadata_entry_identity(entry: object) -> object:
|
||||
return entry if isinstance(entry, str) else (json.dumps(entry, sort_keys=True, default=str),)
|
||||
|
||||
|
||||
def _merged_metadata_entry(prior: object, value: object) -> object:
|
||||
"""``value``, or both combined: lists and tuples without repeats, and mappings with ``value``'s keys winning, as
|
||||
``add_policy_sources_to_metadata`` combines policy sources on one request."""
|
||||
prior_mapping: Final = as_str_mapping(prior)
|
||||
value_mapping: Final = as_str_mapping(value)
|
||||
if prior_mapping is not None and value_mapping is not None:
|
||||
return {**prior_mapping, **value_mapping}
|
||||
if not isinstance(prior, list | tuple) or not isinstance(value, list | tuple):
|
||||
return value
|
||||
entries: Final = (*_GUARDRAIL_ENTRIES.validate_python(prior), *_GUARDRAIL_ENTRIES.validate_python(value))
|
||||
return tuple({_metadata_entry_identity(entry): entry for entry in entries}.values())
|
||||
|
||||
|
||||
def _is_sendable_header_value(value: str) -> bool:
|
||||
return all(ord(character) < 256 for character in value) and bool(
|
||||
_SENDABLE_HEADER_VALUE.fullmatch(value.encode("latin-1"))
|
||||
)
|
||||
|
||||
|
||||
def _with_header_in_budget(kept: Mapping[str, str], header: tuple[str, str]) -> Mapping[str, str]:
|
||||
size: Final = sum(len(name) + len(value) for name, value in (*kept.items(), header))
|
||||
return MappingProxyType({**kept, header[0]: header[1]}) if size <= _GUARDRAIL_HEADERS_MAX_BYTES else kept
|
||||
|
||||
|
||||
def _with_header_in_response_limit(headers: Mapping[str, str], header: tuple[str, str]) -> Mapping[str, str]:
|
||||
merged: Final = {**headers, header[0]: header[1]}
|
||||
size: Final = sum(len(name) + len(value) + _HEADER_FRAMING_BYTES for name, value in merged.items())
|
||||
return merged if size <= _RESPONSE_HEADERS_MAX_BYTES else headers
|
||||
|
||||
|
||||
def _guardrail_headers_with_room(
|
||||
guardrail_headers: Mapping[str, str], response_headers: Mapping[str, str]
|
||||
) -> Mapping[str, str]:
|
||||
"""``response_headers`` with each of ``guardrail_headers`` that fits, so adding them can't cost a client the
|
||||
response. One that doesn't fit leaves the header of that name as it was."""
|
||||
return reduce(_with_header_in_response_limit, guardrail_headers.items(), response_headers)
|
||||
|
||||
|
||||
def _add_callback_headers(
|
||||
response_headers: dict[str, str], # mutable-ok: adds to the response's headers in place, like the dict's update
|
||||
callback_headers: Mapping[str, str],
|
||||
own_guardrail_headers: Mapping[str, str],
|
||||
) -> None:
|
||||
"""Add a callback's headers to ``response_headers``, less a guardrail header this proxy set, ``own_guardrail_headers``,
|
||||
that they then leave no room for. The upstream's guardrail headers are kept as they came."""
|
||||
guardrail_headers: Final = {
|
||||
name: value
|
||||
for name, value in own_guardrail_headers.items()
|
||||
if response_headers.get(name) == value and name not in callback_headers
|
||||
}
|
||||
for name in guardrail_headers:
|
||||
response_headers.pop(name)
|
||||
response_headers.update(callback_headers)
|
||||
response_headers.update(_guardrail_headers_with_room(guardrail_headers, response_headers))
|
||||
|
||||
|
||||
def _recorded_guardrail_headers(
|
||||
parsed_body: object = None, kwargs: object = None, post_call_data: object = None
|
||||
) -> Mapping[str, str]:
|
||||
"""The ``x-litellm-applied-guardrails`` family of response headers for what guardrails recorded on this request.
|
||||
|
||||
Pre-call guardrails record on the request body's metadata, which ``init_kwargs_for_pass_through_endpoint`` moves to
|
||||
``kwargs["litellm_params"]``, and post-call guardrails record on ``post_call_data``, so list entries such as the
|
||||
applied guardrail names are combined. A ``litellm_params`` object sent in the request body is never read. A header
|
||||
that would take these headers past ``_GUARDRAIL_HEADERS_MAX_BYTES``, or whose recorded value can't be joined, is
|
||||
left out.
|
||||
"""
|
||||
scopes: Final = (
|
||||
as_str_mapping(parsed_body) or {},
|
||||
as_str_mapping((as_str_mapping(kwargs) or {}).get("litellm_params")) or {},
|
||||
as_str_mapping(post_call_data) or {},
|
||||
)
|
||||
metadatas: Final = tuple(
|
||||
as_str_mapping(scope.get(metadata_key)) or {} for scope, metadata_key in product(scopes, sorted(_METADATA_KEYS))
|
||||
)
|
||||
recorded: Final = chain.from_iterable(
|
||||
((field, metadata[field]) for field in sorted(_GUARDRAIL_HEADER_METADATA_FIELDS & metadata.keys()))
|
||||
for metadata in metadatas
|
||||
)
|
||||
merged: Final = {
|
||||
field: reduce(_merged_metadata_entry, (value for _, value in entries))
|
||||
for field, entries in groupby(sorted(recorded, key=lambda entry: entry[0]), key=lambda entry: entry[0])
|
||||
}
|
||||
headers: Final = _RESPONSE_HEADERS.validate_python(
|
||||
get_logging_caching_headers({"metadata": _joinable_header_metadata(merged)}) or {}
|
||||
)
|
||||
sendable: Final = (
|
||||
(name, value)
|
||||
for name, value in headers.items()
|
||||
if len(value) <= _GUARDRAIL_HEADERS_MAX_BYTES and _is_sendable_header_value(value)
|
||||
)
|
||||
no_headers: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
return reduce(_with_header_in_budget, sendable, no_headers)
|
||||
|
||||
|
||||
def _guardrail_response_headers(
|
||||
parsed_body: object = None, kwargs: object = None, post_call_data: object = None
|
||||
) -> Mapping[str, str]:
|
||||
try:
|
||||
return _recorded_guardrail_headers(parsed_body, kwargs, post_call_data)
|
||||
except (TypeError, ValueError, RecursionError) as error:
|
||||
verbose_proxy_logger.warning(
|
||||
"Leaving out guardrail response headers: a recorded value can't be encoded (%s)", type(error).__name__
|
||||
)
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
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())
|
||||
|
|
@ -469,6 +635,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
headers: httpx.Headers,
|
||||
litellm_call_id: str | None = None,
|
||||
custom_headers: Mapping[str, str] | None = None,
|
||||
guardrail_request_data: Sequence[object] = (),
|
||||
earlier_guardrail_headers: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> dict:
|
||||
# Exclude headers that uvicorn writes itself (server, date) and
|
||||
# encoding/length headers that don't survive re-serialization.
|
||||
|
|
@ -494,6 +662,13 @@ 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)
|
||||
guardrail_headers: Final = _guardrail_response_headers(*guardrail_request_data)
|
||||
combined_guardrail_headers: Final = reduce(
|
||||
_with_header_in_budget,
|
||||
(header for header in earlier_guardrail_headers.items() if header[0] not in guardrail_headers),
|
||||
guardrail_headers,
|
||||
)
|
||||
return_headers.update(_guardrail_headers_with_room(combined_guardrail_headers, return_headers))
|
||||
|
||||
return return_headers
|
||||
|
||||
|
|
@ -1482,6 +1657,29 @@ async def _log_passthrough_upstream_failure(
|
|||
return relay_response
|
||||
|
||||
|
||||
async def _log_passthrough_upstream_failure_keeping_guardrail_headers(
|
||||
response: httpx.Response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_payload: dict[str, object], # mutable-ok: the failure hooks lift fields onto request_data in place
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
parsed_body: Mapping[str, object] | None,
|
||||
kwargs: Mapping[str, object] | None,
|
||||
) -> Mapping[str, str]:
|
||||
"""Log an upstream error like ``_log_passthrough_upstream_failure`` and return the guardrail response headers,
|
||||
read before the failure hooks rebuild ``kwargs``' metadata without the guardrail fields. A successful response
|
||||
has no failure hooks to read around, so it returns none."""
|
||||
guardrail_headers: Final[Mapping[str, str]] = (
|
||||
_guardrail_response_headers(parsed_body, kwargs) if response.status_code >= 400 else {}
|
||||
)
|
||||
await _log_passthrough_upstream_failure(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=request_payload,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return guardrail_headers
|
||||
|
||||
|
||||
async def _relay_reporting_failures(
|
||||
stream: AsyncGenerator[bytes, None],
|
||||
upstream_status: int,
|
||||
|
|
@ -1681,7 +1879,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_fields(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)
|
||||
|
|
@ -2031,6 +2229,14 @@ async def pass_through_request(
|
|||
)
|
||||
apply_upstream_reported_operation(logging_obj, response.headers)
|
||||
|
||||
# Built before the failure hooks rebuild kwargs' metadata without the guardrail fields. The relay
|
||||
# response differs only in body-framing headers, which get_response_headers drops anyway.
|
||||
streamed_guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
earlier_guardrail_headers=streamed_guardrail_headers,
|
||||
)
|
||||
relay_response: Final = await _log_passthrough_upstream_failure(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2048,10 +2254,6 @@ async def pass_through_request(
|
|||
)
|
||||
|
||||
# Call response headers hook for streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2059,7 +2261,7 @@ async def pass_through_request(
|
|||
request_headers=dict(redact_credential_headers(request.headers)),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
_add_callback_headers(_response_headers, callback_headers, streamed_guardrail_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict
|
||||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
|
|
@ -2133,6 +2335,14 @@ async def pass_through_request(
|
|||
logging_obj.stream = True
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
|
||||
# Built before the failure hooks rebuild kwargs' metadata without the guardrail fields. The relay
|
||||
# response differs only in body-framing headers, which get_response_headers drops anyway.
|
||||
detected_stream_guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
earlier_guardrail_headers=detected_stream_guardrail_headers,
|
||||
)
|
||||
detected_relay_response: Final = await _log_passthrough_upstream_failure(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2150,10 +2360,6 @@ async def pass_through_request(
|
|||
)
|
||||
|
||||
# Call response headers hook for detected streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=detected_relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2161,7 +2367,7 @@ async def pass_through_request(
|
|||
request_headers=dict(redact_credential_headers(request.headers)),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
_add_callback_headers(_response_headers, callback_headers, detected_stream_guardrail_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict
|
||||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
|
|
@ -2229,6 +2435,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
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -2252,11 +2459,13 @@ async def pass_through_request(
|
|||
upstream_model=upstream_model,
|
||||
)
|
||||
failure_request_payload["response_body"] = response_body
|
||||
await _log_passthrough_upstream_failure(
|
||||
upstream_guardrail_headers: Final = await _log_passthrough_upstream_failure_keeping_guardrail_headers(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=failure_request_payload,
|
||||
request_payload=failure_request_payload, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
logging_obj=logging_obj,
|
||||
parsed_body=_parsed_body, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
)
|
||||
|
||||
policy_pipelines: Final = _post_call_selection_metadata(
|
||||
|
|
@ -2412,6 +2621,8 @@ async def pass_through_request(
|
|||
response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
custom_headers=custom_headers,
|
||||
earlier_guardrail_headers=upstream_guardrail_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)
|
||||
|
|
@ -2440,6 +2651,7 @@ async def pass_through_request(
|
|||
e.guardrail_name,
|
||||
str(e.message or "")[:200],
|
||||
)
|
||||
modified_response_headers: Final = _guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
try:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2470,7 +2682,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,
|
||||
**modified_response_headers,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
|
|
@ -2480,6 +2695,8 @@ async def pass_through_request(
|
|||
cache_key=None,
|
||||
api_base=str(url._uri_reference) if url else None,
|
||||
)
|
||||
guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
|
||||
custom_headers.update(guardrail_headers) # pyright: ignore[reportUnknownMemberType] # untyped header dict
|
||||
if CustomGuardrail._is_guardrail_intervention(e):
|
||||
verbose_proxy_logger.warning(
|
||||
"pass_through_endpoint: request blocked by guardrail - %s",
|
||||
|
|
@ -2516,6 +2733,7 @@ async def pass_through_request(
|
|||
#########################################################
|
||||
|
||||
if isinstance(e, ProxyException):
|
||||
e.headers = {**e.headers, **guardrail_headers}
|
||||
raise
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,362 @@
|
|||
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"
|
||||
NON_LATIN1: Final = "passthrough-header-pre\u2014em"
|
||||
BLOCKER: Final = "passthrough-header-block"
|
||||
MISSING_MODEL: Final = "passthrough-header-missing-model"
|
||||
PILLAR: Final = "passthrough-header-pillar"
|
||||
TAG_POLICY: Final = "passthrough-header-tag-policy"
|
||||
POLICY_TAG_PREFIX: Final = "passthrough-header-tag-"
|
||||
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"
|
||||
CALLBACK_HEADERS_MODEL: Final = "passthrough-header-callback-headers"
|
||||
UNREQUESTED_STREAM_MODEL: Final = "passthrough-header-callback-headers-unrequested-stream"
|
||||
CALLBACK_HEADER: Final = "x-callback-padding"
|
||||
CALLBACK_HEADER_VALUE: Final = "c" * 12_000
|
||||
_CALLBACK_HEADERS_MODULE: Final = f"""
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
class LargeHeadersHook(CustomLogger):
|
||||
async def async_post_call_response_headers_hook(self, data, user_api_key_dict, response, *args, **kwargs):
|
||||
if str(data.get("model")).startswith("{CALLBACK_HEADERS_MODEL}"):
|
||||
return {{"{CALLBACK_HEADER}": "{CALLBACK_HEADER_VALUE}"}}
|
||||
return None
|
||||
|
||||
|
||||
proxy_handler_instance = LargeHeadersHook()
|
||||
"""
|
||||
_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 _block(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode())
|
||||
|
||||
|
||||
def _pillar_allow(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"flagged": False, "session_id": "pillar-vendor-session"}).encode())
|
||||
|
||||
|
||||
def _openai_reply(request: Request) -> Reply:
|
||||
payload: Final = _JSON.validate_json(request.body)
|
||||
if payload.get("model") == MISSING_MODEL:
|
||||
return Reply(
|
||||
status=404,
|
||||
body=json.dumps({"error": {"message": "model not found", "type": "invalid_request_error"}}).encode(),
|
||||
)
|
||||
return chat_reply(
|
||||
"chatcmpl-passthrough-header",
|
||||
"gpt-5.6",
|
||||
"Hello",
|
||||
stream=payload.get("stream") is True or payload.get("model") == UNREQUESTED_STREAM_MODEL,
|
||||
)
|
||||
|
||||
|
||||
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))
|
||||
blocking_vendor: Final = stack.enter_context(wire_server(_block))
|
||||
pillar_vendor: Final = stack.enter_context(wire_server(_pillar_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()))
|
||||
(directory / "passthrough_large_headers_hook.py").write_text(_CALLBACK_HEADERS_MODULE)
|
||||
path: Final = directory / "passthrough-applied-guardrails-header.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"litellm_settings": {
|
||||
**_JSON.validate_python(config.get("litellm_settings", {})),
|
||||
"callbacks": ["passthrough_large_headers_hook.proxy_handler_instance"],
|
||||
},
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": api_base,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
}
|
||||
for name, mode, api_base in (
|
||||
(PRE_CALL, "pre_call", vendor.url),
|
||||
(POST_CALL, "post_call", vendor.url),
|
||||
(NON_LATIN1, "pre_call", vendor.url),
|
||||
(BLOCKER, "pre_call", blocking_vendor.url),
|
||||
)
|
||||
]
|
||||
+ [
|
||||
{
|
||||
"guardrail_name": PILLAR,
|
||||
"litellm_params": {
|
||||
"guardrail": "pillar",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": pillar_vendor.url,
|
||||
"api_key": "synthetic-pillar-key",
|
||||
},
|
||||
}
|
||||
],
|
||||
"policies": {TAG_POLICY: {"guardrails": {"add": [PRE_CALL]}}},
|
||||
"policy_attachments": [{"policy": TAG_POLICY, "tags": [f"{POLICY_TAG_PREFIX}*"]}],
|
||||
}
|
||||
)
|
||||
)
|
||||
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}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["metadata", "litellm_params"])
|
||||
def test_client_cannot_claim_a_guardrail_ran(rig: HeaderRig, field: str) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
claim: Final[dict[str, JsonValue]] = {
|
||||
"applied_guardrails": ["forged-guardrail"],
|
||||
"guardrail_scan_ids": ["forged-scan"],
|
||||
}
|
||||
forged: Final = {**_route_body(OPENAI_ROUTE), field: claim if field == "metadata" else {"metadata": claim}}
|
||||
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
|
||||
|
||||
|
||||
def test_guardrail_name_a_header_cannot_carry_is_left_out(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), NON_LATIN1)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert HEADER not in response.headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
def test_upstream_error_names_the_guardrail_that_ran(rig: HeaderRig, stream: bool) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
body: Final = {**_route_body(OPENAI_ROUTE), "model": MISSING_MODEL, "stream": stream}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, body, PRE_CALL)
|
||||
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
|
||||
|
||||
def test_blocked_request_names_the_blocking_guardrail(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
upstream: Final = rig.upstreams[OPENAI_ROUTE]
|
||||
upstream.drain()
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), PRE_CALL, BLOCKER)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.headers.get(HEADER) == f"{PRE_CALL},{BLOCKER}"
|
||||
assert upstream.drain() == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"])
|
||||
def test_client_claims_beside_a_real_guardrail_are_dropped(rig: HeaderRig, field: str) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
claim: Final[dict[str, JsonValue]] = {
|
||||
"applied_guardrails": ["forged-guardrail"],
|
||||
"guardrail_scan_ids": ["forged-scan"],
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), field: claim}, PRE_CALL)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
assert "forged-scan" not in response.headers.get("x-litellm-guardrail-scan-id", "")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("claim", ["none", 7, ""])
|
||||
def test_non_list_client_claim_does_not_fail_a_guarded_request(rig: HeaderRig, claim: JsonValue) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
body: Final = {**_route_body(OPENAI_ROUTE), "metadata": {"applied_guardrails": claim}}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, body, PRE_CALL)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
|
||||
|
||||
def test_pillar_headers_stay_off_pass_through_responses(rig: HeaderRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
forged: Final[dict[str, JsonValue]] = {"pillar_session_id_response": "client-session", "pillar_scanners": {"x": 1}}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(
|
||||
gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), "metadata": forged}, PILLAR
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == PILLAR
|
||||
assert [name for name, _ in response.headers.multi_items() if name.startswith("x-pillar-")] == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("suffix", ["team\n", "x\r\nX-Injected: yes"])
|
||||
def test_header_value_with_a_line_break_is_left_out(rig: HeaderRig, suffix: str) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
**_route_body(OPENAI_ROUTE),
|
||||
"metadata": {"tags": [f"{POLICY_TAG_PREFIX}{suffix}"]},
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(HEADER) == PRE_CALL
|
||||
assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY
|
||||
assert "x-litellm-policy-sources" not in response.headers
|
||||
assert "x-injected" not in response.headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "stream"),
|
||||
[(CALLBACK_HEADERS_MODEL, True), (UNREQUESTED_STREAM_MODEL, False)],
|
||||
ids=["requested-stream", "unrequested-stream"],
|
||||
)
|
||||
@pytest.mark.parametrize(("names", "kept"), [(300, False), (100, True)], ids=["no-room", "room"])
|
||||
def test_a_callback_header_on_a_stream_leaves_room_or_costs_the_guardrail_header(
|
||||
rig: HeaderRig, model: str, stream: bool, names: int, kept: bool
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
**_route_body(OPENAI_ROUTE),
|
||||
"model": model,
|
||||
"stream": stream,
|
||||
"guardrails": [f"caller-named-{index}" for index in range(names)],
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers.get(CALLBACK_HEADER) == CALLBACK_HEADER_VALUE
|
||||
assert (HEADER in response.headers) is kept
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
def test_response_to_a_body_naming_thousands_of_guardrails_stays_readable(rig: HeaderRig, stream: bool) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
**_route_body(OPENAI_ROUTE),
|
||||
"stream": stream,
|
||||
"guardrails": [f"caller-named-{index}" for index in range(10_000)],
|
||||
"metadata": {"tags": [f"{POLICY_TAG_PREFIX}prod"]},
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
response: Final = _send(gateway, scenario, OPENAI_ROUTE, body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert HEADER not in response.headers
|
||||
assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY
|
||||
|
|
@ -717,6 +717,20 @@ def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGua
|
|||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(rig.key_post_call_vendor.drain()) == 1
|
||||
assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",")
|
||||
|
||||
|
||||
def test_response_a_policy_pipeline_modifies_names_the_step_that_ran(rig: AuthGuardrailRig) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
_drain(rig)
|
||||
|
||||
response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_MODIFIES), key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["error"]["type"] == "content_filter"
|
||||
assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",")
|
||||
|
||||
|
||||
def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRig) -> None:
|
||||
|
|
|
|||
|
|
@ -11,13 +11,14 @@ import zlib
|
|||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from io import BytesIO
|
||||
from types import MappingProxyType, ModuleType, SimpleNamespace
|
||||
from typing import Final, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
from uuid import UUID
|
||||
|
||||
import h11
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
|
@ -67,6 +68,10 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
WebsocketFrameModelVerdict,
|
||||
_build_passthrough_failure_request_payload,
|
||||
_registered_pass_through_routes,
|
||||
_add_callback_headers,
|
||||
_guardrail_response_headers,
|
||||
_log_passthrough_upstream_failure_keeping_guardrail_headers,
|
||||
_strip_client_guardrail_fields,
|
||||
_truncate_upstream_error_body,
|
||||
_update_metadata_with_tags_in_header,
|
||||
_with_trace_context,
|
||||
|
|
@ -505,6 +510,61 @@ async def test_pass_through_request_preserves_proxy_exception_status():
|
|||
assert exc.value.param == "limit"
|
||||
|
||||
|
||||
class _RecordsThenRaisesProxyException(CustomGuardrail):
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict[str, object],
|
||||
call_type: str,
|
||||
) -> None:
|
||||
data["metadata"] = {"applied_guardrails": ["passed-guard"]} # mutable-ok: a guardrail records on the live body
|
||||
raise ProxyException(message="blocked by guardrail", type="None", param="None", code=400)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_proxy_exception_carries_the_guardrails_that_ran():
|
||||
guardrail: Final = _RecordsThenRaisesProxyException(
|
||||
guardrail_name="passed-guard", event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
mock_request: Final = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "messages": []}')
|
||||
mock_request.headers = Headers({"content-type": "application/json"})
|
||||
mock_request.query_params = QueryParams({})
|
||||
|
||||
with (
|
||||
pytest.MonkeyPatch.context() as litellm_patch,
|
||||
patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())
|
||||
),
|
||||
):
|
||||
litellm_patch.setattr(litellm, "callbacks", [guardrail])
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await pass_through_request(
|
||||
request=mock_request,
|
||||
target="http://test.com/v1/chat/completions",
|
||||
custom_headers={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
)
|
||||
|
||||
assert exc.value.headers["x-litellm-applied-guardrails"] == "passed-guard"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_response_keeps_no_guardrail_headers_from_before_the_post_call_pass():
|
||||
headers: Final = await _log_passthrough_upstream_failure_keeping_guardrail_headers(
|
||||
response=httpx.Response(status_code=200),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
request_payload={},
|
||||
logging_obj=MagicMock(),
|
||||
parsed_body={"metadata": {"applied_guardrails": ["pre-call-guard"]}},
|
||||
kwargs=None,
|
||||
)
|
||||
|
||||
assert headers == {}
|
||||
|
||||
|
||||
def test_is_langfuse_route():
|
||||
"""
|
||||
Test that the is_langfuse_route method correctly identifies Langfuse routes
|
||||
|
|
@ -13229,13 +13289,23 @@ async def test_passthrough_header_and_callback_share_finalized_cost(
|
|||
|
||||
client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(upstream), timeout=None)
|
||||
try:
|
||||
request: Final = Request({
|
||||
"type": "http", "method": "GET", "path": "/passthrough-relay/results",
|
||||
"headers": [], "query_string": b"", "scheme": "http", "server": ("test", 80),
|
||||
})
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/passthrough-relay/results",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
"scheme": "http",
|
||||
"server": ("test", 80),
|
||||
}
|
||||
)
|
||||
result: Final = await pass_through_request(
|
||||
request=request, target="http://cost-upstream.invalid/results", custom_headers={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed"), cost_per_request=flat_cost,
|
||||
request=request,
|
||||
target="http://cost-upstream.invalid/results",
|
||||
custom_headers={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed"),
|
||||
cost_per_request=flat_cost,
|
||||
)
|
||||
payload: Final = await asyncio.wait_for(observed, timeout=10)
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
|
|
@ -13790,9 +13860,7 @@ class TestInheritAuthGuardrails:
|
|||
async def test_request_body_guardrails_are_added_to_the_key_guardrails(self):
|
||||
auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
|
||||
|
||||
pre_call, _ = await self._run(
|
||||
self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth
|
||||
)
|
||||
pre_call, _ = await self._run(self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth)
|
||||
|
||||
assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"]
|
||||
assert "guardrails" not in pre_call[0]
|
||||
|
|
@ -13844,3 +13912,301 @@ 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"]}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-applied-guardrails"] == "pre,post"
|
||||
assert headers["x-litellm-guardrail-scan-id"] == "scan-1"
|
||||
|
||||
def test_scan_ids_recorded_as_tuples_before_and_after_the_call_are_both_kept(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{},
|
||||
{"litellm_params": {"metadata": {"guardrail_scan_ids": ("pre-scan",)}}},
|
||||
{"metadata": {"guardrail_scan_ids": ("post-scan",)}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-guardrail-scan-id"] == "pre-scan,post-scan"
|
||||
|
||||
def test_policy_sources_recorded_before_and_after_the_call_are_both_kept(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{},
|
||||
{"litellm_params": {"metadata": {"policy_sources": {"pre-policy": "key", "shared": "team"}}}},
|
||||
{"metadata": {"policy_sources": {"post-policy": "tag", "shared": "key"}}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-policy-sources"] == "pre-policy=key; shared=key; post-policy=tag"
|
||||
|
||||
def test_a_litellm_params_object_in_the_request_body_is_never_read(self):
|
||||
forged: Final = {"metadata": {"applied_guardrails": ["forged"], "guardrail_scan_ids": ["fake"]}}
|
||||
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{"model": "gpt-5.6", "litellm_params": forged},
|
||||
{"litellm_params": {"metadata": {}}},
|
||||
{"model": "gpt-5.6", "litellm_params": forged, "metadata": {"guardrails": ["post"]}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers == {"content-type": "application/json"}
|
||||
|
||||
def test_guardrail_name_outside_latin1_is_left_out_of_the_headers(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{},
|
||||
{
|
||||
"litellm_params": {
|
||||
"metadata": {"applied_guardrails": ["pii\u2014guard"], "guardrail_scan_ids": ["s1"]}
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers == {"content-type": "application/json", "x-litellm-guardrail-scan-id": "s1"}
|
||||
assert Response(content=b"{}", headers=headers).status_code == 200
|
||||
|
||||
def test_metadata_outside_the_header_fields_is_never_read(self):
|
||||
class _Unreadable:
|
||||
__hash__ = object.__hash__
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
raise AssertionError("non-header metadata was read")
|
||||
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"content-type": "application/json"}),
|
||||
guardrail_request_data=(
|
||||
{"metadata": {"unused": [_Unreadable()]}},
|
||||
{"litellm_params": {"metadata": {"unused": [_Unreadable()]}}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers == {"content-type": "application/json"}
|
||||
|
||||
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_fields(body)
|
||||
|
||||
assert body == {"metadata": {"pillar_flagged": True, "tags": ["team-a"]}, "litellm_metadata": {}}
|
||||
|
||||
def test_pillar_fields_recorded_on_the_request_add_no_pillar_headers(self):
|
||||
headers: Final = _guardrail_response_headers(
|
||||
{"metadata": {"pillar_session_id_response": "client-session"}},
|
||||
{
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"applied_guardrails": ["pillar"],
|
||||
"pillar_flagged": False,
|
||||
"pillar_response_headers": {"x-pillar-flagged": "false", "x-pillar-session-id": "vendor"},
|
||||
"_pillar_response_headers_trusted": True,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert headers == {"x-litellm-applied-guardrails": "pillar"}
|
||||
|
||||
@pytest.mark.parametrize("tag", ["prod-team\n", "prod-x\r\nX-Injected: yes", "prod\x00"])
|
||||
def test_header_value_an_http_response_cannot_carry_is_left_out(self, tag: str):
|
||||
headers: Final = _guardrail_response_headers(
|
||||
{},
|
||||
{
|
||||
"litellm_params": {
|
||||
"metadata": {"applied_policies": ["tag-policy"], "policy_sources": {"tag-policy": f"tag:{tag}"}}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert headers == {"x-litellm-applied-policies": "tag-policy"}
|
||||
assert h11.Response(status_code=200, headers=[*headers.items()]).status_code == 200
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_bytes", "names", "kept"),
|
||||
[(12000, 500, False), (12000, 1, True), (8000, 500, True), (15340, 1, False)],
|
||||
ids=["large-and-large", "large-and-small", "both-within-the-limit", "upstream-at-the-limit"],
|
||||
)
|
||||
def test_guardrail_headers_leave_the_whole_response_within_a_clients_limit(self, upstream_bytes, names, kept):
|
||||
upstream: Final = "u" * upstream_bytes
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"x-upstream": upstream}),
|
||||
guardrail_request_data=(
|
||||
{},
|
||||
{"litellm_params": {"metadata": {"applied_guardrails": [f"guardrail-{i:03d}" for i in range(names)]}}},
|
||||
),
|
||||
)
|
||||
|
||||
assert headers["x-upstream"] == upstream
|
||||
assert ("x-litellm-applied-guardrails" in headers) is kept
|
||||
assert sum(len(name) + len(value) for name, value in headers.items()) <= 15360 or not kept
|
||||
|
||||
def test_many_small_upstream_headers_count_with_their_framing(self):
|
||||
upstream: Final = httpx.Headers({f"x-{index:04d}": "v" for index in range(1_500)})
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=upstream,
|
||||
guardrail_request_data=({}, {"litellm_params": {"metadata": {"applied_guardrails": ["guard"]}}}),
|
||||
)
|
||||
|
||||
assert "x-litellm-applied-guardrails" not in headers
|
||||
assert len(headers) == 1_500
|
||||
|
||||
def test_an_upstream_guardrail_header_a_larger_one_cannot_replace_still_counts(self):
|
||||
names: Final = [f"guardrail-{index:03d}" for index in range(280)]
|
||||
upstream: Final = {"x-upstream": "u" * 10_000, "x-litellm-applied-policies": "p" * 4_000}
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers(upstream),
|
||||
guardrail_request_data=(
|
||||
{},
|
||||
{"litellm_params": {"metadata": {"applied_guardrails": names, "applied_policies": names}}},
|
||||
),
|
||||
)
|
||||
|
||||
assert sum(len(name) + len(value) + 4 for name, value in headers.items()) <= 15360
|
||||
assert headers["x-upstream"] == upstream["x-upstream"]
|
||||
|
||||
def test_an_upstream_guardrail_header_the_proxy_replaces_takes_no_room(self):
|
||||
headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=httpx.Headers({"x-litellm-applied-guardrails": "u" * 15_400}),
|
||||
guardrail_request_data=({}, {"litellm_params": {"metadata": {"applied_guardrails": ["guard"]}}}),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-applied-guardrails"] == "guard"
|
||||
|
||||
@staticmethod
|
||||
def _upstream_with_guardrail_headers() -> dict[str, str]:
|
||||
return {
|
||||
"x-litellm-applied-guardrails": "upstream-guard",
|
||||
"x-litellm-guardrail-scan-id": "upstream-scan",
|
||||
"x-upstream-0": "o" * 8_000,
|
||||
"x-upstream-1": "o" * 8_000,
|
||||
}
|
||||
|
||||
def test_upstream_guardrail_headers_survive_a_callback_header_when_the_proxy_recorded_nothing(self):
|
||||
headers: Final = self._upstream_with_guardrail_headers()
|
||||
|
||||
_add_callback_headers(headers, {"x-callback": "c"}, {})
|
||||
|
||||
assert headers == {**self._upstream_with_guardrail_headers(), "x-callback": "c"}
|
||||
|
||||
def test_the_proxys_guardrail_header_gives_way_to_a_callback_header_with_no_room(self):
|
||||
headers: Final = {"x-litellm-applied-guardrails": "guard", "x-upstream": "o" * 15_000}
|
||||
|
||||
_add_callback_headers(headers, {"x-callback": "c" * 300}, {"x-litellm-applied-guardrails": "guard"})
|
||||
|
||||
assert headers == {"x-upstream": "o" * 15_000, "x-callback": "c" * 300}
|
||||
|
||||
def test_the_proxys_guardrail_header_stays_beside_a_callback_header_with_room(self):
|
||||
headers: Final = {"x-litellm-applied-guardrails": "guard", "x-upstream": "o"}
|
||||
|
||||
_add_callback_headers(headers, {"x-callback": "c"}, {"x-litellm-applied-guardrails": "guard"})
|
||||
|
||||
assert headers == {"x-litellm-applied-guardrails": "guard", "x-upstream": "o", "x-callback": "c"}
|
||||
|
||||
def test_guardrail_header_past_the_size_budget_is_left_out(self):
|
||||
names: Final = [f"g{index}" for index in range(2_000)]
|
||||
|
||||
assert _guardrail_response_headers({"metadata": {"applied_guardrails": names}}) == {}
|
||||
|
||||
@pytest.mark.parametrize("value", [7, "name", {"a": 1}, ["ok", 7], [None]])
|
||||
@pytest.mark.parametrize("field", ["applied_guardrails", "applied_policies", "guardrail_scan_ids"])
|
||||
def test_recorded_value_that_is_not_a_list_of_names_costs_the_header(self, field, value):
|
||||
assert _guardrail_response_headers(post_call_data={"metadata": {field: value}}) == {}
|
||||
|
||||
def test_bad_post_call_value_costs_the_header_of_the_names_recorded_before_it(self):
|
||||
assert (
|
||||
_guardrail_response_headers(
|
||||
{"metadata": {"applied_guardrails": ["pre-call-guard"]}},
|
||||
post_call_data={"metadata": {"applied_guardrails": 7}},
|
||||
)
|
||||
== {}
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"entry",
|
||||
[{"s1"}, {"at": datetime(2026, 1, 1, tzinfo=timezone.utc)}, object()],
|
||||
ids=["set", "datetime", "object"],
|
||||
)
|
||||
def test_scan_metadata_json_cannot_encode_costs_the_headers_not_the_response(self, entry):
|
||||
assert _guardrail_response_headers(post_call_data={"metadata": {"guardrail_scan_metadata": [entry]}}) == {}
|
||||
|
||||
@pytest.mark.parametrize("field", ["applied_guardrails", "guardrail_scan_ids", "guardrail_scan_metadata"])
|
||||
def test_mixed_key_entry_recorded_before_and_after_the_call_costs_the_headers_not_the_response(self, field):
|
||||
entries: Final = [{1: "a", "b": 2}]
|
||||
|
||||
assert (
|
||||
_guardrail_response_headers({"metadata": {field: entries}}, post_call_data={"metadata": {field: entries}})
|
||||
== {}
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("control", ["\x01", "\x1f", "\x7f", "\x00", "\n"])
|
||||
def test_guardrail_name_with_a_control_character_stays_out_of_the_header(self, control):
|
||||
assert _guardrail_response_headers({"metadata": {"applied_guardrails": [f"a{control}b"]}}) == {}
|
||||
|
||||
def test_guardrail_name_with_a_latin1_character_and_a_space_stays_in_the_header(self):
|
||||
assert _guardrail_response_headers({"metadata": {"applied_guardrails": ["caf\xe9 guard"]}}) == {
|
||||
"x-litellm-applied-guardrails": "caf\xe9 guard"
|
||||
}
|
||||
|
||||
def test_guardrail_headers_stay_in_budget_across_the_earlier_and_the_final_set(self):
|
||||
final_value: Final = "n" * 8000
|
||||
response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
httpx.Headers(),
|
||||
guardrail_request_data=({"metadata": {"applied_guardrails": [final_value]}},),
|
||||
earlier_guardrail_headers={"x-litellm-policy-sources": "s" * 7900, "x-litellm-applied-guardrails": "old"},
|
||||
)
|
||||
|
||||
assert response_headers == {"x-litellm-applied-guardrails": final_value}
|
||||
|
||||
def test_earlier_guardrail_header_within_the_budget_is_kept(self):
|
||||
response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
httpx.Headers(),
|
||||
guardrail_request_data=({"metadata": {"applied_guardrails": ["post-call"]}},),
|
||||
earlier_guardrail_headers={"x-litellm-policy-sources": "p=q"},
|
||||
)
|
||||
|
||||
assert response_headers == {
|
||||
"x-litellm-applied-guardrails": "post-call",
|
||||
"x-litellm-policy-sources": "p=q",
|
||||
}
|
||||
|
||||
def test_guardrail_headers_past_the_size_budget_are_left_out_so_a_client_can_read_the_response(self):
|
||||
names: Final = [f"caller-named-{index}" for index in range(10_000)]
|
||||
headers: Final = _guardrail_response_headers(
|
||||
{}, {"litellm_params": {"metadata": {"applied_guardrails": names, "applied_policies": ["tag-policy"]}}}
|
||||
)
|
||||
server: Final = h11.Connection(h11.SERVER)
|
||||
server.receive_data(b"GET / HTTP/1.1\r\nhost: proxy\r\n\r\n")
|
||||
server.next_event()
|
||||
client: Final = h11.Connection(h11.CLIENT)
|
||||
client.send(h11.Request(method="GET", target="/", headers=[("host", "proxy")]))
|
||||
client.receive_data(server.send(h11.Response(status_code=200, headers=[*headers.items()])) or b"")
|
||||
|
||||
assert headers == {"x-litellm-applied-policies": "tag-policy"}
|
||||
assert isinstance(client.next_event(), h11.Response)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue