diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e478096f7ad..e99ddb7776c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -69,6 +69,14 @@ from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + VERTEX_REQUEST_LABELS_MAX, + is_vertex_labels_header, + resolve_vertex_request_labels, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * @@ -2633,6 +2641,42 @@ def _without_dropped_anthropic_betas( return MappingProxyType({name: value for name, value in upstream_values if value is not None}) +_VERTEX_BODY_LABELED_METHODS: Final = (":generateContent", ":streamGenerateContent", ":predict") +_VERTEX_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) + + +def _is_vertex_anthropic_route(endpoint: str) -> bool: + return "/publishers/anthropic/" in f"/{endpoint}" + + +def _vertex_body_with_resolved_labels( + label_fields: tuple[str, ...], identity: Mapping[str, object], body: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_vertex_request_labels( + label_fields, body.get("labels"), {"litellm_metadata": identity}, VERTEX_REQUEST_LABELS_MAX + ) + unlabeled: Final = {name: value for name, value in body.items() if name != "labels"} + return {**unlabeled, "labels": dict(resolved)} if resolved else unlabeled + + +def _upstream_vertex_headers( + dropped_anthropic_betas: frozenset[str], + label_fields: tuple[str, ...], + identity: Mapping[str, object], + headers: Mapping[str, str], +) -> Mapping[str, str]: + without_dropped_betas: Final = _without_dropped_anthropic_betas(dropped_anthropic_betas, headers) + if not label_fields: + return without_dropped_betas + labels_header: Final = vertex_partner_labels_header( + label_fields, None, without_dropped_betas, {"litellm_metadata": identity} + ) + unlabeled: Final = { + name: value for name, value in without_dropped_betas.items() if not is_vertex_labels_header(name) + } + return MappingProxyType(unlabeled if labels_header is None else {**unlabeled, VERTEX_LABELS_HEADER: labels_header}) + + def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]: if not _is_vertex_anthropic_count_tokens_route(endpoint): return headers @@ -3081,6 +3125,19 @@ async def _base_vertex_proxy_route( if is_streaming_request: target += "?alt=sse" + label_fields: Final = vertex_request_label_fields() + identity: Final[Mapping[str, object]] = ( + get_passthrough_router_request_metadata(user_api_key_dict) if label_fields else MappingProxyType({}) + ) + if label_fields and encoded_endpoint.endswith(_VERTEX_BODY_LABELED_METHODS): + body: Final = _VERTEX_REQUEST_BODY.validate_python(await read_request_body(request)) + if body: + setattr( + request.state, + LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + _vertex_body_with_resolved_labels(label_fields, identity, body), + ) + ## CREATE PASS-THROUGH endpoint_func: Final = create_pass_through_route( endpoint=endpoint, @@ -3088,7 +3145,10 @@ async def _base_vertex_proxy_route( custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, upstream_headers_filter=partial( - _without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers) + _upstream_vertex_headers, + _vertex_dropped_anthropic_betas(general_settings, request.headers), + label_fields if _is_vertex_anthropic_route(encoded_endpoint) else (), + identity, ), ) # dynamically construct pass-through endpoint based on incoming path diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index e654c6d444b..f2f8931a38d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import base64 import importlib import json from collections.abc import Iterator @@ -555,3 +556,133 @@ def test_vertex_pass_through_forwards_a_request_without_betas_despite_a_malforme assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) assert "anthropic-beta" not in upstream.calls.last.request.headers + + +_VERTEX_GEMINI_PATH: Final = "v1/projects/caller-project/locations/global/publishers/google/models/gemini-3.8-flash" +_VERTEX_GEMINI_BODY: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} +_VERTEX_LABELED_KEY: Final = "sk-billing-svc" +_VERTEX_CALLER_LABELS: Final = {"user_api_key_team_alias": "forged", "env": "prod"} +_VERTEX_RESOLVED_LABELS: Final = { + "user_api_key_alias": "billing-svc", + "user_api_key_team_alias": "ai_platform", + "env": "prod", +} + + +def _labels_header(labels: dict[str, str]) -> str: + return base64.b64encode(json.dumps(labels).encode()).decode() + + +@pytest.fixture +def labeled_key(monkeypatch: pytest.MonkeyPatch) -> str: + from fastapi import Request + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def billing_service_auth(request: Request, api_key: str) -> UserAPIKeyAuth: + assert api_key == _VERTEX_LABELED_KEY + return UserAPIKeyAuth(api_key=api_key, key_alias="billing-svc", team_alias="AI Platform") + + monkeypatch.setattr(proxy_server, "user_custom_auth", billing_service_auth) + return _VERTEX_LABELED_KEY + + +@pytest.mark.parametrize("method", [":generateContent", ":streamGenerateContent", ":predict"]) +@pytest.mark.parametrize( + ("label_fields", "sent_labels", "upstream_labels"), + [ + pytest.param( + ["user_api_key_alias", "user_api_key_team_alias"], + _VERTEX_CALLER_LABELS, + _VERTEX_RESOLVED_LABELS, + id="identity-replaces-forged-label", + ), + pytest.param(["user_api_key_org_id"], {"user_api_key_team_alias": "forged"}, None, id="only-forged-label"), + pytest.param(None, _VERTEX_CALLER_LABELS, _VERTEX_CALLER_LABELS, id="setting-off"), + ], +) +def test_vertex_pass_through_labels_the_request_body_from_the_key_identity( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + labeled_key: str, + method: str, + label_fields: list[str] | None, + sent_labels: dict[str, str], + upstream_labels: dict[str, str] | None, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", label_fields) + upstream: Final = respx_mock.post( + url__startswith=f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_PATH}{method}" + ).mock(return_value=httpx.Response(200, json={})) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_PATH}{method}", + json={**_VERTEX_GEMINI_BODY, "labels": sent_labels}, + headers={"x-litellm-api-key": labeled_key, "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}"}, + ) + + assert response.status_code == 200 + expected_labels: Final = {} if upstream_labels is None else {"labels": upstream_labels} + assert json.loads(upstream.calls.last.request.content) == {**_VERTEX_GEMINI_BODY, **expected_labels} + + +@pytest.mark.parametrize( + ("label_fields", "caller_headers", "upstream_labels"), + [ + pytest.param( + ["user_api_key_alias", "user_api_key_team_alias"], + {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_CALLER_LABELS)}, + _VERTEX_RESOLVED_LABELS, + id="identity-replaces-forged-header", + ), + pytest.param( + ["user_api_key_alias", "user_api_key_team_alias"], + {}, + {"user_api_key_alias": "billing-svc", "user_api_key_team_alias": "ai_platform"}, + id="identity-without-caller-header", + ), + pytest.param( + ["user_api_key_org_id"], + {"x-pass-x-vertex-ai-labels": _labels_header({"user_api_key_team_alias": "forged"})}, + None, + id="only-forged-header", + ), + pytest.param( + None, + {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_CALLER_LABELS)}, + _VERTEX_CALLER_LABELS, + id="setting-off", + ), + ], +) +def test_vertex_pass_through_labels_the_claude_request_header_from_the_key_identity( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + labeled_key: str, + label_fields: list[str] | None, + caller_headers: dict[str, str], + upstream_labels: dict[str, str] | None, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", label_fields) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_CALLER_PATH}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_CALLER_PATH}", + json=_VERTEX_CLAUDE_BODY, + headers={ + "x-litellm-api-key": labeled_key, + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + **caller_headers, + }, + ) + + assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) + outbound: Final = upstream.calls.last.request + sent_headers: Final = tuple(json.loads(base64.b64decode(value)) for value in outbound.headers.get_list("x-vertex-ai-labels")) + assert sent_headers == (() if upstream_labels is None else (upstream_labels,)) + assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY