mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge af2cf4d706 into 7f1b0a4633
This commit is contained in:
commit
6c76e27f6a
4 changed files with 790 additions and 47 deletions
|
|
@ -1,12 +1,14 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
from collections.abc import Mapping
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import request_timeout
|
||||
|
||||
# Import the adapter for fallback to completion format
|
||||
|
|
@ -40,6 +42,7 @@ else:
|
|||
# Initialize any necessary instances or variables here
|
||||
base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
_GENERATE_CONTENT_RESULT: Final = TypeAdapter(object)
|
||||
_FORWARDED_HEADERS: Final = TypeAdapter[Mapping[str, str]](Mapping[str, str])
|
||||
#################################################
|
||||
|
||||
|
||||
|
|
@ -236,6 +239,30 @@ def _merge_native_request_fields(
|
|||
return {**native_request_fields, **(extra_body or {})}
|
||||
|
||||
|
||||
def _validated_client_headers(client_headers: object) -> Mapping[str, str] | None:
|
||||
if client_headers is None:
|
||||
return None
|
||||
try:
|
||||
return _FORWARDED_HEADERS.validate_python(client_headers)
|
||||
except ValidationError as exc:
|
||||
verbose_logger.debug("Ignoring forwarded `headers` that are not a str-to-str mapping: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _merge_client_forwarded_headers(
|
||||
extra_headers: Mapping[str, object] | None,
|
||||
client_headers: object,
|
||||
) -> dict[str, object] | None: # mutable-ok: downstream handlers expect dict | None
|
||||
forwarded: Final = _validated_client_headers(client_headers)
|
||||
if not forwarded:
|
||||
return dict(extra_headers) if extra_headers is not None else None
|
||||
explicit_names: Final = frozenset(name.lower() for name in (extra_headers or ()))
|
||||
return {
|
||||
**{name: value for name, value in forwarded.items() if name.lower() not in explicit_names},
|
||||
**(extra_headers or {}),
|
||||
}
|
||||
|
||||
|
||||
@client
|
||||
async def agenerate_content(
|
||||
model: str,
|
||||
|
|
@ -349,6 +376,7 @@ def generate_content(
|
|||
|
||||
# Extract systemInstruction from kwargs to pass to handler
|
||||
system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
|
||||
merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers"))
|
||||
|
||||
# Check if we should use the adapter (when provider config is None)
|
||||
if setup_result.generate_content_provider_config is None:
|
||||
|
|
@ -360,7 +388,7 @@ def generate_content(
|
|||
tools=tools,
|
||||
_is_async=_is_async,
|
||||
litellm_params=setup_result.litellm_params,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -375,7 +403,7 @@ def generate_content(
|
|||
custom_llm_provider=setup_result.custom_llm_provider,
|
||||
litellm_params=setup_result.litellm_params,
|
||||
logging_obj=setup_result.litellm_logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
|
||||
timeout=timeout or request_timeout,
|
||||
_is_async=_is_async,
|
||||
|
|
@ -442,6 +470,7 @@ async def agenerate_content_stream(
|
|||
|
||||
# Extract systemInstruction from kwargs to pass to handler
|
||||
system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
|
||||
merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers"))
|
||||
|
||||
# Check if we should use the adapter (when provider config is None)
|
||||
if setup_result.generate_content_provider_config is None:
|
||||
|
|
@ -456,7 +485,7 @@ async def agenerate_content_stream(
|
|||
litellm_params=setup_result.litellm_params,
|
||||
tools=tools,
|
||||
stream=True,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -472,7 +501,7 @@ async def agenerate_content_stream(
|
|||
custom_llm_provider=setup_result.custom_llm_provider,
|
||||
litellm_params=setup_result.litellm_params,
|
||||
logging_obj=setup_result.litellm_logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
|
||||
timeout=timeout or request_timeout,
|
||||
_is_async=True,
|
||||
|
|
@ -534,6 +563,7 @@ def generate_content_stream(
|
|||
|
||||
# Extract systemInstruction from kwargs to pass to handler
|
||||
system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
|
||||
merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers"))
|
||||
|
||||
# Check if we should use the adapter (when provider config is None)
|
||||
if setup_result.generate_content_provider_config is None:
|
||||
|
|
@ -548,7 +578,7 @@ def generate_content_stream(
|
|||
_is_async=_is_async,
|
||||
litellm_params=setup_result.litellm_params,
|
||||
stream=True,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -563,7 +593,7 @@ def generate_content_stream(
|
|||
custom_llm_provider=setup_result.custom_llm_provider,
|
||||
litellm_params=setup_result.litellm_params,
|
||||
logging_obj=setup_result.litellm_logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
extra_headers=merged_extra_headers,
|
||||
extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
|
||||
timeout=timeout or request_timeout,
|
||||
_is_async=_is_async,
|
||||
|
|
|
|||
|
|
@ -1498,12 +1498,29 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
|||
)
|
||||
|
||||
|
||||
def litellm_auth_header_name(headers: Headers) -> str:
|
||||
if "x-litellm-api-key" in headers:
|
||||
return "x-litellm-api-key"
|
||||
if "authorization" in headers:
|
||||
return "authorization"
|
||||
return "x-api-key"
|
||||
_PROXY_AUTH_HEADER_PRECEDENCE: Final = tuple(
|
||||
header.value.lower()
|
||||
for header in (
|
||||
SpecialHeaders.custom_litellm_api_key,
|
||||
SpecialHeaders.openai_authorization,
|
||||
SpecialHeaders.azure_authorization,
|
||||
SpecialHeaders.anthropic_authorization,
|
||||
SpecialHeaders.google_ai_studio_authorization,
|
||||
SpecialHeaders.azure_apim_authorization,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _strip_bearer(value: str) -> str:
|
||||
stripped: Final = value.strip()
|
||||
return stripped[7:].strip() if stripped[:7].lower() == "bearer " else stripped
|
||||
|
||||
|
||||
def litellm_auth_header_name(headers: Mapping[str, str], litellm_key_header_name: str | None = None) -> str:
|
||||
present: Final = {name.lower() for name, value in headers.items() if value}
|
||||
if litellm_key_header_name is not None and litellm_key_header_name.lower() in present:
|
||||
return litellm_key_header_name.lower()
|
||||
return next((name for name in _PROXY_AUTH_HEADER_PRECEDENCE if name in present), "x-api-key")
|
||||
|
||||
|
||||
def forwarded_anthropic_oauth_token(headers: Headers) -> str | None:
|
||||
|
|
@ -1521,11 +1538,11 @@ def forwarded_anthropic_oauth_token(headers: Headers) -> str | None:
|
|||
|
||||
|
||||
def clean_headers(
|
||||
headers: Headers,
|
||||
headers: Mapping[str, str],
|
||||
litellm_key_header_name: str | None = None,
|
||||
forward_llm_provider_auth_headers: bool = False,
|
||||
authenticated_with_header: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
) -> dict[str, str]: # mutable-ok: downstream callers mutate returned headers dict
|
||||
"""
|
||||
Removes litellm api key from headers
|
||||
|
||||
|
|
@ -1542,36 +1559,35 @@ def clean_headers(
|
|||
from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key
|
||||
from litellm.llms.chatgpt.common_utils import is_chatgpt_oauth_key
|
||||
|
||||
clean_headers: Final = {}
|
||||
litellm_key_lower: Final = litellm_key_header_name.lower() if litellm_key_header_name is not None else None
|
||||
for header, value in headers.items():
|
||||
header_lower = header.lower()
|
||||
auth_header_lower: Final = authenticated_with_header.lower() if authenticated_with_header is not None else None
|
||||
custom_excluded: Final = frozenset({litellm_key_lower, auth_header_lower})
|
||||
provider_excluded: Final = frozenset({"authorization", "x-litellm-api-key", *custom_excluded})
|
||||
proxy_auth_credential: Final = (
|
||||
next(
|
||||
(
|
||||
credential
|
||||
for name, raw in headers.items()
|
||||
if name.lower() == auth_header_lower and (credential := _strip_bearer(raw))
|
||||
),
|
||||
None,
|
||||
)
|
||||
if auth_header_lower is not None
|
||||
else None
|
||||
)
|
||||
|
||||
def _keep(header: str, value: str) -> bool:
|
||||
header_lower: Final = header.lower()
|
||||
if header_lower == "authorization" and (is_anthropic_oauth_key(value) or is_chatgpt_oauth_key(value)):
|
||||
if authenticated_with_header is None or authenticated_with_header.lower() != "authorization":
|
||||
clean_headers[header] = value
|
||||
continue
|
||||
# Special handling for x-api-key: forward it based on authenticated_with_header
|
||||
elif header_lower == "x-api-key":
|
||||
if forward_llm_provider_auth_headers and (
|
||||
authenticated_with_header is None or authenticated_with_header.lower() != "x-api-key"
|
||||
):
|
||||
clean_headers[header] = value
|
||||
elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE:
|
||||
if litellm_key_lower and header_lower == litellm_key_lower:
|
||||
continue
|
||||
if header_lower == "authorization":
|
||||
continue
|
||||
# Never forward x-litellm-api-key (it's for proxy auth only)
|
||||
if header_lower == "x-litellm-api-key":
|
||||
continue
|
||||
clean_headers[header] = value
|
||||
# Check if header should be excluded: either in special headers cache or matches custom litellm key
|
||||
elif header_lower not in _SPECIAL_HEADERS_CACHE and (
|
||||
litellm_key_lower is None or header_lower != litellm_key_lower
|
||||
):
|
||||
clean_headers[header] = value
|
||||
return clean_headers
|
||||
return auth_header_lower != "authorization" and _strip_bearer(value) != proxy_auth_credential
|
||||
if forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE:
|
||||
if header_lower in provider_excluded:
|
||||
return False
|
||||
credential: Final = _strip_bearer(value)
|
||||
return bool(credential) and credential != proxy_auth_credential
|
||||
return header_lower not in _SPECIAL_HEADERS_CACHE and header_lower not in custom_excluded
|
||||
|
||||
return {header: value for header, value in headers.items() if _keep(header, value)}
|
||||
|
||||
|
||||
def _request_fields(value: object) -> Mapping[str, object]:
|
||||
|
|
@ -2669,13 +2685,14 @@ async def add_litellm_data_to_request(
|
|||
forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False)
|
||||
if not forward_llm_auth:
|
||||
forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False)
|
||||
litellm_key_header_name: Final = (
|
||||
general_settings.get("litellm_key_header_name") if general_settings is not None else None
|
||||
)
|
||||
_headers: Final[dict[str, str]] = clean_headers(
|
||||
request.headers,
|
||||
litellm_key_header_name=(
|
||||
general_settings.get("litellm_key_header_name") if general_settings is not None else None
|
||||
),
|
||||
litellm_key_header_name=litellm_key_header_name,
|
||||
forward_llm_provider_auth_headers=forward_llm_auth,
|
||||
authenticated_with_header=litellm_auth_header_name(request.headers),
|
||||
authenticated_with_header=litellm_auth_header_name(request.headers, litellm_key_header_name),
|
||||
)
|
||||
_strip_untrusted_request_header_controls(
|
||||
_headers,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Test to verify the Google GenAI generate_content adapter functionality
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -382,3 +382,447 @@ async def test_vertex_generate_content_spend_follows_the_deployment_location(
|
|||
recorded: Final = await _logged_vertex_generate_content_cost(location, stream=stream)
|
||||
|
||||
assert recorded == pytest.approx(expected), (location, stream)
|
||||
|
||||
|
||||
class _WireCapture:
|
||||
def __init__(self, url: str) -> None:
|
||||
self.url: Final = url
|
||||
self._items: Final[list[tuple[str, Mapping[str, str]]]] = [] # mutable-ok: test wire log
|
||||
|
||||
def record(self, path: str, headers: Mapping[str, str]) -> None:
|
||||
self._items.append((path, headers))
|
||||
|
||||
@property
|
||||
def requests(self) -> tuple[tuple[str, Mapping[str, str]], ...]:
|
||||
return tuple(self._items)
|
||||
|
||||
|
||||
def _mock_wire_reply(path: str, stream_requested: bool) -> tuple[bytes, str]:
|
||||
if "streamGenerateContent" in path:
|
||||
return (
|
||||
b'data: {"candidates":[{"content":{"role":"model","parts":[{"text":"pong"}]},'
|
||||
b'"finishReason":"STOP"}],'
|
||||
b'"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":1,"totalTokenCount":5}}\n\n',
|
||||
"text/event-stream",
|
||||
)
|
||||
if ":generateContent" in path:
|
||||
return (
|
||||
b'{"candidates":[{"content":{"role":"model","parts":[{"text":"pong"}]},'
|
||||
b'"finishReason":"STOP"}],'
|
||||
b'"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":1,"totalTokenCount":5}}',
|
||||
"application/json",
|
||||
)
|
||||
if stream_requested:
|
||||
return (
|
||||
b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,'
|
||||
b'"model":"gpt-4o-mini","choices":[{"index":0,'
|
||||
b'"delta":{"role":"assistant","content":"pong"},"finish_reason":"stop"}]}\n\n'
|
||||
b"data: [DONE]\n\n",
|
||||
"text/event-stream",
|
||||
)
|
||||
return (
|
||||
b'{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-4o-mini",'
|
||||
b'"choices":[{"index":0,"message":{"role":"assistant","content":"pong"},'
|
||||
b'"finish_reason":"stop"}],'
|
||||
b'"usage":{"prompt_tokens":4,"completion_tokens":1,"total_tokens":5}}',
|
||||
"application/json",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def header_capture_wire() -> Iterator[_WireCapture]:
|
||||
"""
|
||||
Loopback upstream that records outbound wire headers.
|
||||
|
||||
A real socket (allowed by ``tests/unit/conftest.py`` for loopback hosts) is used instead of
|
||||
``respx`` because async calls go through litellm's aiohttp transport, which ``respx`` cannot intercept.
|
||||
"""
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
holder: Final[list[_WireCapture]] = [] # mutable-ok: single-slot fixture bootstrap reference
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
length: Final = int(self.headers.get("Content-Length", "0"))
|
||||
raw_body: Final = self.rfile.read(length)
|
||||
holder[0].record(self.path, {k.lower(): v for k, v in self.headers.items()})
|
||||
payload: Final = json.loads(raw_body.decode("utf-8")) if raw_body else {}
|
||||
body, content_type = _mock_wire_reply(self.path, bool(payload.get("stream")))
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", content_type)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return
|
||||
|
||||
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
capture: Final = _WireCapture(f"http://127.0.0.1:{server.server_address[1]}")
|
||||
holder.append(capture)
|
||||
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield capture
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
def _extract_chunk_texts(payload: object) -> tuple[str, ...]:
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
obj_map: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object])
|
||||
map_seq: Final = TypeAdapter[tuple[Mapping[str, object], ...]](tuple[Mapping[str, object], ...])
|
||||
parsed: Final = obj_map.validate_python(payload.model_dump() if hasattr(payload, "model_dump") else payload)
|
||||
candidates: Final = map_seq.validate_python(parsed.get("candidates", ()))
|
||||
if not candidates:
|
||||
return ()
|
||||
content: Final = obj_map.validate_python(candidates[0].get("content", {}))
|
||||
parts: Final = map_seq.validate_python(content.get("parts", ()))
|
||||
return tuple(str(part["text"]) for part in parts if isinstance(part.get("text"), str) and part["text"])
|
||||
|
||||
|
||||
def _sse_payload_texts(line: str) -> tuple[str, ...]:
|
||||
if not line.startswith("data: "):
|
||||
return ()
|
||||
payload: Final = line[6:].strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
return ()
|
||||
return _extract_chunk_texts(json.loads(payload))
|
||||
|
||||
|
||||
def _extract_sse_texts(chunks: tuple[bytes, ...]) -> tuple[str, ...]:
|
||||
from itertools import chain
|
||||
|
||||
lines: Final = chain.from_iterable(raw.decode("utf-8").splitlines() for raw in chunks)
|
||||
return tuple(chain.from_iterable(_sse_payload_texts(line) for line in lines))
|
||||
|
||||
|
||||
_UPSTREAM_KEY: Final = "upstream-secret"
|
||||
|
||||
|
||||
def _target_route(target: str, wire: _WireCapture) -> tuple[str, str]:
|
||||
if target == "native":
|
||||
return "gemini/gemini-2.0-flash", wire.url
|
||||
return "openai/gpt-4o-mini", f"{wire.url}/v1"
|
||||
|
||||
|
||||
async def _invoke_generate_content(
|
||||
entrypoint: str,
|
||||
*,
|
||||
model: str,
|
||||
api_base: str,
|
||||
headers: Mapping[str, str],
|
||||
extra_headers: Mapping[str, str],
|
||||
) -> tuple[str, ...]:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.google_genai.main import (
|
||||
agenerate_content,
|
||||
agenerate_content_stream,
|
||||
generate_content,
|
||||
generate_content_stream,
|
||||
)
|
||||
|
||||
contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}]
|
||||
if entrypoint == "sync":
|
||||
return _extract_chunk_texts(
|
||||
generate_content(
|
||||
model=model,
|
||||
contents=contents,
|
||||
api_base=api_base,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=headers,
|
||||
extra_headers=dict(extra_headers),
|
||||
)
|
||||
)
|
||||
if entrypoint == "async":
|
||||
return _extract_chunk_texts(
|
||||
await agenerate_content(
|
||||
model=model,
|
||||
contents=contents,
|
||||
api_base=api_base,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=headers,
|
||||
extra_headers=dict(extra_headers),
|
||||
)
|
||||
)
|
||||
stream: Final = (
|
||||
generate_content_stream(
|
||||
model=model,
|
||||
contents=contents,
|
||||
api_base=api_base,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=headers,
|
||||
extra_headers=dict(extra_headers),
|
||||
)
|
||||
if entrypoint == "sync_stream"
|
||||
else await agenerate_content_stream(
|
||||
model=model,
|
||||
contents=contents,
|
||||
api_base=api_base,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=headers,
|
||||
extra_headers=dict(extra_headers),
|
||||
)
|
||||
)
|
||||
as_bytes: Final = TypeAdapter(bytes)
|
||||
# native sync streams implement both protocols; adapter sync streams are async-only
|
||||
if entrypoint == "sync_stream" and isinstance(stream, Iterator):
|
||||
return _extract_sse_texts(tuple(as_bytes.validate_python(chunk) for chunk in stream))
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
return _extract_sse_texts(tuple([as_bytes.validate_python(chunk) async for chunk in stream]))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("entrypoint", ["sync", "async", "sync_stream", "async_stream"])
|
||||
@pytest.mark.parametrize("target", ["native", "adapter"])
|
||||
async def test_generate_content_forwards_client_headers_behind_explicit_extra_headers(
|
||||
header_capture_wire: _WireCapture,
|
||||
target: str,
|
||||
entrypoint: str,
|
||||
) -> None:
|
||||
model, api_base = _target_route(target, header_capture_wire)
|
||||
forwarded: Final = {"x-team-id": "platform", "x-custom": "from-proxy", "anthropic-beta": "context-1m-2025-08-07"}
|
||||
explicit: Final = {"X-Custom": "from-caller"}
|
||||
|
||||
texts: Final = await _invoke_generate_content(
|
||||
entrypoint, model=model, api_base=api_base, headers=forwarded, extra_headers=explicit
|
||||
)
|
||||
|
||||
assert texts == ("pong",)
|
||||
assert forwarded == {"x-team-id": "platform", "x-custom": "from-proxy", "anthropic-beta": "context-1m-2025-08-07"}
|
||||
assert explicit == {"X-Custom": "from-caller"}
|
||||
assert len(header_capture_wire.requests) == 1
|
||||
wire_headers: Final = header_capture_wire.requests[0][1]
|
||||
assert wire_headers["x-team-id"] == "platform"
|
||||
assert wire_headers["anthropic-beta"] == "context-1m-2025-08-07"
|
||||
assert wire_headers["x-custom"] == "from-caller"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_headers, client_headers, expected",
|
||||
[
|
||||
(None, None, None),
|
||||
(None, {}, None),
|
||||
({}, None, {}),
|
||||
({"x-b": "2"}, None, {"x-b": "2"}),
|
||||
(None, {"x-a": "1"}, {"x-a": "1"}),
|
||||
({}, {"x-a": "1"}, {"x-a": "1"}),
|
||||
({"X-Custom": "from-caller"}, {"x-custom": "from-proxy", "x-a": "1"}, {"x-a": "1", "X-Custom": "from-caller"}),
|
||||
({"x-b": "2"}, "not-a-mapping", {"x-b": "2"}),
|
||||
(None, "not-a-mapping", None),
|
||||
({"x-b": "2"}, {"x-a": 1}, {"x-b": "2"}),
|
||||
],
|
||||
ids=[
|
||||
"nothing",
|
||||
"empty-forwarded",
|
||||
"empty-explicit",
|
||||
"explicit-only",
|
||||
"forwarded-only",
|
||||
"empty-explicit-with-forwarded",
|
||||
"case-insensitive-precedence",
|
||||
"non-mapping-forwarded-kept-explicit",
|
||||
"non-mapping-forwarded",
|
||||
"non-str-value-forwarded",
|
||||
],
|
||||
)
|
||||
def test_merge_client_forwarded_headers(
|
||||
extra_headers: Mapping[str, str] | None,
|
||||
client_headers: object,
|
||||
expected: Mapping[str, str] | None,
|
||||
) -> None:
|
||||
from litellm.google_genai.main import _merge_client_forwarded_headers
|
||||
|
||||
merged: Final = _merge_client_forwarded_headers(extra_headers, client_headers)
|
||||
|
||||
assert merged == expected
|
||||
assert merged is None or (merged is not extra_headers and merged is not client_headers)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("target", ["native", "adapter"])
|
||||
async def test_router_deployment_headers_reach_generate_content(
|
||||
header_capture_wire: _WireCapture,
|
||||
target: str,
|
||||
) -> None:
|
||||
from litellm import Router
|
||||
|
||||
model, api_base = _target_route(target, header_capture_wire)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "genai-deployment",
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_base": api_base,
|
||||
"api_key": _UPSTREAM_KEY,
|
||||
"headers": {"x-deployment-tag": "blue"},
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
resp: Final = await router.agenerate_content(
|
||||
model="genai-deployment",
|
||||
contents=[{"role": "user", "parts": [{"text": "ping"}]}],
|
||||
)
|
||||
|
||||
assert _extract_chunk_texts(resp) == ("pong",)
|
||||
assert len(header_capture_wire.requests) == 1
|
||||
assert header_capture_wire.requests[0][1]["x-deployment-tag"] == "blue"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_proxy_auth(
|
||||
header_capture_wire: _WireCapture,
|
||||
) -> None:
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
from litellm.google_genai.main import agenerate_content
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
clean_headers,
|
||||
litellm_auth_header_name,
|
||||
)
|
||||
|
||||
incoming_headers: Final = Headers(
|
||||
{
|
||||
"x-goog-api-key": "sk-litellm-virtual-key",
|
||||
"authorization": "Bearer sk-litellm-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
"user-agent": "google-genai-sdk/1.0",
|
||||
}
|
||||
)
|
||||
cleaned: Final = clean_headers(
|
||||
incoming_headers,
|
||||
litellm_key_header_name=None,
|
||||
forward_llm_provider_auth_headers=False,
|
||||
authenticated_with_header=litellm_auth_header_name(incoming_headers),
|
||||
)
|
||||
backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
||||
headers=cleaned,
|
||||
request_data={"model": "gemini/gemini-2.0-flash"},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"),
|
||||
general_settings={"forward_client_headers_to_llm_api": True},
|
||||
)
|
||||
|
||||
resp: Final = await agenerate_content(
|
||||
model="gemini/gemini-2.0-flash",
|
||||
contents=[{"role": "user", "parts": [{"text": "ping"}]}],
|
||||
api_base=header_capture_wire.url,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=backend_data.get("headers"),
|
||||
)
|
||||
|
||||
assert _extract_chunk_texts(resp) == ("pong",)
|
||||
assert len(header_capture_wire.requests) == 1
|
||||
wire_headers: Final = header_capture_wire.requests[0][1]
|
||||
assert wire_headers["x-goog-api-key"] == _UPSTREAM_KEY
|
||||
assert wire_headers["x-cost-center"] == "core-platform"
|
||||
assert "authorization" not in wire_headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("incoming", "expected_upstream_key"),
|
||||
[
|
||||
({"x-goog-api-key": "sk-litellm-virtual-key", "x-cost-center": "core-platform"}, _UPSTREAM_KEY),
|
||||
(
|
||||
{
|
||||
"authorization": "Bearer sk-litellm-virtual-key",
|
||||
"x-goog-api-key": "sk-litellm-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
_UPSTREAM_KEY,
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-api-key": "",
|
||||
"x-goog-api-key": "sk-litellm-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
_UPSTREAM_KEY,
|
||||
),
|
||||
(
|
||||
{
|
||||
"authorization": "Bearer sk-litellm-virtual-key",
|
||||
"x-goog-api-key": "",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
_UPSTREAM_KEY,
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-litellm-api-key": "sk-litellm-virtual-key",
|
||||
"x-goog-api-key": "AIza-caller-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
"AIza-caller-key",
|
||||
),
|
||||
(
|
||||
{
|
||||
"authorization": "Bearer sk-litellm-virtual-key",
|
||||
"x-goog-api-key": "AIza-caller-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
"AIza-caller-key",
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"proxy-key-in-goog-header",
|
||||
"sdk-dual-bearer-and-goog-proxy-key",
|
||||
"empty-higher-precedence-header",
|
||||
"empty-goog-header-beside-bearer-key",
|
||||
"caller-key-beside-litellm-key",
|
||||
"caller-key-beside-bearer-key",
|
||||
],
|
||||
)
|
||||
async def test_provider_auth_forwarding_never_sends_proxy_key_to_generate_content(
|
||||
header_capture_wire: _WireCapture,
|
||||
incoming: Mapping[str, str],
|
||||
expected_upstream_key: str,
|
||||
) -> None:
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
from litellm.google_genai.main import agenerate_content
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
clean_headers,
|
||||
litellm_auth_header_name,
|
||||
)
|
||||
|
||||
incoming_headers: Final = Headers(incoming)
|
||||
cleaned: Final = clean_headers(
|
||||
incoming_headers,
|
||||
forward_llm_provider_auth_headers=True,
|
||||
authenticated_with_header=litellm_auth_header_name(incoming_headers),
|
||||
)
|
||||
backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
||||
headers=cleaned,
|
||||
request_data={"model": "gemini/gemini-2.0-flash"},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"),
|
||||
general_settings={"forward_client_headers_to_llm_api": True},
|
||||
)
|
||||
|
||||
resp: Final = await agenerate_content(
|
||||
model="gemini/gemini-2.0-flash",
|
||||
contents=[{"role": "user", "parts": [{"text": "ping"}]}],
|
||||
api_base=header_capture_wire.url,
|
||||
api_key=_UPSTREAM_KEY,
|
||||
headers=backend_data.get("headers"),
|
||||
)
|
||||
|
||||
assert _extract_chunk_texts(resp) == ("pong",)
|
||||
assert len(header_capture_wire.requests) == 1
|
||||
wire_headers: Final = header_capture_wire.requests[0][1]
|
||||
assert wire_headers["x-goog-api-key"] == expected_upstream_key
|
||||
assert wire_headers["x-cost-center"] == "core-platform"
|
||||
assert "x-litellm-api-key" not in wire_headers
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import copy
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
|
@ -58,6 +59,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
forwarded_anthropic_oauth_token,
|
||||
get_dynamic_logging_metadata,
|
||||
get_metadata_variable_name,
|
||||
litellm_auth_header_name,
|
||||
move_guardrails_to_metadata,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
|
|
@ -6447,6 +6449,152 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth
|
|||
assert "x-api-key" not in result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "litellm_key_header_name", "expected"),
|
||||
[
|
||||
({"x-litellm-api-key": "sk-proxy", "authorization": "Bearer sk-other"}, None, "x-litellm-api-key"),
|
||||
({"x-litellm-api-key": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"),
|
||||
({"authorization": "Bearer sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "authorization"),
|
||||
({"authorization": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"),
|
||||
({"api-key": "sk-proxy", "x-api-key": "sk-ant-caller"}, None, "api-key"),
|
||||
({"x-api-key": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"),
|
||||
({"x-api-key": "sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "x-api-key"),
|
||||
({"x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"),
|
||||
({"X-Goog-Api-Key": "sk-proxy"}, None, "x-goog-api-key"),
|
||||
({"ocp-apim-subscription-key": "sk-proxy"}, None, "ocp-apim-subscription-key"),
|
||||
({"X-Custom-Key": "sk-proxy", "authorization": "Bearer sk-other"}, "x-custom-key", "x-custom-key"),
|
||||
({"x-custom-key": "", "x-goog-api-key": "sk-proxy"}, "x-custom-key", "x-goog-api-key"),
|
||||
({"authorization": "Bearer sk-proxy"}, "x-custom-key", "authorization"),
|
||||
({"content-type": "application/json"}, None, "x-api-key"),
|
||||
],
|
||||
ids=[
|
||||
"litellm-key-beats-authorization",
|
||||
"empty-litellm-key-skipped",
|
||||
"authorization-beats-goog",
|
||||
"empty-authorization-skipped",
|
||||
"azure-beats-anthropic",
|
||||
"empty-anthropic-skipped",
|
||||
"anthropic-beats-goog",
|
||||
"goog",
|
||||
"goog-mixed-case-dict",
|
||||
"azure-apim",
|
||||
"custom-header-mixed-case-dict",
|
||||
"empty-custom-header-skipped",
|
||||
"custom-header-absent",
|
||||
"no-auth-header",
|
||||
],
|
||||
)
|
||||
def test_litellm_auth_header_name_follows_user_api_key_auth_precedence(
|
||||
headers: Mapping[str, str], litellm_key_header_name: str | None, expected: str
|
||||
) -> None:
|
||||
assert litellm_auth_header_name(headers, litellm_key_header_name) == expected
|
||||
assert litellm_auth_header_name(Headers(headers), litellm_key_header_name) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "expected"),
|
||||
[
|
||||
(
|
||||
{"x-goog-api-key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"},
|
||||
{"x-cost-center": "core-platform"},
|
||||
),
|
||||
(
|
||||
{"X-Goog-Api-Key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"},
|
||||
{"x-cost-center": "core-platform"},
|
||||
),
|
||||
(
|
||||
{
|
||||
"authorization": "Bearer sk-proxy-virtual-key",
|
||||
"x-goog-api-key": "sk-proxy-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
{"x-cost-center": "core-platform"},
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-litellm-api-key": "sk-proxy-virtual-key",
|
||||
"x-goog-api-key": "sk-proxy-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
{"x-cost-center": "core-platform"},
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-api-key": "",
|
||||
"x-goog-api-key": "sk-proxy-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
{"x-cost-center": "core-platform"},
|
||||
),
|
||||
(
|
||||
{"x-litellm-api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"},
|
||||
{"x-goog-api-key": "AIza-caller-key"},
|
||||
),
|
||||
(
|
||||
{"authorization": "Bearer sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"},
|
||||
{"x-goog-api-key": "AIza-caller-key"},
|
||||
),
|
||||
(
|
||||
{"api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"},
|
||||
{"x-goog-api-key": "AIza-caller-key"},
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"goog-was-proxy-auth",
|
||||
"goog-mixed-case-was-proxy-auth",
|
||||
"duplicate-proxy-key-in-bearer-and-goog",
|
||||
"duplicate-proxy-key-in-litellm-and-goog",
|
||||
"empty-provider-header-and-goog-proxy-auth",
|
||||
"goog-beside-litellm-key",
|
||||
"goog-beside-bearer-proxy-key",
|
||||
"goog-beside-azure-proxy-key",
|
||||
],
|
||||
)
|
||||
def test_clean_headers_byok_strips_only_the_header_used_for_proxy_auth(
|
||||
headers: Mapping[str, str], expected: Mapping[str, str]
|
||||
) -> None:
|
||||
assert (
|
||||
clean_headers(
|
||||
headers,
|
||||
forward_llm_provider_auth_headers=True,
|
||||
authenticated_with_header=litellm_auth_header_name(headers),
|
||||
)
|
||||
== expected
|
||||
)
|
||||
incoming: Final = Headers(headers)
|
||||
assert (
|
||||
clean_headers(
|
||||
incoming,
|
||||
forward_llm_provider_auth_headers=True,
|
||||
authenticated_with_header=litellm_auth_header_name(incoming),
|
||||
)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
def test_clean_headers_strips_authenticated_custom_header_without_litellm_key_name() -> None:
|
||||
cleaned: Final = clean_headers(
|
||||
{"X-Custom-Proxy-Key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"},
|
||||
forward_llm_provider_auth_headers=True,
|
||||
authenticated_with_header="x-custom-proxy-key",
|
||||
)
|
||||
|
||||
assert cleaned == {"x-cost-center": "core-platform"}
|
||||
|
||||
|
||||
def test_clean_headers_byok_strips_custom_litellm_key_header_named_x_api_key() -> None:
|
||||
incoming: Final = Headers({"x-api-key": "sk-proxy-virtual-key", "authorization": "Bearer sk-ant-api03-caller"})
|
||||
|
||||
result: Final = clean_headers(
|
||||
incoming,
|
||||
litellm_key_header_name="x-api-key",
|
||||
forward_llm_provider_auth_headers=True,
|
||||
authenticated_with_header=litellm_auth_header_name(incoming, "x-api-key"),
|
||||
)
|
||||
|
||||
assert "x-api-key" not in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Team guardrail + global policy regression tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -7965,6 +8113,110 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_
|
|||
assert value not in logged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_headers", "general_settings", "expected_forwarded_auth"),
|
||||
[
|
||||
(
|
||||
{"x-goog-api-key": "sk-virtual-key", "x-cost-center": "core-platform"},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{},
|
||||
),
|
||||
(
|
||||
{"X-Goog-Api-Key": "sk-virtual-key", "x-cost-center": "core-platform"},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{},
|
||||
),
|
||||
(
|
||||
{
|
||||
"authorization": "Bearer sk-virtual-key",
|
||||
"x-goog-api-key": "sk-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{},
|
||||
),
|
||||
(
|
||||
{
|
||||
"Authorization": "Bearer sk-virtual-key",
|
||||
"X-Goog-Api-Key": "sk-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{},
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-api-key": "",
|
||||
"x-goog-api-key": "sk-virtual-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{},
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-litellm-api-key": "Bearer sk-virtual-key",
|
||||
"x-goog-api-key": "AIza-caller-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
FORWARD_PROVIDER_AUTH_HEADERS,
|
||||
{"x-goog-api-key": "AIza-caller-key"},
|
||||
),
|
||||
(
|
||||
{"x-api-key": "sk-virtual-key", "authorization": "Bearer sk-other", "x-cost-center": "core-platform"},
|
||||
{**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-api-key"},
|
||||
{},
|
||||
),
|
||||
(
|
||||
{
|
||||
"x-custom-proxy-key": "sk-virtual-key",
|
||||
"x-goog-api-key": "AIza-caller-key",
|
||||
"x-cost-center": "core-platform",
|
||||
},
|
||||
{**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-custom-proxy-key"},
|
||||
{"x-goog-api-key": "AIza-caller-key"},
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"goog-was-proxy-auth",
|
||||
"goog-mixed-case-was-proxy-auth",
|
||||
"duplicate-proxy-key-in-bearer-and-goog",
|
||||
"duplicate-proxy-key-mixed-case",
|
||||
"empty-higher-precedence-header",
|
||||
"goog-beside-litellm-key",
|
||||
"custom-litellm-key-header",
|
||||
"custom-litellm-key-header-preserves-goog-byok",
|
||||
],
|
||||
)
|
||||
async def test_add_litellm_data_to_request_never_forwards_the_header_used_for_proxy_auth(
|
||||
request_headers: Mapping[str, str],
|
||||
general_settings: Mapping[str, object],
|
||||
expected_forwarded_auth: Mapping[str, str],
|
||||
) -> None:
|
||||
request_mock: Final = _make_request_mock(
|
||||
"/v1/chat/completions", {"Content-Type": "application/json", **request_headers}
|
||||
)
|
||||
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data={"model": "gemini-2.0-flash", "messages": [{"role": "user", "content": "hello"}]},
|
||||
request=request_mock,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings=dict(general_settings),
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
forwarded_auth: Final = {
|
||||
name: value
|
||||
for name, value in updated["headers"].items()
|
||||
if name.lower() in {"x-goog-api-key", "x-api-key", "x-litellm-api-key"}
|
||||
}
|
||||
assert forwarded_auth == expected_forwarded_auth
|
||||
assert updated["headers"]["x-cost-center"] == "core-platform"
|
||||
assert "sk-virtual-key" not in str(updated.get("api_key"))
|
||||
|
||||
|
||||
_SERVICE_ACCOUNT_TOKEN_REFERENCE: Final = "oidc/file//var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue