From 60ce0a06116ecb268e8833edc2ac9ed4e7e4cf38 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 28 Oct 2026 08:41:50 -0700 Subject: [PATCH 01/13] fix(vertex_ai): send Vertex Claude labels as X-Vertex-AI-Labels header instead of the request body (#48753) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit db5b94014371a52230b20bfa58f79d71a357b6db) --- .../anthropic/transformation.py | 21 +++ .../test_vertex_partner_labels_wire.py | 120 ++++++++++++++++++ ...partner_models_anthropic_transformation.py | 28 ++++ 3 files changed, 169 insertions(+) create mode 100644 tests/integration/providers/test_vertex_partner_labels_wire.py diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 29253359003..363e5596cb9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,5 +1,7 @@ # What is this? ## Handler file for calling claude-3 on vertex ai +import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -19,6 +21,19 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels" +_LABELS: Final = TypeAdapter(Mapping[object, object]) + + +def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None: + if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names): + return None + string_labels: Final = { + key: value + for key, value in _LABELS.validate_python(labels).items() + if isinstance(key, str) and isinstance(value, str) + } + return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None class VertexAIError(Exception): @@ -102,6 +117,12 @@ class VertexAIAnthropicConfig(AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: + params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) + labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers)) + if "labels" in params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + if labels_header is not None: + headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation data: Final = super().transform_request( model=model, messages=messages, diff --git a/tests/integration/providers/test_vertex_partner_labels_wire.py b/tests/integration/providers/test_vertex_partner_labels_wire.py new file mode 100644 index 00000000000..ebdf3f2022e --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_labels_wire.py @@ -0,0 +1,120 @@ +import base64 +import json +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-5-5" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_BACKEND}" +_TARGET: Final = f"{_MODEL_PATH}:rawPredict" +_STREAM_TARGET: Final = f"{_MODEL_PATH}:streamRawPredict?alt=sse" +_LABELS_HEADER: Final = "x-vertex-ai-labels" +_REPLY_TEXT: Final = "labels control" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_LABELS: Final[dict[str, JsonValue]] = {"team": "triage", "env": "prod"} + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_vertex_labels", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_REPLY, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, +) +_SSE: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _EVENTS) + + +def _peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() + ) + if request.target == _STREAM_TARGET: + return Reply(content_type="text/event-stream", chunks=_SSE) + if request.target == _TARGET: + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, api_base), + num_retries=0, + ) + + +def _chat(gateway: Gateway, wire: Wire, scenario: Scenario, stream: bool, **extra: JsonValue) -> Request: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Reply ok"}], + "max_tokens": 16, + "stream": stream, + **extra, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = tuple(request for request in wire.drain() if request.target in (_TARGET, _STREAM_TARGET)) + assert sent.target == (_STREAM_TARGET if stream else _TARGET), sent.target + return sent + + +def _decoded_labels(sent: Request) -> JsonValue: + return json.loads(base64.b64decode(sent.headers[_LABELS_HEADER])) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_regression_vertex_claude_labels_move_from_body_to_header(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, stream, labels=_LABELS) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _decoded_labels(sent) == _LABELS, sent.headers + + +def test_guard_vertex_claude_without_labels_sends_no_labels_header(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, False) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _LABELS_HEADER not in sent.headers, sent.headers + + +def test_regression_vertex_claude_caller_labels_header_wins_over_labels_param(gateway: Gateway) -> None: + own_header: Final = base64.b64encode(json.dumps({"team": "caller"}).encode()).decode() + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat( + gateway, wire, scenario, False, labels=_LABELS, extra_headers={"X-Vertex-AI-Labels": own_header} + ) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert sent.headers[_LABELS_HEADER] == own_header, sent.headers diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index ace24c8a3f4..63a7b208a8e 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -1,4 +1,5 @@ +import base64 import copy import json from collections.abc import Mapping @@ -40,6 +41,33 @@ def test_vertex_ai_anthropic_thinking_param(model, expected_thinking): assert "thinking" not in supported_openai_params +@pytest.mark.parametrize( + ("labels", "headers", "expected_header"), + [ + ({"team": "triage", "env": "prod"}, {}, {"team": "triage", "env": "prod"}), + ({"team": "triage", "cost": 3}, {}, {"team": "triage"}), + ({"cost": 3}, {}, None), + ("team=triage", {}, None), + ({"team": "triage"}, {"x-vertex-ai-labels": "e30="}, {}), + ], + ids=["string_labels", "non_string_values_dropped", "no_string_values", "not_a_mapping", "caller_header_wins"], +) +def test_vertex_anthropic_labels_move_from_body_to_header( + labels: object, headers: dict[str, str], expected_header: dict[str, str] | None +) -> None: + body: Final = VertexAIAnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"max_tokens": 16, "anthropic_version": "vertex-2023-10-16", "labels": labels}, + litellm_params={}, + headers=headers, + ) + sent_labels: Final = next((value for name, value in headers.items() if name.lower() == "x-vertex-ai-labels"), None) + + assert "labels" not in body + assert (None if sent_labels is None else json.loads(base64.b64decode(sent_labels))) == expected_header + + @pytest.mark.parametrize( ("stream", "operation"), [(False, "rawPredict"), (True, "streamRawPredict")], From c84e2e731cfb9cdda5c2c2542598597cbf52e2ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:14:26 -0700 Subject: [PATCH 02/13] feat(vertex_ai): label Vertex requests with the caller's key, team, org, and project Adds the opt-in litellm.vertex_request_label_fields setting, the Vertex twin of bedrock_request_metadata_fields. When set, Gemini, embeddings, Imagen and Gemini image generation, rerank, and Claude on Vertex (chat and /v1/messages) send the listed proxy identity fields and spend_logs_metadata as request labels, identity first, with the user_api_key_ prefix reserved so a caller cannot forge or evict them. Identity text is rewritten to fit Google's label rules. Gemini image generation now forwards labels at all. (cherry picked from commit ae73db38557c271ad39052e434da96dd4eaaab6b) --- litellm/__init__.py | 1 + litellm/llms/vertex_ai/common_utils.py | 19 +- .../vertex_gemini_transformation.py | 15 +- litellm/llms/vertex_ai/request_labels.py | 158 ++++++++++ .../llms/vertex_ai/rerank/transformation.py | 6 +- .../transformation.py | 26 +- .../anthropic/transformation.py | 36 ++- ...test_vertex_request_label_identity_wire.py | 261 ++++++++++++++++ .../llms/vertex_ai/test_request_labels.py | 291 ++++++++++++++++++ 9 files changed, 790 insertions(+), 23 deletions(-) create mode 100644 litellm/llms/vertex_ai/request_labels.py create mode 100644 tests/integration/providers/test_vertex_request_label_identity_wire.py create mode 100644 tests/unit/llms/vertex_ai/test_request_labels.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 8054fad57fe..837ef640103 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -237,6 +237,7 @@ overwrite_user_with_key_hash: bool = ( bedrock_request_metadata_fields: Optional[Sequence[str]] = ( None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata` ) +vertex_request_label_fields: Optional[Sequence[str]] = None store_audit_logs: bool | None = None skip_system_message_in_guardrail: bool = False skip_tool_message_in_guardrail: bool = False diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 3c9a8429882..364dd31605a 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -17,6 +17,11 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_REQUEST_LABELS_MAX, + resolve_vertex_request_labels, + vertex_request_label_fields, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ( VERTEX_AI_PROVIDER_METADATA_FIELDS, @@ -164,6 +169,9 @@ def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> N hidden_params.pop(field, None) +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + def vertex_request_labels_from_litellm_params( litellm_params: dict | None, ) -> dict[str, str] | None: @@ -199,7 +207,8 @@ def pop_vertex_request_labels( """ Resolve labels from optional ``labels`` (Gemini-style) and/or ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]`` - (``requester_metadata``). Pops ``labels`` from optional_params when present. + (``requester_metadata``), plus the proxy identity ``litellm.vertex_request_label_fields`` allows. + Pops ``labels`` from optional_params when present. """ labels: dict[str, str] | None = None if optional_params is not None and "labels" in optional_params: @@ -208,7 +217,13 @@ def pop_vertex_request_labels( labels = {k: v for k, v in raw.items() if isinstance(v, str)} if not labels: labels = vertex_request_labels_from_litellm_params(litellm_params) - return labels if labels else None + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return labels if labels else None + resolved: Final = resolve_vertex_request_labels( + allowed_fields, labels, _LITELLM_PARAMS.validate_python(litellm_params or {}), VERTEX_REQUEST_LABELS_MAX + ) + return dict(resolved) if resolved else None class VertexAIModelRoute(str, Enum): diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index fe94e5d17e0..0655de9fc37 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -15,7 +15,7 @@ from litellm.llms.gemini.common_utils import ( map_gemini_image_tools_params, ) from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, pop_vertex_request_labels from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -253,12 +253,13 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): "generationConfig": generation_config, } - if tools := optional_params.get("tools"): - request_body["tools"] = tools - if tool_config := optional_params.get("toolConfig"): - request_body["toolConfig"] = tool_config - - return request_body + params: Final = _JSON_OBJECT.validate_python(optional_params) + optional_fields: Final = { + "tools": params.get("tools"), + "toolConfig": params.get("toolConfig"), + "labels": pop_vertex_request_labels(optional_params, litellm_params), + } + return {**request_body, **{key: value for key, value in optional_fields.items() if value}} def transform_image_generation_response( self, diff --git a/litellm/llms/vertex_ai/request_labels.py b/litellm/llms/vertex_ai/request_labels.py new file mode 100644 index 00000000000..f1e0b3da99b --- /dev/null +++ b/litellm/llms/vertex_ai/request_labels.py @@ -0,0 +1,158 @@ +""" +Vertex AI forwards request ``labels`` to Cloud Billing, so proxy identity becomes labels only when the operator sets +``litellm.vertex_request_label_fields``. Identity resolves first and the ``user_api_key_`` prefix is reserved, so a +caller can neither evict nor forge it. Identity and ``spend_logs_metadata`` text is lowercased, every character Google's +label rules disallow becomes ``_``, and it is cut to 63 characters, so an alias like ``AI Platform`` still bills under a +compliant ``ai_platform`` label. Caller labels pass through verbatim, as they do with the setting off. +""" + +from __future__ import annotations + +import base64 +import json +import unicodedata +from collections.abc import Mapping +from itertools import chain +from typing import Final, TypeAlias + +from pydantic import TypeAdapter + +import litellm + +VERTEX_REQUEST_LABELS_MAX: Final = 64 +VERTEX_PARTNER_REQUEST_LABELS_MAX: Final = 32 +VERTEX_REQUEST_LABEL_MAX_CHARS: Final = 63 +VERTEX_REQUEST_LABEL_IDENTITY_PREFIX: Final = "user_api_key_" +VERTEX_REQUEST_LABEL_CLIENT_FIELD: Final = "spend_logs_metadata" +VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels" + +_METADATA_PARAM_NAMES: Final = ("metadata", "litellm_metadata") +_KEY_START_CATEGORIES: Final = frozenset(("Ll", "Lo")) +_LABEL_CHARACTER_CATEGORIES: Final = frozenset(("Ll", "Lo", "Nd", "Nl", "No")) +_ANY_MAPPING: Final = TypeAdapter(Mapping[object, object]) +_ANY_SEQUENCE: Final = TypeAdapter(tuple[object, ...]) + +Label: TypeAlias = tuple[str, str] + + +def vertex_request_label_fields() -> tuple[str, ...]: + configured: Final[object] = litellm.vertex_request_label_fields + if not isinstance(configured, (list, tuple)): + return () + return tuple(field for field in _ANY_SEQUENCE.validate_python(configured) if isinstance(field, str)) + + +def _is_label_character(character: str) -> bool: + return character in "_-" or unicodedata.category(character) in _LABEL_CHARACTER_CATEGORIES + + +def _normalized(text: str) -> str: + compliant: Final = (character if _is_label_character(character) else "_" for character in text.lower()) + return "".join(compliant)[:VERTEX_REQUEST_LABEL_MAX_CHARS] + + +def _normalized_label(key: str, value: str) -> Label | None: + normalized_key: Final = _normalized(key) + if not normalized_key or unicodedata.category(normalized_key[0]) not in _KEY_START_CATEGORIES: + return None + return normalized_key, _normalized(value) + + +def _is_reserved(key: str) -> bool: + return key.strip().lower().startswith(VERTEX_REQUEST_LABEL_IDENTITY_PREFIX) + + +def _metadata_sources(litellm_params: Mapping[str, object]) -> tuple[Mapping[object, object], ...]: + candidates: Final = (litellm_params.get(name) for name in _METADATA_PARAM_NAMES) + return tuple(_ANY_MAPPING.validate_python(source) for source in candidates if isinstance(source, Mapping)) + + +def _text_labels(source: object) -> tuple[Label, ...]: + if not isinstance(source, Mapping): + return () + pairs: Final = _ANY_MAPPING.validate_python(source).items() + return tuple((key, value) for key, value in pairs if isinstance(key, str) and isinstance(value, str)) + + +def _first_text(sources: tuple[Mapping[object, object], ...], field: str) -> str | None: + return next((value for source in sources if isinstance(value := source.get(field), str)), None) + + +def _identity_label(sources: tuple[Mapping[object, object], ...], field: str) -> Label | None: + value: Final = _first_text(sources, field) + if value is None or not _is_reserved(field): + return None + return _normalized_label(field, value) + + +def _identity_labels( + sources: tuple[Mapping[object, object], ...], allowed_fields: tuple[str, ...] +) -> tuple[Label, ...]: + candidates: Final = (_identity_label(sources, field) for field in allowed_fields) + return tuple(label for label in candidates if label is not None) + + +def _spend_logs_labels(sources: tuple[Mapping[object, object], ...]) -> tuple[Label, ...]: + pairs: Final = chain.from_iterable( + _text_labels(source.get(VERTEX_REQUEST_LABEL_CLIENT_FIELD)) for source in sources + ) + candidates: Final = (_normalized_label(key, value) for key, value in pairs) + return tuple(label for label in candidates if label is not None) + + +def _first_occurrences(labels: tuple[Label, ...], max_labels: int) -> tuple[Label, ...]: + first_value: Final = dict(reversed(labels)) + keys: Final = tuple(dict.fromkeys(key for key, _ in labels)) + return tuple((key, first_value[key]) for key in keys[:max_labels]) + + +def resolve_vertex_request_labels( + allowed_fields: tuple[str, ...], + caller_labels: object, + litellm_params: Mapping[str, object], + max_labels: int, +) -> tuple[Label, ...]: + """ + The labels to send for a Vertex request once ``allowed_fields`` (``vertex_request_label_fields()``) is non-empty. + Identity comes first, then the caller's string labels minus the reserved prefix, then ``spend_logs_metadata`` + pairs when allowed, each key keeping its first value, up to ``max_labels``. + """ + sources: Final = _metadata_sources(litellm_params) + spend_logs: Final = _spend_logs_labels(sources) if VERTEX_REQUEST_LABEL_CLIENT_FIELD in allowed_fields else () + client: Final = tuple(label for label in (*_text_labels(caller_labels), *spend_logs) if not _is_reserved(label[0])) + return _first_occurrences((*_identity_labels(sources, allowed_fields), *client), max_labels) + + +def is_vertex_labels_header(name: str) -> bool: + return name.lower() == VERTEX_LABELS_HEADER.lower() + + +def _decoded_labels_header(encoded: object) -> object: + if not isinstance(encoded, str): + return None + try: + return _ANY_MAPPING.validate_json(base64.b64decode(encoded, validate=True)) + except ValueError: + return None + + +def vertex_partner_labels_header( + allowed_fields: tuple[str, ...], + caller_labels: object, + outbound_headers: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> str | None: + """ + The ``X-Vertex-AI-Labels`` value for a partner model call once ``allowed_fields`` is non-empty. A labels header the + caller already set stands in for ``caller_labels``, as it does with the setting off, and the result replaces it. + """ + caller_header: Final = next( + (value for name, value in outbound_headers.items() if is_vertex_labels_header(name)), None + ) + resolved: Final = resolve_vertex_request_labels( + allowed_fields, + caller_labels if caller_header is None else _decoded_labels_header(caller_header), + litellm_params, + VERTEX_PARTNER_REQUEST_LABELS_MAX, + ) + return base64.b64encode(json.dumps(dict(resolved)).encode()).decode() if resolved else None diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index dce4d2f2a87..ba0ed47bd5a 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -14,9 +14,7 @@ import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig -from litellm.llms.vertex_ai.common_utils import ( - vertex_request_labels_from_litellm_params, -) +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( @@ -151,7 +149,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # When return_documents is False, we want to ignore record details (return only IDs) request_data["ignoreRecordDetailsInResponse"] = not return_documents - user_labels: Final = vertex_request_labels_from_litellm_params(litellm_params) + user_labels: Final = pop_vertex_request_labels(None, litellm_params) if user_labels: request_data["userLabels"] = user_labels diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 5d748ad6e74..7b954594884 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -5,6 +6,12 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, messages_carry_output_config, ) +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + is_vertex_labels_header, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS, @@ -17,6 +24,18 @@ from ....vertex_llm_base import VertexBase from ..output_params_utils import sanitize_vertex_anthropic_output_params +def _labels_header_replacement( + request_headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> tuple[tuple[str, ...], str | None]: + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return (), None + return ( + tuple(filter(is_vertex_labels_header, request_headers)), + vertex_partner_labels_header(allowed_fields, None, request_headers, litellm_params), + ) + + class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase): @property def custom_llm_provider(self) -> str | None: @@ -118,7 +137,12 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if beta_values: headers["anthropic-beta"] = ",".join(beta_values) - return headers, api_base + replaced_headers, labels_header = _labels_header_replacement(headers, litellm_params) + kept_headers: Final = {name: value for name, value in headers.items() if name not in replaced_headers} + labelled_headers: Final = ( + kept_headers if labels_header is None else {**kept_headers, VERTEX_LABELS_HEADER: labels_header} + ) + return labelled_headers, api_base def get_complete_url( self, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 363e5596cb9..99488ae02ea 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -11,6 +11,12 @@ from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + is_vertex_labels_header, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -21,12 +27,11 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) -_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels" _LABELS: Final = TypeAdapter(Mapping[object, object]) def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None: - if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names): + if not isinstance(labels, dict) or any(map(is_vertex_labels_header, header_names)): return None string_labels: Final = { key: value @@ -118,11 +123,27 @@ class VertexAIAnthropicConfig(AnthropicConfig): headers: dict, ) -> dict: params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) - labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers)) + outbound_headers: Final = _NON_DEFAULT_PARAMS.validate_python(headers) + allowed_fields: Final = vertex_request_label_fields() + labels_header, replaced_headers = ( + ( + vertex_partner_labels_header( + allowed_fields, + params.get("labels"), + outbound_headers, + _NON_DEFAULT_PARAMS.validate_python(litellm_params), + ), + tuple(filter(is_vertex_labels_header, outbound_headers)), + ) + if allowed_fields + else (_vertex_labels_header(params.get("labels"), outbound_headers), ()) + ) if "labels" in params: del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + for name in replaced_headers: + del headers[name] # rebind-ok: the resolved identity labels replace a caller-set labels header if labels_header is not None: - headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation + headers[VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation data: Final = super().transform_request( model=model, messages=messages, @@ -158,11 +179,8 @@ class VertexAIAnthropicConfig(AnthropicConfig): extra_headers: Final = optional_params.get("extra_headers") or {} anthropic_beta_value: Final = extra_headers.get("anthropic-beta", "") - if isinstance(anthropic_beta_value, str) and anthropic_beta_value: - for beta in anthropic_beta_value.split(","): - beta = beta.strip() - if beta: - beta_set.add(beta) + if isinstance(anthropic_beta_value, str): + beta_set.update(filter(None, (beta.strip() for beta in anthropic_beta_value.split(",")))) elif isinstance(anthropic_beta_value, list): beta_set.update(anthropic_beta_value) diff --git a/tests/integration/providers/test_vertex_request_label_identity_wire.py b/tests/integration/providers/test_vertex_request_label_identity_wire.py new file mode 100644 index 00000000000..2d4584cdb0d --- /dev/null +++ b/tests/integration/providers/test_vertex_request_label_identity_wire.py @@ -0,0 +1,261 @@ +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, Scenario +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_CHAT_BACKEND: Final = "gemini-3.7-flash" +_EMBED_BACKEND: Final = "text-embedding-005" +_CLAUDE_BACKEND: Final = "claude-haiku-5-5" +_MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models" +_CHAT_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:generateContent" +_STREAM_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:streamGenerateContent?alt=sse" +_EMBED_TARGET: Final = f"{_MODELS_PATH}/{_EMBED_BACKEND}:predict" +_CLAUDE_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_CLAUDE_BACKEND}:rawPredict" +) +_LABELS_HEADER: Final = "x-vertex-ai-labels" +_REPLY_TEXT: Final = "labels ok" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GEMINI_REPLY: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"role": "model", "parts": [{"text": _REPLY_TEXT}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + "modelVersion": _CHAT_BACKEND, +} +_EMBED_REPLY: Final[dict[str, JsonValue]] = { + "predictions": [{"embeddings": {"values": [0.5, 0.25], "statistics": {"token_count": 2}}}] +} +_CLAUDE_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_labels", + "type": "message", + "role": "assistant", + "model": _CLAUDE_BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, +} + + +def _vertex_peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() + ) + if request.target == _STREAM_TARGET: + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_GEMINI_REPLY)}\n\n".encode(),)) + if request.target == _CHAT_TARGET: + return Reply(body=json.dumps(_GEMINI_REPLY).encode()) + if request.target == _EMBED_TARGET: + return Reply(body=json.dumps(_EMBED_REPLY).encode()) + if request.target == _CLAUDE_TARGET: + return Reply(body=json.dumps(_CLAUDE_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +def _deployment(scenario: Scenario, wire: Wire, backend: str) -> str: + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, wire.url), + num_retries=0, + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + chat_model: str + embed_model: str + claude_model: str + key: str + identity: Mapping[str, str] + uppercase_team_key: str + uppercase_team_identity: Mapping[str, str] + + def request_sent(self, marker: str) -> Request: + (request,) = tuple(request for request in self.wire.drain() if marker.encode() in request.body) + return request + + def labels_sent(self, marker: str) -> JsonValue: + return _JSON_OBJECT.validate_json(self.request_sent(marker).body).get("labels") + + def header_labels_sent(self, marker: str) -> JsonValue: + headers: Final = self.request_sent(marker).headers + (encoded,) = tuple(value for name, value in headers.items() if name.lower() == _LABELS_HEADER) + return _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + directory: Final = tmp_path_factory.mktemp("vertex-request-labels") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **base_config, + "litellm_settings": { + **base_config.get("litellm_settings", {}), + "vertex_request_label_fields": [ + "user_api_key_alias", + "user_api_key_team_alias", + "user_api_key_org_alias", + "user_api_key_project_alias", + "spend_logs_metadata", + ], + }, + } + config_path: Final = directory / "vertex-request-labels.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + wire_server(_vertex_peer) as wire, + owned_proxy(session_gateway, directory, {}, config=config_path) as proxy, + proxy.scenario() as scenario, + ): + chat_model: Final = _deployment(scenario, wire, _CHAT_BACKEND) + embed_model: Final = _deployment(scenario, wire, _EMBED_BACKEND) + claude_model: Final = _deployment(scenario, wire, _CLAUDE_BACKEND) + org_alias: Final = f"ml-org-{suffix}" + team_alias: Final = f"ai-platform-{suffix}" + project_alias: Final = f"support-bot-{suffix}" + key_alias: Final = f"chargeback-{suffix}" + uppercase_key_alias: Final = f"support-{suffix}" + team: Final = scenario.team( + team_alias=team_alias, + organization_id=scenario.organization( + organization_alias=org_alias, models=[chat_model, embed_model, claude_model] + ), + ) + yield Rig( + proxy=proxy, + wire=wire, + chat_model=chat_model, + embed_model=embed_model, + claude_model=claude_model, + key=scenario.key( + key_alias=key_alias, + team_id=team, + project_id=scenario.project(team, project_alias=project_alias), + metadata={"spend_logs_metadata": {"cost-center": "cc-1234", "Owner": "Platform Team"}}, + ), + identity={ + "user_api_key_alias": key_alias, + "user_api_key_team_alias": team_alias, + "user_api_key_org_alias": org_alias, + "user_api_key_project_alias": project_alias, + }, + uppercase_team_key=scenario.key( + key_alias=uppercase_key_alias, + team_id=scenario.team(team_alias=f"AI Platform {suffix}"), + ), + uppercase_team_identity={ + "user_api_key_alias": uppercase_key_alias, + "user_api_key_team_alias": f"ai_platform_{suffix}", + }, + ) + + +def _chat(rig: Rig, key: str, marker: str, stream: bool = False, **extra: JsonValue) -> JsonValue: + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": rig.chat_model, "messages": [{"role": "user", "content": marker}], "stream": stream, **extra}, + key=key, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + return rig.labels_sent(marker) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_chat_labels_carry_the_key_team_org_and_project_identity_with_no_client_labels(rig: Rig, stream: bool) -> None: + labels: Final = _chat(rig, rig.key, f"identity-{uuid.uuid4().hex}", stream) + + assert labels == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"}, labels + + +def test_caller_cannot_forge_identity_and_its_labels_rank_above_spend_logs_metadata(rig: Rig) -> None: + labels: Final = _chat( + rig, + rig.key, + f"forged-{uuid.uuid4().hex}", + labels={"user_api_key_team_alias": "forged", "cost-center": "caller", "ticket": "t-42"}, + ) + + assert labels == {**rig.identity, "cost-center": "caller", "ticket": "t-42", "owner": "platform_team"}, labels + + +def test_request_metadata_labels_still_reach_vertex_under_the_identity(rig: Rig) -> None: + labels: Final = _chat(rig, rig.key, f"requester-{uuid.uuid4().hex}", metadata={"env": "prod"}) + + assert labels == {**rig.identity, "env": "prod", "cost-center": "cc-1234", "owner": "platform_team"}, labels + + +def test_team_alias_outside_googles_label_rules_is_rewritten_to_a_compliant_label(rig: Rig) -> None: + labels: Final = _chat(rig, rig.uppercase_team_key, f"uppercase-{uuid.uuid4().hex}") + + assert labels == rig.uppercase_team_identity, labels + + +def test_embedding_labels_carry_the_caller_identity(rig: Rig) -> None: + marker: Final = f"embed-{uuid.uuid4().hex}" + response: Final = rig.proxy.request( + "POST", "/v1/embeddings", {"model": rig.embed_model, "input": [marker]}, key=rig.key + ) + + assert response.status_code == 200, response.text + assert rig.labels_sent(marker) == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"} + + +def test_claude_chat_labels_header_carries_the_identity_and_drops_a_forged_label(rig: Rig) -> None: + marker: Final = f"claude-chat-{uuid.uuid4().hex}" + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": rig.claude_model, + "messages": [{"role": "user", "content": marker}], + "labels": {"user_api_key_team_alias": "forged", "ticket": "t-42"}, + }, + key=rig.key, + ) + + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert rig.header_labels_sent(marker) == { + **rig.identity, + "ticket": "t-42", + "cost-center": "cc-1234", + "owner": "platform_team", + } + + +def test_claude_messages_labels_header_carries_the_identity(rig: Rig) -> None: + marker: Final = f"claude-messages-{uuid.uuid4().hex}" + response: Final = rig.proxy.request( + "POST", + "/v1/messages", + {"model": rig.claude_model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=rig.key, + ) + + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert rig.header_labels_sent(marker) == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"} diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py new file mode 100644 index 00000000000..6721da0003b --- /dev/null +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -0,0 +1,291 @@ +import base64 +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from pydantic import JsonValue + +import litellm +from litellm.llms.vertex_ai.gemini.transformation import transform_request_body +from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig +from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import VertexAIImagenImageGenerationConfig +from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig +from litellm.llms.vertex_ai.vertex_embeddings.transformation import VertexAITextEmbeddingConfig + +FIELDS: Final = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] +IDENTITY: Final = {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} +LABELS_HEADER: Final = "x-vertex-ai-labels" + +LabelSource = Callable[[dict[str, JsonValue], dict[str, JsonValue]], object] + + +def gemini_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object: + body: Final = transform_request_body( + messages=[{"role": "user", "content": "hi"}], + model="gemini-3.8-flash", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params=dict(litellm_params), + cached_content=None, + ) + return body.get("labels") + + +def embedding_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object: + request: Final = VertexAITextEmbeddingConfig().transform_openai_request_to_vertex_embedding_request( + input=["hi"], optional_params=dict(optional_params), model="text-embedding-005", litellm_params=litellm_params + ) + return request.get("labels") + + +def imagen_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object: + request: Final = VertexAIImagenImageGenerationConfig().transform_image_generation_request( + model="imagen-4.0-generate-001", + prompt="a cat", + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + headers={}, + ) + return request.get("labels") + + +def gemini_image_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object: + request: Final = VertexAIGeminiImageGenerationConfig().transform_image_generation_request( + model="gemini-3.1-flash-image", + prompt="a cat", + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + headers={}, + ) + return request.get("labels") + + +def rerank_labels(litellm_params: dict[str, JsonValue], _: dict[str, JsonValue]) -> object: + request: Final = VertexAIRerankConfig().transform_rerank_request( + model="semantic-ranker-default@latest", + optional_rerank_params={"query": "q", "documents": ["a", "b"]}, + headers={}, + litellm_params=dict(litellm_params), + ) + return request.get("userLabels") + + +def encoded(labels: Mapping[str, str]) -> str: + return base64.b64encode(json.dumps(dict(labels)).encode()).decode() + + +def claude_headers( + litellm_params: dict[str, JsonValue], + optional_params: dict[str, JsonValue], + headers: Mapping[str, str] = MappingProxyType({}), +) -> Mapping[str, str]: + outbound: Final = dict(headers) + VertexAIAnthropicConfig().transform_request( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "hi"}], + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + headers=outbound, + ) + return outbound + + +def labels_header_names(headers: Mapping[str, str]) -> list[str]: + return [name for name in headers if name.lower() == LABELS_HEADER] + + +def decoded_labels(headers: Mapping[str, str]) -> object: + names: Final = labels_header_names(headers) + return json.loads(base64.b64decode(headers[names[0]])) if names else None + + +def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object: + return decoded_labels(claude_headers(litellm_params, optional_params)) + + +LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels] +REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels] +ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels] +CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels] + + +def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]: + return {key: dict(fields)} + + +@pytest.mark.parametrize("setting", [None, []]) +@pytest.mark.parametrize("surface", REQUESTER_METADATA_SURFACES) +def test_setting_off_sends_only_the_callers_labels( + setting: object, surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = metadata( + spend_logs_metadata={"cost-center": "cc-1"}, + requester_metadata={"user_api_key_team_alias": "caller-chosen"}, + **IDENTITY, + ) + + assert surface(params, {}) == {"user_api_key_team_alias": "caller-chosen"} + assert surface(metadata(**IDENTITY), {}) is None + + +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = metadata(spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY) + caller_header: Final = {"X-Vertex-AI-Labels": encoded({"user_api_key_team_alias": "Caller Chosen"})} + + assert claude_labels(params, {"labels": {"user_api_key_team_alias": "Caller Chosen"}}) == { + "user_api_key_team_alias": "Caller Chosen" + } + assert claude_labels(params, {}) is None + assert claude_headers(params, {"labels": {"app": "bot"}}, caller_header) == caller_header + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("surface", ALL_SURFACES) +def test_identity_and_spend_logs_metadata_label_every_vertex_surface( + metadata_key: str, surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + params: Final = metadata(metadata_key, spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY) + + assert surface(params, {}) == {**IDENTITY, "cost-center": "cc-1"} + + +@pytest.mark.parametrize("field_order", [FIELDS, list(reversed(FIELDS))]) +@pytest.mark.parametrize( + ("surface", "cap"), + [ + (gemini_labels, 64), + (embedding_labels, 64), + (imagen_labels, 64), + (gemini_image_labels, 64), + (claude_labels, 32), + ], +) +def test_identity_survives_a_caller_filling_every_label_slot( + field_order: list[str], surface: LabelSource, cap: int, monkeypatch: pytest.MonkeyPatch +): + """64 labels per call for Google models, 32 for partner models, + https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls (2026-10-28)""" + monkeypatch.setattr(litellm, "vertex_request_label_fields", field_order) + caller_labels: Final = {f"caller-{index:02d}": "v" for index in range(cap)} + + labels: Final = surface(metadata(**IDENTITY), {"labels": caller_labels}) + + assert isinstance(labels, dict) + assert len(labels) == cap + assert {key: labels[key] for key in IDENTITY} == IDENTITY + assert all(labels[key] == "v" for key in tuple(caller_labels)[: cap - len(IDENTITY)]) + + +@pytest.mark.parametrize("surface", CALLER_LABEL_SURFACES) +def test_caller_cannot_write_or_override_the_identity_prefix_in_any_case( + surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + forged: Final = { + "user_api_key_team_alias": "forged", + "User_API_Key_Team_Alias": "forged", + " user_api_key_user_id": "forged", + "app": "bot", + } + + assert surface(metadata(**IDENTITY), {"labels": forged}) == {**IDENTITY, "app": "bot"} + assert surface(metadata(user_api_key_alias="chargeback-key"), {"labels": forged}) == { + "user_api_key_alias": "chargeback-key", + "app": "bot", + } + assert surface(metadata(spend_logs_metadata={"USER_API_KEY_TEAM_ALIAS": "forged"}), {}) is None + + +def test_a_caller_set_claude_labels_header_cannot_forge_identity_and_is_replaced(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + caller_header: Final = {"x-Vertex-AI-labels": encoded({"user_api_key_team_alias": "forged", "app": "bot"})} + + sent: Final = claude_headers(metadata(**IDENTITY), {"labels": {"ignored": "x"}}, caller_header) + undecodable: Final = claude_headers(metadata(**IDENTITY), {}, {LABELS_HEADER: "%%"}) + + assert len(labels_header_names(sent)) == 1 + assert decoded_labels(sent) == {**IDENTITY, "app": "bot"} + assert len(labels_header_names(undecodable)) == 1 + assert decoded_labels(undecodable) == IDENTITY + + +@pytest.mark.parametrize("surface", ALL_SURFACES) +def test_identity_and_spend_logs_text_is_rewritten_to_labels_google_accepts( + surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + """Key and value rules from https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls (2026-10-28)""" + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + spend_logs: Final = { + "Owner": "Platform Team", + "1st": "x", + "_private": "x", + "": "x", + "dotted": "a.b", + "too-long": "V" * 64, + "K" * 64: "", + "équipe": "Données-2", + "部門": "営業", + "numerals": "ⅻ½", + } + params: Final = metadata( + user_api_key_alias="Chargeback Key", + user_api_key_team_alias="AI Platform " + "x" * 60, + spend_logs_metadata=spend_logs, + ) + + assert surface(params, {}) == { + "user_api_key_alias": "chargeback_key", + "user_api_key_team_alias": ("ai_platform_" + "x" * 60)[:63], + "owner": "platform_team", + "dotted": "a_b", + "too-long": "v" * 63, + "k" * 63: "", + "équipe": "données-2", + "部門": "営業", + "numerals": "ⅻ½", + } + + +@pytest.mark.parametrize("surface", CALLER_LABEL_SURFACES) +def test_caller_labels_keep_todays_pass_through_and_win_over_spend_logs_metadata( + surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + params: Final = metadata(spend_logs_metadata={"cost-center": "cc-admin", "env": "prod"}, **IDENTITY) + + assert surface(params, {"labels": {"cost-center": "cc-caller", "Mixed": "Case", 1: "dropped"}}) == { + **IDENTITY, + "cost-center": "cc-caller", + "Mixed": "Case", + "env": "prod", + } + + +@pytest.mark.parametrize("surface", ALL_SURFACES) +def test_only_allow_listed_identity_fields_and_opted_in_spend_logs_metadata_become_labels( + surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", ["user_api_key_team_alias", "trace_id"]) + params: Final = metadata(trace_id="t-1", spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY) + + assert surface(params, {}) == {"user_api_key_team_alias": "ai-platform"} + + +def test_metadata_wins_over_litellm_metadata_for_the_same_identity_field(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + params: Final[Mapping[str, JsonValue]] = { + "metadata": {"user_api_key_team_alias": "from-metadata"}, + "litellm_metadata": {"user_api_key_team_alias": "from-litellm-metadata", "user_api_key_alias": "key-a"}, + } + + assert gemini_labels(dict(params), {}) == { + "user_api_key_team_alias": "from-metadata", + "user_api_key_alias": "key-a", + } From dd2eb2764a846848fe22df283f4e2168f68b372b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:10:58 -0700 Subject: [PATCH 03/13] feat(vertex_ai): send the opt-in identity labels on multimodal embeddings, the Google GenAI route, and batches With litellm_settings.vertex_request_label_fields on, the multimodal embedding predict body, the native generateContent and streamGenerateContent body, and the BatchPredictionJob now carry the key's identity labels the same way chat does Caller labels on GenAI are resolved after the extra_body merge, so neither a top-level labels field nor extra_body.labels can write the reserved prefix. Batch labels come only from the proxy's litellm_metadata, never the caller's OpenAI batch metadata. With the setting off every route sends the same body as before (cherry picked from commit a7254fede6b95e00ac84f5ad46e8711595030434) --- litellm/batches/main.py | 1 + .../base_llm/google_genai/transformation.py | 7 + litellm/llms/custom_httpx/llm_http_handler.py | 14 +- litellm/llms/vertex_ai/batches/handler.py | 17 ++- .../llms/vertex_ai/batches/transformation.py | 6 +- .../vertex_ai/google_genai/transformation.py | 15 ++ .../embedding_handler.py | 5 +- .../multimodal_embeddings/transformation.py | 14 +- litellm/types/llms/vertex_ai.py | 5 +- tests/unit/batches/test_main.py | 99 +++++++++++++ .../google_genai/test_google_genai_main.py | 139 ++++++++++++++++++ .../llms/vertex_ai/batches/test_handler.py | 41 ++++++ ..._ai_multimodal_embedding_transformation.py | 76 +++++++++- .../llms/vertex_ai/test_request_labels.py | 64 +++++++- 14 files changed, 489 insertions(+), 14 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9b5630d81ad..6b722db922d 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -359,6 +359,7 @@ def create_batch( max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, custom_endpoint=optional_params.get("custom_endpoint"), + litellm_params=litellm_params, ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 232d9d8eb42..51a6177cfc0 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -75,6 +75,13 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): """ return ("safetySettings", "toolConfig", "cachedContent", "labels", "continuationToken") + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return request_fields + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6737905192e..ce4109e732c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12126,9 +12126,12 @@ class BaseLLMHTTPHandler: generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) + request_fields: Final = generate_content_provider_config.transform_top_level_request_fields( + extra_body, litellm_params + ) - if extra_body: - data.update(extra_body) + if request_fields: + data.update(request_fields) ## LOGGING logging_obj.pre_call( @@ -12241,9 +12244,12 @@ class BaseLLMHTTPHandler: generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) + request_fields: Final = generate_content_provider_config.transform_top_level_request_fields( + extra_body, litellm_params + ) - if extra_body: - data.update(extra_body) + if request_fields: + data.update(request_fields) ## LOGGING logging_obj.pre_call( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index a039347e824..3689b51d9dc 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Coroutine, Sequence +from collections.abc import Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import urlparse @@ -21,6 +21,11 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_REQUEST_LABELS_MAX, + resolve_vertex_request_labels, + vertex_request_label_fields, +) from litellm.llms.vertex_ai.vertex_llm_base import graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -91,6 +96,14 @@ def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatus ) +def _proxy_identity_labels(litellm_params: Mapping[str, object]) -> Mapping[str, str]: + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return {} + proxy_metadata: Final = {"litellm_metadata": litellm_params.get("litellm_metadata")} + return dict(resolve_vertex_request_labels(allowed_fields, None, proxy_metadata, VERTEX_REQUEST_LABELS_MAX)) + + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): super().__init__(*args, **kwargs) @@ -107,6 +120,7 @@ class VertexAIBatchPrediction(VertexLLM): timeout: float | httpx.Timeout, max_retries: int | None, custom_endpoint: bool | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: if custom_endpoint: raise VertexAIError( @@ -135,6 +149,7 @@ class VertexAIBatchPrediction(VertexLLM): request=create_batch_data, vertex_project=vertex_project or project_id, vertex_location=vertex_location or "us-central1", + labels=_proxy_identity_labels(litellm_params or {}), ) ) vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 2460278294c..66f12d5b636 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -205,6 +205,7 @@ class VertexAIBatchTransformation: request: CreateBatchRequest, vertex_project: str | None = None, vertex_location: str | None = None, + labels: Mapping[str, str] | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -223,12 +224,15 @@ class VertexAIBatchTransformation: predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), ) - return VertexAIBatchPredictionJob( + job: Final = VertexAIBatchPredictionJob( inputConfig=input_config, outputConfig=output_config, model=model, displayName=request_display_name, ) + if not labels: + return job + return {**job, "labels": dict(labels)} @classmethod def transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 563a87b862d..78fe644a872 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -2,9 +2,12 @@ Transformation for Calling Google models in their native format. """ +from collections.abc import Mapping from typing import Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.types.router import GenericLiteLLMParams @@ -38,6 +41,18 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): return default_headers + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + if not vertex_request_label_fields(): + return request_fields + fields: Final = request_fields or {} + unlabelled: Final = {name: value for name, value in fields.items() if name != "labels"} + labels: Final = pop_vertex_request_labels({"labels": fields.get("labels")}, dict(litellm_params)) + return {**unlabelled, "labels": labels} if labels else unlabelled + def _camel_to_snake(self, camel_str: str) -> str: """Convert camelCase to snake_case""" import re diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py index 8af05b3ef32..18dd942a4b3 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py @@ -1,7 +1,9 @@ import json +from collections.abc import Mapping from typing import Final, Literal import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -19,6 +21,7 @@ from litellm.types.utils import EmbeddingResponse from .transformation import VertexAIMultimodalEmbeddingConfig vertex_multimodal_embedding_handler: Final = VertexAIMultimodalEmbeddingConfig() +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) class VertexMultimodalEmbedding(VertexLLM): @@ -84,7 +87,7 @@ class VertexMultimodalEmbedding(VertexLLM): sync_handler = client request_data: Final = vertex_multimodal_embedding_handler.transform_embedding_request( - model, input, optional_params, headers + model, input, optional_params, headers, _LITELLM_PARAMS.validate_python(litellm_params) ) headers = vertex_multimodal_embedding_handler.validate_environment( diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index da77fa338db..61028ee982d 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final, cast from httpx import Headers, Response @@ -22,7 +23,8 @@ from litellm.types.utils import ( from litellm.utils import count_characters, is_base64_encoded from ...base_llm.embedding.transformation import BaseEmbeddingConfig -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, pop_vertex_request_labels +from ..request_labels import vertex_request_label_fields class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): @@ -169,8 +171,14 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): input: AllEmbeddingInputValues, optional_params: dict, headers: dict, - ) -> dict: + litellm_params: Mapping[str, object] | None = None, + ) -> dict[str, object]: optional_params = optional_params or {} + labels: Final = ( + pop_vertex_request_labels(optional_params, dict(litellm_params or {})) + if vertex_request_label_fields() + else None + ) request_data: Final = VertexMultimodalEmbeddingRequest(instances=[]) @@ -192,7 +200,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if "outputDimensionality" in optional_params: request_data["parameters"] = {"dimension": optional_params["outputDimensionality"]} - return cast(dict, request_data) + return {**request_data, "labels": labels} if labels else {**request_data} def transform_embedding_response( self, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 0480d6b1f3e..691c2c37e0e 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + NotRequired, ReadOnly, Required, TypedDict, @@ -581,6 +582,7 @@ class Instance(TypedDict, total=False): class VertexMultimodalEmbeddingRequest(TypedDict, total=False): instances: Required[list[Instance]] parameters: dict + labels: ReadOnly[Mapping[str, str]] class VideoEmbedding(TypedDict): @@ -728,6 +730,7 @@ class VertexAIBatchPredictionJob(TypedDict): model: str inputConfig: InputConfig outputConfig: OutputConfig + labels: NotRequired[ReadOnly[Mapping[str, str]]] class VertexBatchPredictionResponse(TypedDict, total=False): diff --git a/tests/unit/batches/test_main.py b/tests/unit/batches/test_main.py index 7106f8c7bc2..8bb434e1442 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -1229,3 +1229,102 @@ def test_retrieve__gemini_routes_to_base_http_handler_with_gemini_config(seams): forwarded = seams.base_http.retrieve_batch.call_args.kwargs assert type(forwarded["provider_config"]).__name__ == "GeminiBatchesConfig" assert forwarded["batch_id"] == "la6" + + +VERTEX_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +VERTEX_IDENTITY: Final = MappingProxyType( + {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} +) +VERTEX_PROXY_METADATA: Final = MappingProxyType({**VERTEX_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +VERTEX_BATCH_FILE: Final = "gs://bucket/publishers/google/models/gemini-3.8-flash/file-uuid" +VERTEX_SENT_JOB: Final = TypeAdapter(Mapping[str, object]) +VERTEX_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) +VERTEX_BATCH_JOBS_URL: Final = ( + "https://us-central1-aiplatform.googleapis.com/v1/projects/proj/locations/us-central1/batchPredictionJobs" +) + + +async def _vertex_batch_job_sent( + monkeypatch: pytest.MonkeyPatch, is_async: bool, batch_metadata: Mapping[str, object] +) -> Mapping[str, object]: + import litellm.llms.vertex_ai.batches.handler as vertex_batches + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + reply: Final = MagicMock( + side_effect=lambda _request: httpx.Response( + 200, + json={ + "name": "projects/proj/locations/us-central1/batchPredictionJobs/123", + "state": "JOB_STATE_PENDING", + "createTime": "2026-10-29T00:00:00Z", + }, + ) + ) + transport: Final = httpx.MockTransport(reply) + handler: Final = vertex_batches.VertexAIBatchPrediction(gcs_bucket_name="bucket") + monkeypatch.setattr(handler, "_ensure_access_token", lambda **_: ("token", "proj")) + monkeypatch.setattr(bm, "vertex_ai_batches_instance", handler) + monkeypatch.setattr( + vertex_batches, "get_httpx_client", lambda **_: HTTPHandler(client=httpx.Client(transport=transport)) + ) + monkeypatch.setattr(vertex_batches, "get_async_httpx_client", lambda **_: AsyncHTTPHandler(transport=transport)) + batch: Final = { + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + "input_file_id": VERTEX_BATCH_FILE, + "custom_llm_provider": "vertex_ai", + "vertex_location": "us-central1", + **batch_metadata, + } + if is_async: + await bm.acreate_batch(**batch) + else: + bm.create_batch(**batch) + + sent: Final[httpx.Request] = reply.call_args.args[0] + assert str(sent.url) == VERTEX_BATCH_JOBS_URL + return VERTEX_SENT_JOB.validate_json(sent.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_job_carries_proxy_identity_never_the_callers_batch_metadata( + monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", VERTEX_LABEL_FIELDS) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, + is_async, + { + "metadata": {"user_api_key_alias": "forged", "user_api_key_team_alias": "forged", "app": "bot"}, + "litellm_metadata": dict(VERTEX_PROXY_METADATA), + }, + ) + + labels: Final = VERTEX_SENT_LABELS.validate_python(job["labels"]) + assert list(labels.items()) == [*VERTEX_IDENTITY.items(), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_metadata_alone_never_labels_the_job(monkeypatch: pytest.MonkeyPatch, is_async: bool): + monkeypatch.setattr(litellm, "vertex_request_label_fields", VERTEX_LABEL_FIELDS) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, is_async, {"metadata": {**VERTEX_IDENTITY, "spend_logs_metadata": "cc-1"}} + ) + + assert "labels" not in job + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_job_with_the_setting_off_is_unlabelled(monkeypatch: pytest.MonkeyPatch, is_async: bool): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, is_async, {"metadata": {"app": "bot"}, "litellm_metadata": dict(VERTEX_PROXY_METADATA)} + ) + + assert "labels" not in job diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 21f42ca2dff..8a03c76f408 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -4,8 +4,12 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final import pytest +from pydantic import TypeAdapter @@ -286,3 +290,138 @@ def test_vertex_generate_content_routes_by_deployment_base_model(base_model, use ) assert (setup.generate_content_provider_config is not None) is uses_native_api + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType( + {**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}, "requester_metadata": {"team": "search"}} +) +_FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +_GENAI_ENTRYPOINTS: Final = ( + "generate_content", + "agenerate_content", + "generate_content_stream", + "agenerate_content_stream", +) +_GEMINI_REPLY: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, +} +_SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def _generate_content_body( + entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object], model: str +) -> Mapping[str, object]: + from unittest.mock import AsyncMock, MagicMock + + import httpx + + from litellm import google_genai + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig + + reply: Final = MagicMock(side_effect=lambda _request: httpx.Response(200, json=_GEMINI_REPLY)) + vertex_token: Final = ("token", "proj") + monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token", MagicMock(return_value=vertex_token)) + monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token_async", AsyncMock(return_value=vertex_token)) + is_async: Final = entrypoint.startswith("a") + transport: Final = httpx.MockTransport(reply) + client: Final = ( + AsyncHTTPHandler(transport=transport) if is_async else HTTPHandler(client=httpx.Client(transport=transport)) + ) + call: Final = getattr(google_genai, entrypoint)( + model=model, + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + client=client, + metadata=dict(_PROXY_METADATA), + **request_fields, + ) + if is_async: + await call + + sent: Final[httpx.Request] = reply.call_args.args[0] + return _SENT_BODY.validate_json(sent.content) + + +async def _vertex_generate_content_body( + entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object] +) -> Mapping[str, object]: + return await _generate_content_body( + entrypoint, + monkeypatch, + {**request_fields, "vertex_location": "us-central1"}, + "vertex_ai/gemini-3.8-flash", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@pytest.mark.parametrize( + "request_fields", + [{"labels": dict(_FORGED)}, {"extra_body": {"labels": dict(_FORGED)}}], + ids=["top-level-labels", "extra-body-labels"], +) +async def test_vertex_generate_content_sends_identity_labels_a_caller_cannot_forge( + entrypoint: str, request_fields: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + + labels: Final = _SENT_LABELS.validate_python(body["labels"]) + assert list(labels.items()) == [*_IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +async def test_vertex_generate_content_without_caller_labels_still_carries_identity( + entrypoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, {}) + + labels: Final = _SENT_LABELS.validate_python(body["labels"]) + assert list(labels.items()) == [*_IDENTITY.items(), ("team", "search"), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@pytest.mark.parametrize( + "request_fields, expected_labels", + [ + ({}, "absent"), + ({"labels": dict(_FORGED)}, _FORGED), + ({"extra_body": {"labels": dict(_FORGED)}}, _FORGED), + ({"labels": "not-a-mapping"}, "not-a-mapping"), + ], + ids=["no-labels", "top-level-labels", "extra-body-labels", "non-mapping-labels"], +) +async def test_vertex_generate_content_with_the_setting_off_sends_only_the_callers_labels( + entrypoint: str, + request_fields: Mapping[str, object], + expected_labels: object, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + + assert body.get("labels", "absent") == expected_labels + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +async def test_ai_studio_generate_content_keeps_the_callers_labels_with_the_setting_on( + entrypoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _generate_content_body( + entrypoint, monkeypatch, {"labels": dict(_FORGED), "api_key": "test-key"}, "gemini/gemini-3.8-flash" + ) + + assert body["labels"] == _FORGED diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 2c246e065a8..97d4a88f16d 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio import json +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -262,6 +263,46 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +@pytest.mark.parametrize( + "create_data, model", + [ + (CREATE_DATA, "publishers/google/models/gemini-1.5-flash-001"), + (ENDPOINT_CREATE_DATA, TUNED_MODEL_RESOURCE), + ], + ids=["publisher-model", "fine-tuned-endpoint"], +) +def test_create_batch_labels_the_job_from_litellm_metadata_never_metadata( + create_data: Mapping[str, str], model: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", ["user_api_key_alias", "user_api_key_team_alias"]) + h: Final = _make_handler() + client: Final = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}.get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()), + ): + h.create_batch( + _is_async=False, + create_batch_data=create_data, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + litellm_params={ + "metadata": {"user_api_key_alias": "forged"}, + "litellm_metadata": {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}, + }, + ) + + sent: Final = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == model + assert sent["labels"] == {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} + + @pytest.mark.parametrize( "api_base, expected", [ diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index edb6e889814..c790de4bc47 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -1,10 +1,16 @@ import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import TypeAdapter - +import litellm +import litellm.main +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.multimodal_embeddings.transformation import ( VertexAIMultimodalEmbeddingConfig, ) @@ -174,3 +180,71 @@ class TestVertexMultimodalEmbedding: headers={}, ) assert "parameters" not in request + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +TODAYS_INSTANCES: Final = [{"text": "a red apple"}] +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def multimodal_embedding_body( + is_async: bool, monkeypatch: pytest.MonkeyPatch, request_labels: Mapping[str, object] +) -> Mapping[str, object]: + reply: Final = MagicMock( + side_effect=lambda _request: httpx.Response(200, json={"predictions": [{"textEmbedding": [0.1, 0.2]}]}) + ) + monkeypatch.setattr(litellm.main.vertex_multimodal_embedding, "_ensure_access_token", lambda **_: ("token", "proj")) + transport: Final = httpx.MockTransport(reply) + embed: Final = { + "model": "vertex_ai/multimodalembedding@001", + "input": "a red apple", + "vertex_location": "us-central1", + "metadata": dict(PROXY_METADATA), + **request_labels, + } + if is_async: + await litellm.aembedding(**embed, client=AsyncHTTPHandler(transport=transport)) + else: + litellm.embedding(**embed, client=HTTPHandler(client=httpx.Client(transport=transport))) + + sent: Final[httpx.Request] = reply.call_args.args[0] + return SENT_BODY.validate_json(sent.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("request_labels", [{}, {"labels": dict(FORGED)}], ids=["no-caller-labels", "caller-labels"]) +async def test_setting_off_sends_todays_multimodal_embedding_body( + is_async: bool, request_labels: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + assert await multimodal_embedding_body(is_async, monkeypatch, request_labels) == {"instances": TODAYS_INSTANCES} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "request_labels, expected_labels", + [ + ({}, [*IDENTITY.items(), ("cost-center", "cc-1")]), + ({"labels": dict(FORGED)}, [*IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")]), + ], + ids=["no-caller-labels", "forged-caller-labels"], +) +async def test_multimodal_embedding_sends_identity_labels_a_caller_cannot_forge( + is_async: bool, + request_labels: Mapping[str, object], + expected_labels: Sequence[tuple[str, str]], + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + + body: Final = await multimodal_embedding_body(is_async, monkeypatch, request_labels) + + assert body["instances"] == TODAYS_INSTANCES + assert list(SENT_LABELS.validate_python(body["labels"]).items()) == list(expected_labels) diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index 6721da0003b..c91694a3c5a 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -9,11 +9,14 @@ from pydantic import JsonValue import litellm from litellm.llms.vertex_ai.gemini.transformation import transform_request_body +from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import VertexAIImagenImageGenerationConfig +from litellm.llms.vertex_ai.multimodal_embeddings.transformation import VertexAIMultimodalEmbeddingConfig from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.llms.vertex_ai.vertex_embeddings.transformation import VertexAITextEmbeddingConfig +from litellm.types.router import GenericLiteLLMParams FIELDS: Final = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] IDENTITY: Final = {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} @@ -63,6 +66,28 @@ def gemini_image_labels(litellm_params: dict[str, JsonValue], optional_params: d return request.get("labels") +def multimodal_embedding_labels( + litellm_params: Mapping[str, JsonValue], optional_params: Mapping[str, JsonValue] +) -> object: + request: Final = VertexAIMultimodalEmbeddingConfig().transform_embedding_request( + model="multimodalembedding@001", + input="a red apple", + optional_params=dict(optional_params), + headers={}, + litellm_params=litellm_params, + ) + return request.get("labels") + + +def google_genai_labels( + litellm_params: Mapping[str, JsonValue], optional_params: Mapping[str, JsonValue] +) -> object: + fields: Final = VertexAIGoogleGenAIConfig().transform_top_level_request_fields( + optional_params or None, GenericLiteLLMParams(**litellm_params) + ) + return (fields or {}).get("labels") + + def rerank_labels(litellm_params: dict[str, JsonValue], _: dict[str, JsonValue]) -> object: request: Final = VertexAIRerankConfig().transform_rerank_request( model="semantic-ranker-default@latest", @@ -108,8 +133,9 @@ def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[st LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels] REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels] -ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels] -CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels] +SETTING_GATED_SURFACES: Final = [multimodal_embedding_labels, google_genai_labels] +ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels, *SETTING_GATED_SURFACES] +CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels, *SETTING_GATED_SURFACES] def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]: @@ -145,6 +171,38 @@ def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: objec assert claude_headers(params, {"labels": {"app": "bot"}}, caller_header) == caller_header +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_sends_no_multimodal_embedding_labels(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = metadata( + spend_logs_metadata={"cost-center": "cc-1"}, requester_metadata={"team": "search"}, **IDENTITY + ) + + assert multimodal_embedding_labels(params, {"labels": {"app": "bot"}}) is None + + +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_forwards_google_genai_fields_untouched(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = GenericLiteLLMParams(**metadata(spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY)) + fields: Final = {"labels": {"user_api_key_team_alias": "Caller Chosen", "app": "bot"}, "cachedContent": "c-1"} + + assert VertexAIGoogleGenAIConfig().transform_top_level_request_fields(fields, params) is fields + assert VertexAIGoogleGenAIConfig().transform_top_level_request_fields(None, params) is None + + +def test_google_genai_keeps_the_other_top_level_fields_and_drops_an_empty_labels_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + config: Final = VertexAIGoogleGenAIConfig() + fields: Final = {"labels": {"user_api_key_alias": "forged"}, "cachedContent": "c-1"} + + assert config.transform_top_level_request_fields(fields, GenericLiteLLMParams()) == {"cachedContent": "c-1"} + assert config.transform_top_level_request_fields(fields, GenericLiteLLMParams(**metadata(**IDENTITY))) == { + "labels": IDENTITY, + "cachedContent": "c-1", + } + + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("surface", ALL_SURFACES) def test_identity_and_spend_logs_metadata_label_every_vertex_surface( @@ -165,6 +223,8 @@ def test_identity_and_spend_logs_metadata_label_every_vertex_surface( (imagen_labels, 64), (gemini_image_labels, 64), (claude_labels, 32), + (multimodal_embedding_labels, 64), + (google_genai_labels, 64), ], ) def test_identity_survives_a_caller_filling_every_label_slot( From 9fe1fe0b6a8ecd7aaf0c7644a9b7564d79b0eac3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 14:32:37 -0800 Subject: [PATCH 04/13] fix(vertex_ai): keep identity labels when a Gemini caller sends extra_body labels (cherry picked from commit f139f7e33f611826719ccf52f0989a0c7d024ab0) --- litellm/llms/vertex_ai/gemini/transformation.py | 5 +++++ .../unit/llms/vertex_ai/test_request_labels.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 895d59965b7..702187e1c12 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.types.files import ( get_file_mime_type_for_file_type, get_file_type_from_extension, @@ -1240,6 +1241,10 @@ def transform_request_body( if labels and custom_llm_provider != LlmProviders.GEMINI: data["labels"] = labels _pop_and_merge_extra_body(data, optional_params) + if custom_llm_provider != LlmProviders.GEMINI and vertex_request_label_fields(): + merged_labels: Final = pop_vertex_request_labels({"labels": data.pop("labels", None)}, litellm_params) + if merged_labels: + data["labels"] = merged_labels _rewrite_google_maps_response_format(data) except Exception as e: raise e diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index c91694a3c5a..cbe990dceb7 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -263,6 +263,23 @@ def test_caller_cannot_write_or_override_the_identity_prefix_in_any_case( assert surface(metadata(spend_logs_metadata={"USER_API_KEY_TEAM_ALIAS": "forged"}), {}) is None +@pytest.mark.parametrize("extra_labels", [{"user_api_key_team_alias": "forged", "app": "bot"}, None]) +def test_gemini_extra_body_labels_cannot_replace_or_drop_identity( + extra_labels: dict[str, JsonValue] | None, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + expected: Final = {**IDENTITY, "app": "bot"} if extra_labels else IDENTITY + + assert gemini_labels(metadata(**IDENTITY), {"extra_body": {"labels": extra_labels}}) == expected + + +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_merges_gemini_extra_body_labels_as_sent(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + optional_params: Final = {"labels": {"env": "prod"}, "extra_body": {"labels": {"user_api_key_team_alias": "x"}}} + + assert gemini_labels(metadata(**IDENTITY), optional_params) == {"env": "prod", "user_api_key_team_alias": "x"} + def test_a_caller_set_claude_labels_header_cannot_forge_identity_and_is_replaced(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) caller_header: Final = {"x-Vertex-AI-labels": encoded({"user_api_key_team_alias": "forged", "app": "bot"})} From 24ba38af06fdbe03bc0491eadae4f1ec39103155 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:17:31 -0800 Subject: [PATCH 05/13] fix(vertex_ai): label Interactions, image edits, Mistral, and the /vertex_ai pass-through, and gate Gemini image labels on the setting (cherry picked from commit 8c2a905e19263152a07a08f52dc022a7a5fd5314) --- litellm/interactions/http_handler.py | 10 +- .../base_llm/interactions/transformation.py | 8 + litellm/llms/vertex_ai/common_utils.py | 54 +++- .../llms/vertex_ai/gemini/transformation.py | 5 +- .../vertex_ai/google_genai/transformation.py | 10 +- .../vertex_gemini_transformation.py | 14 +- .../vertex_gemini_transformation.py | 5 +- .../vertex_ai/interactions/transformation.py | 9 +- .../vertex_ai_partner_models/main.py | 28 +- .../llm_passthrough_endpoints.py | 59 ++++ .../pass_through_endpoints.py | 54 +++- .../google_genai/test_google_genai_main.py | 135 +++++--- ...est_vertex_ai_image_edit_transformation.py | 48 ++- ...t_vertex_ai_interactions_transformation.py | 98 ++++++ .../llms/vertex_ai/test_request_labels.py | 24 +- ...i_partner_models_mistral_transformation.py | 112 +++++++ .../test_llm_pass_through_endpoints.py | 3 + .../test_llm_passthrough_endpoints.py | 306 ++++++++++++++++++ 18 files changed, 879 insertions(+), 103 deletions(-) diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 09aa656e49b..5cb1931259d 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -161,8 +161,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( @@ -256,8 +257,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index b7bbf7bde73..ff69a21d7f0 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -11,6 +11,7 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -118,6 +119,13 @@ class BaseInteractionsAPIConfig(ABC): raise ValueError("api_base is required") return api_base + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return request_fields + @abstractmethod def transform_request( self, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 364dd31605a..93f319a021d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -28,6 +28,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages @@ -173,7 +174,7 @@ _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) def vertex_request_labels_from_litellm_params( - litellm_params: dict | None, + litellm_params: Mapping[str, object] | None, ) -> dict[str, str] | None: """ Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``: @@ -194,38 +195,57 @@ def vertex_request_labels_from_litellm_params( rm = metadata["requester_metadata"] if not isinstance(rm, dict): continue - labels = {k: v for k, v in rm.items() if isinstance(v, str)} + labels = {k: v for k, v in rm.items() if isinstance(k, str) and isinstance(v, str)} if labels: return labels return None +def pop_vertex_caller_labels( + optional_params: dict | None, + litellm_params: Mapping[str, object] | None, +) -> dict[str, str] | None: + """ + Pop the caller's optional ``labels`` (Gemini-style), falling back to ``requester_metadata`` + on ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``. + """ + raw: Final = optional_params.pop("labels", None) if optional_params is not None else None + labels: Final = ( + {k: v for k, v in raw.items() if isinstance(k, str) and isinstance(v, str)} if isinstance(raw, dict) else None + ) + return labels or vertex_request_labels_from_litellm_params(litellm_params) + + def pop_vertex_request_labels( optional_params: dict | None, litellm_params: dict | None, ) -> dict[str, str] | None: """ - Resolve labels from optional ``labels`` (Gemini-style) and/or - ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]`` - (``requester_metadata``), plus the proxy identity ``litellm.vertex_request_label_fields`` allows. - Pops ``labels`` from optional_params when present. + The caller's labels from ``pop_vertex_caller_labels``, plus the proxy identity + ``litellm.vertex_request_label_fields`` allows. """ - labels: dict[str, str] | None = None - if optional_params is not None and "labels" in optional_params: - raw: Final = optional_params.pop("labels") - if isinstance(raw, dict): - labels = {k: v for k, v in raw.items() if isinstance(v, str)} - if not labels: - labels = vertex_request_labels_from_litellm_params(litellm_params) + request_params: Final = _LITELLM_PARAMS.validate_python(litellm_params or {}) + labels: Final = pop_vertex_caller_labels(optional_params, request_params) allowed_fields: Final = vertex_request_label_fields() if not allowed_fields: - return labels if labels else None - resolved: Final = resolve_vertex_request_labels( - allowed_fields, labels, _LITELLM_PARAMS.validate_python(litellm_params or {}), VERTEX_REQUEST_LABELS_MAX - ) + return labels + resolved: Final = resolve_vertex_request_labels(allowed_fields, labels, request_params, VERTEX_REQUEST_LABELS_MAX) return dict(resolved) if resolved else None +def vertex_top_level_request_fields( + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, +) -> Mapping[str, object] | None: + """The caller's top-level body fields with ``labels`` replaced by ``pop_vertex_request_labels`` while the setting is on.""" + if not vertex_request_label_fields(): + return request_fields + fields: Final = request_fields or {} + unlabelled: Final = {name: value for name, value in fields.items() if name != "labels"} + labels: Final = pop_vertex_request_labels({"labels": fields.get("labels")}, dict(litellm_params)) + return {**unlabelled, "labels": labels} if labels else unlabelled + + class VertexAIModelRoute(str, Enum): """Enum for Vertex AI model routing""" diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 702187e1c12..dfde8150b34 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -29,7 +29,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.common_utils import pop_vertex_caller_labels, pop_vertex_request_labels from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.types.files import ( get_file_mime_type_for_file_type, @@ -1186,8 +1186,9 @@ def transform_request_body( optional_params.pop("output_config", None) config_fields: Final = GenerationConfig.__annotations__.keys() + request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) - labels: Final = pop_vertex_request_labels(optional_params, litellm_params) + labels: Final = pop_vertex_caller_labels(optional_params, request_params) filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))} diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 78fe644a872..9982763a846 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -6,8 +6,7 @@ from collections.abc import Mapping from typing import Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig -from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels -from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields +from litellm.llms.vertex_ai.common_utils import vertex_top_level_request_fields from litellm.types.router import GenericLiteLLMParams @@ -46,12 +45,7 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): request_fields: Mapping[str, object] | None, litellm_params: GenericLiteLLMParams, ) -> Mapping[str, object] | None: - if not vertex_request_label_fields(): - return request_fields - fields: Final = request_fields or {} - unlabelled: Final = {name: value for name, value in fields.items() if name != "labels"} - labels: Final = pop_vertex_request_labels({"labels": fields.get("labels")}, dict(litellm_params)) - return {**unlabelled, "labels": labels} if labels else unlabelled + return vertex_top_level_request_fields(request_fields, litellm_params) def _camel_to_snake(self, camel_str: str) -> str: """Convert camelCase to snake_case""" diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 8e12add0316..09c8596043d 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -12,8 +12,9 @@ import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, vertex_top_level_request_fields from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.vertex_ai import ( @@ -190,7 +191,16 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value } - request_body: Final[dict[str, object]] = {"contents": contents, "generationConfig": generation_config} + label_fields: Final = ( + vertex_top_level_request_fields({"labels": litellm_params.get("labels")}, litellm_params) + if vertex_request_label_fields() + else None + ) + request_body: Final[dict[str, object]] = { + "contents": contents, + "generationConfig": generation_config, + **(label_fields or {}), + } payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 0655de9fc37..98056e47aa1 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.gemini.common_utils import ( from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, pop_vertex_request_labels from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( @@ -257,7 +258,9 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): optional_fields: Final = { "tools": params.get("tools"), "toolConfig": params.get("toolConfig"), - "labels": pop_vertex_request_labels(optional_params, litellm_params), + "labels": pop_vertex_request_labels(optional_params, litellm_params) + if vertex_request_label_fields() + else None, } return {**request_body, **{key: value for key, value in optional_fields.items() if value}} diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index 36965fe666f..7b600c4dc57 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -4,7 +4,7 @@ from typing import Final from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig -from litellm.llms.vertex_ai.common_utils import validate_vertex_location +from litellm.llms.vertex_ai.common_utils import validate_vertex_location, vertex_top_level_request_fields from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.router import GenericLiteLLMParams @@ -111,6 +111,13 @@ class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): collection_url: Final = self._target(api_base, params).collection_url return f"{collection_url}?alt=sse" if stream else collection_url + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return vertex_top_level_request_fields(request_fields, litellm_params) + def _interaction_by_id_request( self, interaction_id: str, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index c855a073648..d2459b4afa1 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -1,14 +1,20 @@ # What is this? ## API Handler for calling Vertex AI Partner Models -from collections.abc import Callable +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm import LlmProviders +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.types.llms.vertex_ai import VertexPartnerProvider from litellm.utils import ModelResponse @@ -16,6 +22,7 @@ from ...custom_httpx.llm_http_handler import BaseLLMHTTPHandler from ..vertex_llm_base import VertexBase base_llm_http_handler = BaseLLMHTTPHandler() +_PARAMS: Final = TypeAdapter(Mapping[str, object]) class VertexAIError(Exception): @@ -158,6 +165,21 @@ class VertexAIPartnerModels(VertexBase): model=model, ) + label_fields: Final = vertex_request_label_fields() if partner == VertexPartnerProvider.mistralai else () + labels_header: Final = ( + vertex_partner_labels_header( + label_fields, + _PARAMS.validate_python(optional_params).get("labels"), + _PARAMS.validate_python(headers or {}), + _PARAMS.validate_python(litellm_params), + ) + if label_fields + else None + ) + if label_fields and "labels" in optional_params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + vertex_labels_headers: Final = {} if labels_header is None else {VERTEX_LABELS_HEADER: labels_header} + if "codestral" in model or "mistral" in model: model = model.split("@")[0] @@ -179,6 +201,7 @@ class VertexAIPartnerModels(VertexBase): logger_fn=logger_fn, timeout=timeout, encoding=encoding, + headers=vertex_labels_headers, ) elif "claude" in model: # Build a new dict so we never mutate the shared deployment extra_headers object. @@ -249,6 +272,9 @@ class VertexAIPartnerModels(VertexBase): encoding=encoding, custom_llm_provider="vertex_ai", custom_endpoint=True, + headers={"Content-Type": "application/json", **vertex_labels_headers} + if vertex_labels_headers + else None, ) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 978882ef3f1..22ecbcd0f74 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -68,6 +68,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 * @@ -109,6 +117,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + UpstreamRequest, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -2566,6 +2575,53 @@ def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str] return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) +_VERTEX_BODY_LABELLED_METHODS: Final = frozenset( + {"generateContent", "streamGenerateContent", "predict", "batchPredictionJobs"} +) +_VERTEX_HEADER_LABELLED_METHODS: Final = frozenset({"rawPredict", "streamRawPredict"}) + + +def _vertex_labelled_method(http_method: str, endpoint: str) -> str | None: + if http_method != "POST": + return None + resource: Final = endpoint.rstrip("/").rsplit("/", 1)[-1] + return resource.partition(":")[2] or resource + + +def _with_vertex_body_labels( + allowed_fields: tuple[str, ...], body: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_vertex_request_labels( + allowed_fields, body.get("labels"), litellm_params, VERTEX_REQUEST_LABELS_MAX + ) + unlabelled: Final = {key: value for key, value in body.items() if key != "labels"} + return {**unlabelled, "labels": dict(resolved)} if resolved else unlabelled + + +def _with_vertex_labels_header( + allowed_fields: tuple[str, ...], headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + labels_header: Final = vertex_partner_labels_header(allowed_fields, None, headers, litellm_params) + unlabelled: Final = {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} + return unlabelled if labels_header is None else {**unlabelled, VERTEX_LABELS_HEADER: labels_header} + + +def _with_vertex_request_labels( + vertex_method: str | None, + headers: Mapping[str, str], + body: Mapping[str, object] | None, + litellm_params: Mapping[str, object], +) -> UpstreamRequest: + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return headers, body + if body is not None and (vertex_method in _VERTEX_BODY_LABELLED_METHODS or "labels" in body): + return headers, _with_vertex_body_labels(allowed_fields, body, litellm_params) + if vertex_method in _VERTEX_HEADER_LABELLED_METHODS: + return _with_vertex_labels_header(allowed_fields, headers, litellm_params), body + return headers, body + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], ) -> BaseVertexAIPassThroughHandler: @@ -3012,6 +3068,9 @@ async def _base_vertex_proxy_route( target=target, custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, + upstream_request_filter=partial( + _with_vertex_request_labels, _vertex_labelled_method(request.method, encoded_endpoint) + ), ) # dynamically construct pass-through endpoint based on incoming path try: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a41e2ea068..aaac97360c8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -465,7 +465,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): url: httpx.URL, headers: dict, requested_query_params: dict | None = None, - _parsed_body: dict | None = None, + _parsed_body: Mapping[str, object] | None = None, forward_multipart: bool = False, ) -> httpx.Response: """ @@ -1087,6 +1087,28 @@ from litellm.passthrough.timeout_utils import ( resolve_pass_through_request_timeout, ) +UpstreamRequest: TypeAlias = tuple[Mapping[str, str], Mapping[str, object] | None] +UpstreamRequestFilter: TypeAlias = Callable[ + [Mapping[str, str], Mapping[str, object] | None, Mapping[str, object]], UpstreamRequest +] +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def _unfiltered_request( + headers: Mapping[str, str], body: Mapping[str, object] | None, litellm_params: Mapping[str, object] +) -> UpstreamRequest: + return headers, body + + +def _filtered_upstream_request( + upstream_request_filter: UpstreamRequestFilter, + headers: Mapping[str, str], + body: Mapping[str, object] | None, + litellm_params: Mapping[str, object], +) -> tuple[dict[str, str], Mapping[str, object] | None]: # mutable-ok: the relay helpers take a dict of headers + filtered_headers, filtered_body = upstream_request_filter(headers, body, litellm_params) + return dict(filtered_headers), filtered_body + async def pass_through_request( request: Request, @@ -1103,6 +1125,7 @@ async def pass_through_request( custom_llm_provider: str | None = None, guardrails_config: dict | None = None, timeout: float | None = None, + upstream_request_filter: UpstreamRequestFilter = _unfiltered_request, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1123,6 +1146,8 @@ async def pass_through_request( guardrails_config: Optional field - guardrails configuration for passthrough endpoint timeout: Optional per-endpoint timeout in seconds. Falls back to general_settings.pass_through_request_timeout, then 600s. + upstream_request_filter: Optional last rewrite of the headers and JSON body sent upstream, given the + litellm_params the request is logged with """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1461,13 +1486,20 @@ async def pass_through_request( else: logging_url = str(url) + "?" + requested_query_params_str + outgoing_body: Final = cast("Mapping[str, object] | None", _parsed_body) # cast-ok: json + outbound_headers, upstream_body = _filtered_upstream_request( + upstream_request_filter, + upstream_headers, + outgoing_body, + _LITELLM_PARAMS.validate_python(kwargs["litellm_params"]), + ) logging_obj.pre_call( input=[{"role": "user", "content": safe_dumps(_parsed_body)}], api_key="", additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": upstream_headers, + "headers": outbound_headers, }, ) stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body( @@ -1484,7 +1516,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=upstream_headers, + headers=outbound_headers, requested_query_params=requested_query_params, stream=True, ) @@ -1496,7 +1528,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=upstream_headers, + headers=outbound_headers, content=state_raw_body, ) if state_raw_body is not None @@ -1504,8 +1536,8 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=upstream_headers, - json=_parsed_body, + headers=outbound_headers, + json=upstream_body, ) ) @@ -1582,7 +1614,7 @@ async def pass_through_request( raw_body_request: Final = async_client.build_request( request.method, url, - headers=upstream_headers, + headers=outbound_headers, params=requested_query_params, content=state_raw_body, ) @@ -1592,9 +1624,9 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=upstream_headers, + headers=outbound_headers, requested_query_params=requested_query_params, - _parsed_body=_parsed_body, + _parsed_body=upstream_body, forward_multipart=is_multipart, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -2107,6 +2139,7 @@ def create_pass_through_route( guardrails: dict[str, object] | None = None, config_file_path: str | None = None, timeout: float | None = None, + upstream_request_filter: UpstreamRequestFilter = _unfiltered_request, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2246,6 +2279,7 @@ def create_pass_through_route( custom_llm_provider=custom_llm_provider, guardrails_config=cast(dict | None, param_guardrails), timeout=cast(float | None, param_timeout), + upstream_request_filter=upstream_request_filter, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 8a03c76f408..fd36b2fcab8 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -6,15 +6,16 @@ Test to verify the Google GenAI generate_content adapter functionality import json from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import Final, Literal import pytest from pydantic import TypeAdapter - - - import litellm +from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams @pytest.mark.asyncio @@ -298,11 +299,10 @@ _PROXY_METADATA: Final = MappingProxyType( {**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}, "requester_metadata": {"team": "search"}} ) _FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) -_GENAI_ENTRYPOINTS: Final = ( - "generate_content", - "agenerate_content", - "generate_content_stream", - "agenerate_content_stream", +_HANDLER_CALLS: Final = pytest.mark.parametrize( + ("is_async", "stream"), + [(False, False), (True, False), (False, True), (True, True)], + ids=["sync", "async", "sync-stream", "async-stream"], ) _GEMINI_REPLY: Final = { "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], @@ -312,32 +312,67 @@ _SENT_BODY: Final = TypeAdapter(Mapping[str, object]) _SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) +class _SignedInVertexGoogleGenAIConfig(VertexAIGoogleGenAIConfig): + def _ensure_access_token( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: + return "token", "proj" + + async def _ensure_access_token_async( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: + return "token", "proj" + + async def _generate_content_body( - entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object], model: str + config: GoogleGenAIConfig, + litellm_params: GenericLiteLLMParams, + extra_body: Mapping[str, object], + *, + is_async: bool, + stream: bool, ) -> Mapping[str, object]: - from unittest.mock import AsyncMock, MagicMock + import datetime + from unittest.mock import MagicMock import httpx - from litellm import google_genai + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler reply: Final = MagicMock(side_effect=lambda _request: httpx.Response(200, json=_GEMINI_REPLY)) - vertex_token: Final = ("token", "proj") - monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token", MagicMock(return_value=vertex_token)) - monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token_async", AsyncMock(return_value=vertex_token)) - is_async: Final = entrypoint.startswith("a") transport: Final = httpx.MockTransport(reply) client: Final = ( AsyncHTTPHandler(transport=transport) if is_async else HTTPHandler(client=httpx.Client(transport=transport)) ) - call: Final = getattr(google_genai, entrypoint)( - model=model, + call: Final = BaseLLMHTTPHandler().generate_content_handler( + model="gemini-3.8-flash", contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + generate_content_provider_config=config, + generate_content_config_dict={}, + tools=None, + custom_llm_provider=config.custom_llm_provider, + litellm_params=litellm_params, + logging_obj=Logging( + model="gemini-3.8-flash", + messages=[], + stream=stream, + call_type="agenerate_content" if is_async else "generate_content", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="labels-test", + function_id="labels-test", + ), + extra_body=dict(extra_body), + _is_async=is_async, client=client, - metadata=dict(_PROXY_METADATA), - **request_fields, + stream=stream, ) if is_async: await call @@ -347,81 +382,77 @@ async def _generate_content_body( async def _vertex_generate_content_body( - entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object] + extra_body: Mapping[str, object], *, is_async: bool, stream: bool ) -> Mapping[str, object]: return await _generate_content_body( - entrypoint, - monkeypatch, - {**request_fields, "vertex_location": "us-central1"}, - "vertex_ai/gemini-3.8-flash", + _SignedInVertexGoogleGenAIConfig(), + GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), vertex_location="us-central1"), + extra_body, + is_async=is_async, + stream=stream, ) @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) -@pytest.mark.parametrize( - "request_fields", - [{"labels": dict(_FORGED)}, {"extra_body": {"labels": dict(_FORGED)}}], - ids=["top-level-labels", "extra-body-labels"], -) +@_HANDLER_CALLS async def test_vertex_generate_content_sends_identity_labels_a_caller_cannot_forge( - entrypoint: str, request_fields: Mapping[str, object], monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + body: Final = await _vertex_generate_content_body({"labels": dict(_FORGED)}, is_async=is_async, stream=stream) labels: Final = _SENT_LABELS.validate_python(body["labels"]) assert list(labels.items()) == [*_IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")] @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS async def test_vertex_generate_content_without_caller_labels_still_carries_identity( - entrypoint: str, monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, {}) + body: Final = await _vertex_generate_content_body({}, is_async=is_async, stream=stream) labels: Final = _SENT_LABELS.validate_python(body["labels"]) assert list(labels.items()) == [*_IDENTITY.items(), ("team", "search"), ("cost-center", "cc-1")] @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS @pytest.mark.parametrize( - "request_fields, expected_labels", - [ - ({}, "absent"), - ({"labels": dict(_FORGED)}, _FORGED), - ({"extra_body": {"labels": dict(_FORGED)}}, _FORGED), - ({"labels": "not-a-mapping"}, "not-a-mapping"), - ], - ids=["no-labels", "top-level-labels", "extra-body-labels", "non-mapping-labels"], + "extra_body, expected_labels", + [({}, "absent"), ({"labels": dict(_FORGED)}, _FORGED), ({"labels": "not-a-mapping"}, "not-a-mapping")], + ids=["no-labels", "labels", "non-mapping-labels"], ) async def test_vertex_generate_content_with_the_setting_off_sends_only_the_callers_labels( - entrypoint: str, - request_fields: Mapping[str, object], + is_async: bool, + stream: bool, + extra_body: Mapping[str, object], expected_labels: object, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", None) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + body: Final = await _vertex_generate_content_body(extra_body, is_async=is_async, stream=stream) assert body.get("labels", "absent") == expected_labels @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS async def test_ai_studio_generate_content_keeps_the_callers_labels_with_the_setting_on( - entrypoint: str, monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) body: Final = await _generate_content_body( - entrypoint, monkeypatch, {"labels": dict(_FORGED), "api_key": "test-key"}, "gemini/gemini-3.8-flash" + GoogleGenAIConfig(), + GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), api_key="test-key"), + {"labels": dict(_FORGED)}, + is_async=is_async, + stream=stream, ) assert body["labels"] == _FORGED diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 9b2f8a40e0e..51838c114a1 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -1,14 +1,16 @@ import base64 import json import os +from collections.abc import Mapping from datetime import datetime from io import BytesIO +from types import MappingProxyType from typing import Dict, Final from unittest.mock import MagicMock, patch import httpx import pytest -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging @@ -636,3 +638,47 @@ def test_imagen_image_edit_maps_malformed_predictions_to_a_connection_error_with 500, False, ) + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType({**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +_FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +_SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +def _sent_gemini_edit_body() -> Mapping[str, object]: + sent: Final[list[httpx.Request]] = [] + + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, json={"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}]} + ) + + litellm.image_edit( + model="vertex_ai/gemini-3.1-flash-image", + image=BytesIO(b"\x89PNG\r\n\x1a\nimage"), + prompt="mountains", + api_base="https://vertex.test/edit", + metadata=dict(_PROXY_METADATA), + labels=dict(_FORGED), + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(reply))), + ) + return _SENT_BODY.validate_json(sent[0].content) + + +def test_vertex_gemini_edit_sends_identity_labels_only_while_the_setting_is_on(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_gemini_edit_body() + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _sent_gemini_edit_body() + + assert sorted(setting_off) == ["contents", "generationConfig"] + assert {key: value for key, value in setting_on.items() if key != "labels"} == setting_off + assert list(_SENT_LABELS.validate_python(setting_on["labels"]).items()) == [ + *_IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] diff --git a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py index 3364b1b3872..661f5200601 100644 --- a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py +++ b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -1,7 +1,18 @@ +import json +from collections.abc import Mapping +from datetime import datetime +from types import MappingProxyType +from typing import Final + +import httpx import pytest +from pydantic import TypeAdapter import litellm +from litellm.interactions.http_handler import InteractionsHTTPHandler from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig, ) @@ -229,3 +240,90 @@ class TestInteractionByIdRequests: assert url == f"{GLOBAL_BASE}/abc123:cancel" assert request_body == {} + + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def _created_interaction_body( + config: VertexAIInteractionsConfig, extra_body: Mapping[str, object], *, is_async: bool +) -> Mapping[str, object]: + sent: Final[list[httpx.Request]] = [] # mutable-ok: MockTransport hands each request to a callback + + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"id": "interaction-1", "status": "completed", "outputs": []}) + + transport: Final = httpx.MockTransport(reply) + litellm_params: Final = GenericLiteLLMParams(vertex_project="test-proj", metadata=dict(PROXY_METADATA)) + logging_obj: Final = Logging( + model="gemini-omni-flash-preview", + messages=[], + stream=False, + call_type="acreate_interaction" if is_async else "create_interaction", + start_time=datetime(2026, 1, 1), + litellm_call_id="labels-test", + function_id="labels-test", + ) + handler: Final = InteractionsHTTPHandler() + if is_async: + await handler.async_create_interaction( + interactions_api_config=config, + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + logging_obj=logging_obj, + model="gemini-omni-flash-preview", + input="hi", + extra_body=extra_body, + client=AsyncHTTPHandler(transport=transport), + ) + else: + handler.create_interaction( + interactions_api_config=config, + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + logging_obj=logging_obj, + model="gemini-omni-flash-preview", + input="hi", + extra_body=extra_body, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + return SENT_BODY.validate_python(json.loads(sent[0].content)) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_create_replaces_forged_extra_body_labels_with_identity_while_the_setting_is_on( + config: VertexAIInteractionsConfig, monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + + body: Final = await _created_interaction_body(config, {"labels": dict(FORGED)}, is_async=is_async) + + assert list(SENT_LABELS.validate_python(body["labels"]).items()) == [ + *IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_create_sends_caller_labels_as_is_while_the_setting_is_off( + config: VertexAIInteractionsConfig, monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + with_labels: Final = await _created_interaction_body(config, {"labels": dict(FORGED)}, is_async=is_async) + without_labels: Final = await _created_interaction_body(config, {}, is_async=is_async) + + assert with_labels["labels"] == dict(FORGED) + assert "labels" not in without_labels diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index cbe990dceb7..0ebbbd2013f 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -131,11 +131,12 @@ def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[st return decoded_labels(claude_headers(litellm_params, optional_params)) -LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels] +LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels] REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels] -SETTING_GATED_SURFACES: Final = [multimodal_embedding_labels, google_genai_labels] +SETTING_GATED_SURFACES: Final = [gemini_image_labels, multimodal_embedding_labels, google_genai_labels] ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels, *SETTING_GATED_SURFACES] CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels, *SETTING_GATED_SURFACES] +BODY_SURFACES_ADDED_BY_THE_SETTING: Final = [gemini_image_labels, multimodal_embedding_labels] def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]: @@ -172,13 +173,16 @@ def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: objec @pytest.mark.parametrize("setting", [None, []]) -def test_setting_off_sends_no_multimodal_embedding_labels(setting: object, monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("surface", BODY_SURFACES_ADDED_BY_THE_SETTING) +def test_setting_off_sends_no_labels_where_none_were_sent_before( + setting: object, surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) params: Final = metadata( spend_logs_metadata={"cost-center": "cc-1"}, requester_metadata={"team": "search"}, **IDENTITY ) - assert multimodal_embedding_labels(params, {"labels": {"app": "bot"}}) is None + assert surface(params, {"labels": {"app": "bot"}}) is None @pytest.mark.parametrize("setting", [None, []]) @@ -273,6 +277,18 @@ def test_gemini_extra_body_labels_cannot_replace_or_drop_identity( assert gemini_labels(metadata(**IDENTITY), {"extra_body": {"labels": extra_labels}}) == expected +def test_gemini_extra_body_labels_outrank_spend_logs_metadata_at_the_cap(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + spend_logs: Final = {f"spend-{index:02d}": "v" for index in range(64)} + caller_labels: Final = {"labels": {"env": "prod"}, "extra_body": {"labels": {"app": "bot"}}} + + labels: Final = gemini_labels(metadata(**IDENTITY, spend_logs_metadata=spend_logs), caller_labels) + + assert isinstance(labels, dict) + assert len(labels) == 64 + assert {key: labels[key] for key in (*IDENTITY, "env", "app")} == {**IDENTITY, "env": "prod", "app": "bot"} + + @pytest.mark.parametrize("setting", [None, []]) def test_setting_off_merges_gemini_extra_body_labels_as_sent(setting: object, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index 15df8e47af3..4a754325be5 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,15 @@ +import base64 +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +import httpx import pytest +from pydantic import TypeAdapter import litellm +import litellm.main +from litellm.llms.custom_httpx.http_handler import HTTPHandler @pytest.fixture @@ -30,3 +39,106 @@ def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model drop_params=True, ) assert "reasoning_effort" not in dropped + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +CHAT_REPLY: Final = MappingProxyType( + { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": "mistral-large-2411", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } +) +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +def _transport(sent: list[httpx.Request]) -> httpx.MockTransport: + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json=dict(CHAT_REPLY)) + + return httpx.MockTransport(reply) + + +def _send_mistral_chat(transport: httpx.MockTransport, monkeypatch: pytest.MonkeyPatch) -> None: + litellm.completion( + model="vertex_ai/mistral-large-2411", + messages=[{"role": "user", "content": "hi"}], + vertex_location="us-central1", + metadata=dict(PROXY_METADATA), + labels=dict(FORGED), + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + +def _send_codestral_fim(transport: httpx.MockTransport, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=transport))) + monkeypatch.setitem( + litellm.known_tokenizer_config, + "codestral-2", + {"status": "success", "tokenizer": {"chat_template": "{{ messages[0]['content'] }}"}}, + ) + litellm.text_completion( + model="vertex_ai/codestral-2", + prompt="def add(", + vertex_location="us-central1", + metadata=dict(PROXY_METADATA), + labels=dict(FORGED), + ) + + +MISTRAL_CALLS: Final = pytest.mark.parametrize( + "send, todays_body", + [ + ( + _send_mistral_chat, + {"model": "mistral-large-2411", "messages": [{"role": "user", "content": "hi"}], "stream": False}, + ), + (_send_codestral_fim, {"model": "codestral-2", "prompt": "def add("}), + ], + ids=["chat", "codestral-fim"], +) + + +def _sent_mistral_request( + send: Callable[[httpx.MockTransport, pytest.MonkeyPatch], None], monkeypatch: pytest.MonkeyPatch +) -> httpx.Request: + monkeypatch.setattr( + litellm.main.vertex_partner_models_chat_completion, "_ensure_access_token", lambda **_: ("token", "proj") + ) + sent: Final[list[httpx.Request]] = [] + send(_transport(sent), monkeypatch) + assert [request.url.path.rsplit(":", 1)[-1] for request in sent] == ["rawPredict"] + return sent[0] + + +@MISTRAL_CALLS +def test_mistral_labels_header_carries_identity_only_while_the_setting_is_on( + send: Callable[[httpx.MockTransport, pytest.MonkeyPatch], None], + todays_body: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_mistral_request(send, monkeypatch) + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + setting_on: Final = _sent_mistral_request(send, monkeypatch) + + assert "x-vertex-ai-labels" not in setting_off.headers + assert SENT_BODY.validate_json(setting_off.content) == {**todays_body, "labels": dict(FORGED)} + assert list(SENT_LABELS.validate_json(base64.b64decode(setting_on.headers["x-vertex-ai-labels"])).items()) == [ + *IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] + assert SENT_BODY.validate_json(setting_on.content) == todays_body + assert (setting_on.headers["content-type"], setting_on.headers["authorization"]) == ( + "application/json", + "Bearer token", + ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..e5e792a832c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -369,6 +369,7 @@ class TestVertexAIPassThroughHandler: target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.asyncio @@ -468,6 +469,7 @@ class TestVertexAIPassThroughHandler: target=f"https://aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.parametrize( @@ -563,6 +565,7 @@ class TestVertexAIPassThroughHandler: target=f"https://{default_location}-aiplatform.googleapis.com/v1/projects/{default_project}/locations/{default_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {default_credentials}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.asyncio 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 98ea2269252..98269df8fa2 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,10 +1,17 @@ import asyncio +import base64 import importlib import json +from collections.abc import Iterator, Mapping from typing import Final +import httpx import pytest +import responses import respx +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi.testclient import TestClient import litellm from litellm import Router @@ -257,3 +264,302 @@ async def test_gemini_pass_through_endpoint_forwards_key_and_body( assert outbound.url.path == f"/{endpoint}" assert outbound.url.params["key"] == "gemini-test-key" assert json.loads(outbound.content) == json.loads(body) + + +_VERTEX_DEPLOYMENT_PATH: Final = ( + "v1/projects/deployment-project/locations/global/publishers/anthropic/models/claude-sonnet-4-5:rawPredict" +) +_VERTEX_CALLER_PATH: Final = ( + "v1/projects/caller-project/locations/global/publishers/anthropic/models/claude-haiku-4-5:rawPredict" +) +_VERTEX_DEPLOYMENT_PROJECT: Final = "deployment-project" +_VERTEX_TOKEN_URI: Final = "https://oauth2.googleapis.com/token" +_VERTEX_MINTED_TOKEN: Final = "minted-vertex-token" +_VERTEX_CALLER_GOOGLE_TOKEN: Final = "caller-google-token" +_VERTEX_CLAUDE_BODY: Final = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} +_VERTEX_CLAUDE_REPLY: Final = {"id": "msg_vertex", "type": "message", "content": []} + + +def _vertex_service_account_json() -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _VERTEX_DEPLOYMENT_PROJECT, + "private_key_id": "test", + "private_key": private_key, + "client_email": f"test@{_VERTEX_DEPLOYMENT_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "token_uri": _VERTEX_TOKEN_URI, + } + ) + + +@pytest.fixture +def vertex_proxy( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, isolate_passthrough_endpoint_router_state: None +) -> Iterator[TestClient]: + from litellm.proxy import proxy_server + + router: Final = Router( + model_list=[ + { + "model_name": "claude-sonnet-4-5", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-4-5", + "vertex_project": _VERTEX_DEPLOYMENT_PROJECT, + "vertex_location": "global", + "vertex_credentials": _vertex_service_account_json(), + "use_in_pass_through": True, + }, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + with responses.RequestsMock(assert_all_requests_are_fired=False) as google_oauth: + google_oauth.post( + _VERTEX_TOKEN_URI, + json={"access_token": _VERTEX_MINTED_TOKEN, "expires_in": 3600, "token_type": "Bearer"}, + ) + yield TestClient(proxy_server.app) + + +_VERTEX_GEMINI_MODEL_PATH: Final = "v1/projects/deployment-project/locations/global/publishers/google/models/gemini-2.5-flash" +_VERTEX_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_VERTEX_IDENTITY_LABELS: Final = {"user_api_key_alias": "search-bot", "user_api_key_team_alias": "search-team"} +_VERTEX_SPEND_LOGS_LABELS: Final = {"cost_center": "cc-42"} +_VERTEX_FORGED_LABELS: Final = {"user_api_key_team_alias": "other-team", "env": "prod", "app": "chat"} +_VERTEX_GEMINI_REPLY: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]} + + +@pytest.fixture +def search_team_caller(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def authenticate(request: object, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key=api_key, + key_alias="search-bot", + team_id="team-search", + team_alias="search-team", + metadata={"spend_logs_metadata": _VERTEX_SPEND_LOGS_LABELS}, + ) + + monkeypatch.setattr(proxy_server, "user_custom_auth", authenticate) + + +def _gemini_upstream(respx_mock: "respx.MockRouter", action: str) -> respx.Route: + reply: Final = ( + httpx.Response(200, json=_VERTEX_GEMINI_REPLY) + if action == "generateContent" + else httpx.Response( + 200, text=f"data: {json.dumps(_VERTEX_GEMINI_REPLY)}\n\n", headers={"content-type": "text/event-stream"} + ) + ) + return respx_mock.post(url__startswith=f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:").mock( + return_value=reply + ) + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("action", ["generateContent", "streamGenerateContent"]) +@pytest.mark.parametrize( + ("caller_labels", "expected_labels"), + [ + pytest.param( + _VERTEX_FORGED_LABELS, + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}, + id="forged-team-replaced-caller-labels-kept", + ), + pytest.param(None, {**_VERTEX_IDENTITY_LABELS, **_VERTEX_SPEND_LOGS_LABELS}, id="identity-added"), + ], +) +def test_vertex_gemini_pass_through_bills_the_callers_own_identity( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + action: str, + caller_labels: Mapping[str, str] | None, + expected_labels: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = _gemini_upstream(respx_mock, action) + body: Final = { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + **({} if caller_labels is None else {"labels": caller_labels}), + } + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:{action}", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + sent: Final = json.loads(upstream.calls.last.request.content) + assert list(sent["labels"].items()) == list(expected_labels.items()) + assert {key: value for key, value in sent.items() if key != "labels"} == { + key: value for key, value in body.items() if key != "labels" + } + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("action", ["generateContent", "streamGenerateContent"]) +def test_vertex_gemini_pass_through_forwards_the_body_untouched_with_the_setting_off( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", action: str +) -> None: + upstream: Final = _gemini_upstream(respx_mock, action) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:{action}", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert list(json.loads(upstream.calls.last.request.content).items()) == list(body.items()) + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize( + "endpoint", + [ + f"{_VERTEX_GEMINI_MODEL_PATH}:predict", + "v1/projects/deployment-project/locations/global/batchPredictionJobs", + ], + ids=["predict", "batch-prediction-job"], +) +def test_vertex_pass_through_adds_identity_labels_to_predict_and_batch_jobs( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={}) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{endpoint}", json={"instances": [{"content": "hi"}]}, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert json.loads(upstream.calls.last.request.content) == { + "instances": [{"content": "hi"}], + "labels": {**_VERTEX_IDENTITY_LABELS, **_VERTEX_SPEND_LOGS_LABELS}, + } + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_pass_through_replaces_forged_labels_on_any_endpoint( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:countTokens").mock( + return_value=httpx.Response(200, json={"totalTokens": 1}) + ) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:countTokens", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert list(json.loads(upstream.calls.last.request.content)["labels"].items()) == list( + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}.items() + ) + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_pass_through_adds_no_labels_where_the_caller_sent_none_outside_billed_methods( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:countTokens").mock( + return_value=httpx.Response(200, json={"totalTokens": 1}) + ) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:countTokens", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert json.loads(upstream.calls.last.request.content) == body + assert "x-vertex-ai-labels" not in upstream.calls.last.request.headers + + +def _labels_header(labels: Mapping[str, str]) -> str: + return base64.b64encode(json.dumps(dict(labels)).encode()).decode() + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize( + ("model_path", "caller_headers"), + [ + pytest.param( + _VERTEX_DEPLOYMENT_PATH, + {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_FORGED_LABELS)}, + id="x-pass-header-on-deployment-credentials", + ), + pytest.param( + _VERTEX_CALLER_PATH, + { + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + "X-Vertex-AI-Labels": _labels_header(_VERTEX_FORGED_LABELS), + }, + id="forwarded-header-on-caller-credentials", + ), + ], +) +def test_vertex_claude_pass_through_replaces_a_forged_labels_header( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + model_path: str, + caller_headers: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{model_path}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{model_path}", json=_VERTEX_CLAUDE_BODY, headers={"x-litellm-api-key": "sk-caller", **caller_headers} + ) + + assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) + outbound: Final = upstream.calls.last.request + sent_headers: Final = outbound.headers.get_list("x-vertex-ai-labels") + assert len(sent_headers) == 1 + assert list(json.loads(base64.b64decode(sent_headers[0])).items()) == list( + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}.items() + ) + assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_claude_pass_through_forwards_the_callers_labels_header_with_the_setting_off( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter" +) -> None: + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_CALLER_PATH}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + forged_header: Final = _labels_header(_VERTEX_FORGED_LABELS) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_CALLER_PATH}", + json=_VERTEX_CLAUDE_BODY, + headers={ + "x-litellm-api-key": "sk-caller", + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + "X-Vertex-AI-Labels": forged_header, + }, + ) + + assert response.status_code == 200 + assert upstream.calls.last.request.headers.get_list("x-vertex-ai-labels") == [forged_header] From 93ac5b5e5a105d4caa0b9a5624f93c80188edca6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 18:42:29 -0800 Subject: [PATCH 06/13] fix(vertex_ai): move body labels into the labels header on Vertex rawPredict pass-through calls (cherry picked from commit e587aa19179a224b575bb68f1d0d93472b334551) --- .../llm_passthrough_endpoints.py | 21 ++++++++++------ .../test_llm_passthrough_endpoints.py | 25 +++++++++++++++++++ 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 22ecbcd0f74..a897404a954 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2599,11 +2599,18 @@ def _with_vertex_body_labels( def _with_vertex_labels_header( - allowed_fields: tuple[str, ...], headers: Mapping[str, str], litellm_params: Mapping[str, object] -) -> Mapping[str, str]: - labels_header: Final = vertex_partner_labels_header(allowed_fields, None, headers, litellm_params) - unlabelled: Final = {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} - return unlabelled if labels_header is None else {**unlabelled, VERTEX_LABELS_HEADER: labels_header} + allowed_fields: tuple[str, ...], + headers: Mapping[str, str], + body: Mapping[str, object] | None, + litellm_params: Mapping[str, object], +) -> UpstreamRequest: + caller_labels: Final = body.get("labels") if body is not None else None + labels_header: Final = vertex_partner_labels_header(allowed_fields, caller_labels, headers, litellm_params) + unlabelled_headers: Final = {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} + unlabelled_body: Final = None if body is None else {key: value for key, value in body.items() if key != "labels"} + if labels_header is None: + return unlabelled_headers, unlabelled_body + return {**unlabelled_headers, VERTEX_LABELS_HEADER: labels_header}, unlabelled_body def _with_vertex_request_labels( @@ -2615,10 +2622,10 @@ def _with_vertex_request_labels( allowed_fields: Final = vertex_request_label_fields() if not allowed_fields: return headers, body + if vertex_method in _VERTEX_HEADER_LABELLED_METHODS: + return _with_vertex_labels_header(allowed_fields, headers, body, litellm_params) if body is not None and (vertex_method in _VERTEX_BODY_LABELLED_METHODS or "labels" in body): return headers, _with_vertex_body_labels(allowed_fields, body, litellm_params) - if vertex_method in _VERTEX_HEADER_LABELLED_METHODS: - return _with_vertex_labels_header(allowed_fields, headers, litellm_params), body return headers, body 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 98269df8fa2..325ade7dab1 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 @@ -542,6 +542,31 @@ def test_vertex_claude_pass_through_replaces_a_forged_labels_header( assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_claude_pass_through_moves_body_labels_into_the_labels_header( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_DEPLOYMENT_PATH}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_DEPLOYMENT_PATH}", + json={**_VERTEX_CLAUDE_BODY, "labels": _VERTEX_FORGED_LABELS}, + headers={"x-litellm-api-key": "sk-caller"}, + ) + + assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) + outbound: Final = upstream.calls.last.request + sent_headers: Final = outbound.headers.get_list("x-vertex-ai-labels") + assert len(sent_headers) == 1 + assert list(json.loads(base64.b64decode(sent_headers[0])).items()) == list( + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}.items() + ) + assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY + + @pytest.mark.usefixtures("search_team_caller") def test_vertex_claude_pass_through_forwards_the_callers_labels_header_with_the_setting_off( vertex_proxy: TestClient, respx_mock: "respx.MockRouter" From 87c731c4530a25c9f4fbea747b19525d77650936 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:32:36 -0800 Subject: [PATCH 07/13] fix(vertex_ai): send Mistral extra_body labels in the labels header instead of the body Labels a caller nests in extra_body now merge over top-level labels before the X-Vertex-AI-Labels header is built, the way the Gemini body merges them, and leave the outbound body. The token-fetch stubs in the Mistral and multimodal embedding tests are scoped to their instance so later tests see the real method (cherry picked from commit bd27e19b5c4a48a3a50eae5051d8be20ef1217a3) --- .../vertex_ai_partner_models/main.py | 17 ++++++++- ..._ai_multimodal_embedding_transformation.py | 16 ++++---- ...i_partner_models_mistral_transformation.py | 38 +++++++++++++++++-- 3 files changed, 58 insertions(+), 13 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index d2459b4afa1..accf24f877a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -25,6 +25,15 @@ base_llm_http_handler = BaseLLMHTTPHandler() _PARAMS: Final = TypeAdapter(Mapping[str, object]) +def _mapping(source: object) -> Mapping[str, object]: + return _PARAMS.validate_python(source) if isinstance(source, Mapping) else {} + + +def _caller_labels(optional_params: Mapping[str, object]) -> Mapping[str, object]: + extra_body_labels: Final = _mapping(optional_params.get("extra_body")).get("labels") + return {**_mapping(optional_params.get("labels")), **_mapping(extra_body_labels)} + + class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -166,10 +175,11 @@ class VertexAIPartnerModels(VertexBase): ) label_fields: Final = vertex_request_label_fields() if partner == VertexPartnerProvider.mistralai else () + params: Final = _PARAMS.validate_python(optional_params) labels_header: Final = ( vertex_partner_labels_header( label_fields, - _PARAMS.validate_python(optional_params).get("labels"), + _caller_labels(params), _PARAMS.validate_python(headers or {}), _PARAMS.validate_python(litellm_params), ) @@ -178,6 +188,11 @@ class VertexAIPartnerModels(VertexBase): ) if label_fields and "labels" in optional_params: del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + extra_body: Final = _mapping(params.get("extra_body")) + if label_fields and "labels" in extra_body: + optional_params["extra_body"] = { # rebind-ok: a copy, so a shared deployment extra_body stays intact + key: value for key, value in extra_body.items() if key != "labels" + } vertex_labels_headers: Final = {} if labels_header is None else {VERTEX_LABELS_HEADER: labels_header} if "codestral" in model or "mistral" in model: diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index c790de4bc47..61455a06009 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -192,12 +192,11 @@ SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) async def multimodal_embedding_body( - is_async: bool, monkeypatch: pytest.MonkeyPatch, request_labels: Mapping[str, object] + is_async: bool, request_labels: Mapping[str, object] ) -> Mapping[str, object]: reply: Final = MagicMock( side_effect=lambda _request: httpx.Response(200, json={"predictions": [{"textEmbedding": [0.1, 0.2]}]}) ) - monkeypatch.setattr(litellm.main.vertex_multimodal_embedding, "_ensure_access_token", lambda **_: ("token", "proj")) transport: Final = httpx.MockTransport(reply) embed: Final = { "model": "vertex_ai/multimodalembedding@001", @@ -206,10 +205,11 @@ async def multimodal_embedding_body( "metadata": dict(PROXY_METADATA), **request_labels, } - if is_async: - await litellm.aembedding(**embed, client=AsyncHTTPHandler(transport=transport)) - else: - litellm.embedding(**embed, client=HTTPHandler(client=httpx.Client(transport=transport))) + with patch.object(litellm.main.vertex_multimodal_embedding, "_ensure_access_token", return_value=("token", "proj")): + if is_async: + await litellm.aembedding(**embed, client=AsyncHTTPHandler(transport=transport)) + else: + litellm.embedding(**embed, client=HTTPHandler(client=httpx.Client(transport=transport))) sent: Final[httpx.Request] = reply.call_args.args[0] return SENT_BODY.validate_json(sent.content) @@ -223,7 +223,7 @@ async def test_setting_off_sends_todays_multimodal_embedding_body( ): monkeypatch.setattr(litellm, "vertex_request_label_fields", None) - assert await multimodal_embedding_body(is_async, monkeypatch, request_labels) == {"instances": TODAYS_INSTANCES} + assert await multimodal_embedding_body(is_async, request_labels) == {"instances": TODAYS_INSTANCES} @pytest.mark.asyncio @@ -244,7 +244,7 @@ async def test_multimodal_embedding_sends_identity_labels_a_caller_cannot_forge( ): monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) - body: Final = await multimodal_embedding_body(is_async, monkeypatch, request_labels) + body: Final = await multimodal_embedding_body(is_async, request_labels) assert body["instances"] == TODAYS_INSTANCES assert list(SENT_LABELS.validate_python(body["labels"]).items()) == list(expected_labels) diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index 4a754325be5..97850690e5e 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -2,6 +2,7 @@ import base64 from collections.abc import Callable, Mapping from types import MappingProxyType from typing import Final +from unittest.mock import patch import httpx import pytest @@ -110,11 +111,11 @@ MISTRAL_CALLS: Final = pytest.mark.parametrize( def _sent_mistral_request( send: Callable[[httpx.MockTransport, pytest.MonkeyPatch], None], monkeypatch: pytest.MonkeyPatch ) -> httpx.Request: - monkeypatch.setattr( - litellm.main.vertex_partner_models_chat_completion, "_ensure_access_token", lambda **_: ("token", "proj") - ) sent: Final[list[httpx.Request]] = [] - send(_transport(sent), monkeypatch) + with patch.object( + litellm.main.vertex_partner_models_chat_completion, "_ensure_access_token", return_value=("token", "proj") + ): + send(_transport(sent), monkeypatch) assert [request.url.path.rsplit(":", 1)[-1] for request in sent] == ["rawPredict"] return sent[0] @@ -142,3 +143,32 @@ def test_mistral_labels_header_carries_identity_only_while_the_setting_is_on( "application/json", "Bearer token", ) + + +def test_mistral_extra_body_labels_join_the_labels_header_instead_of_the_body(monkeypatch: pytest.MonkeyPatch): + def send(transport: httpx.MockTransport, _: pytest.MonkeyPatch) -> None: + litellm.completion( + model="vertex_ai/mistral-large-2411", + messages=[{"role": "user", "content": "hi"}], + vertex_location="us-central1", + metadata=dict(PROXY_METADATA), + labels={"app": "bot"}, + extra_body={"labels": {"app": "agent", "env": "prod"}, "safe_prompt": True}, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + sent: Final = _sent_mistral_request(send, monkeypatch) + + assert list(SENT_LABELS.validate_json(base64.b64decode(sent.headers["x-vertex-ai-labels"])).items()) == [ + *IDENTITY.items(), + ("app", "agent"), + ("env", "prod"), + ("cost-center", "cc-1"), + ] + assert SENT_BODY.validate_json(sent.content) == { + "model": "mistral-large-2411", + "messages": [{"role": "user", "content": "hi"}], + "stream": False, + "safe_prompt": True, + } From 1a4c6ca36a2878cc2638f7684012d23406f18e15 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 08:36:30 -0800 Subject: [PATCH 08/13] fix(vertex_ai): limit pass-through labels to POST and to publishers that take the labels header rawPredict to custom endpoints and publishers other than Anthropic and Mistral, and PATCH or PUT bodies, now reach Vertex as the caller sent them. The pass-through drops a caller's labels header where it writes body labels, a single field written as text turns the setting on for that field, and an undecodable Claude labels header falls back to the caller's labels (cherry picked from commit ec655f4ce6f5bc6e13cd7ecb1a35bccaf7a5e4d6) --- litellm/llms/vertex_ai/request_labels.py | 10 +- .../llm_passthrough_endpoints.py | 53 ++++++--- .../llms/vertex_ai/test_request_labels.py | 21 ++++ .../test_llm_passthrough_endpoints.py | 102 ++++++++++++++++++ 4 files changed, 168 insertions(+), 18 deletions(-) diff --git a/litellm/llms/vertex_ai/request_labels.py b/litellm/llms/vertex_ai/request_labels.py index f1e0b3da99b..605023754bc 100644 --- a/litellm/llms/vertex_ai/request_labels.py +++ b/litellm/llms/vertex_ai/request_labels.py @@ -37,6 +37,8 @@ Label: TypeAlias = tuple[str, str] def vertex_request_label_fields() -> tuple[str, ...]: configured: Final[object] = litellm.vertex_request_label_fields + if isinstance(configured, str): + return (configured,) if configured else () if not isinstance(configured, (list, tuple)): return () return tuple(field for field in _ANY_SEQUENCE.validate_python(configured) if isinstance(field, str)) @@ -143,15 +145,17 @@ def vertex_partner_labels_header( litellm_params: Mapping[str, object], ) -> str | None: """ - The ``X-Vertex-AI-Labels`` value for a partner model call once ``allowed_fields`` is non-empty. A labels header the - caller already set stands in for ``caller_labels``, as it does with the setting off, and the result replaces it. + The ``X-Vertex-AI-Labels`` value for a partner model call once ``allowed_fields`` is non-empty. A decodable labels + header the caller already set stands in for ``caller_labels``, as it does with the setting off, and the result + replaces it. """ caller_header: Final = next( (value for name, value in outbound_headers.items() if is_vertex_labels_header(name)), None ) + decoded_header: Final = _decoded_labels_header(caller_header) resolved: Final = resolve_vertex_request_labels( allowed_fields, - caller_labels if caller_header is None else _decoded_labels_header(caller_header), + caller_labels if decoded_header is None else decoded_header, litellm_params, VERTEX_PARTNER_REQUEST_LABELS_MAX, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a897404a954..938305d9f2d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -18,8 +18,9 @@ import sys from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial +from itertools import pairwise from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -2579,23 +2580,41 @@ _VERTEX_BODY_LABELLED_METHODS: Final = frozenset( {"generateContent", "streamGenerateContent", "predict", "batchPredictionJobs"} ) _VERTEX_HEADER_LABELLED_METHODS: Final = frozenset({"rawPredict", "streamRawPredict"}) +_VERTEX_HEADER_LABELLED_PUBLISHERS: Final = frozenset({"anthropic", "mistralai"}) + +VertexLabelPlacement: TypeAlias = Literal["header", "body", "caller_body"] -def _vertex_labelled_method(http_method: str, endpoint: str) -> str | None: +def _vertex_publisher(segments: tuple[str, ...]) -> str | None: + return next((value for name, value in pairwise(segments) if name == "publishers"), None) + + +def _vertex_label_placement(http_method: str, endpoint: str) -> VertexLabelPlacement | None: if http_method != "POST": return None - resource: Final = endpoint.rstrip("/").rsplit("/", 1)[-1] - return resource.partition(":")[2] or resource + segments: Final = tuple(endpoint.strip("/").split("/")) + method: Final = segments[-1].partition(":")[2] or segments[-1] + if method in _VERTEX_HEADER_LABELLED_METHODS: + return "header" if _vertex_publisher(segments) in _VERTEX_HEADER_LABELLED_PUBLISHERS else None + return "body" if method in _VERTEX_BODY_LABELLED_METHODS else "caller_body" + + +def _without_vertex_labels_header(headers: Mapping[str, str]) -> Mapping[str, str]: + return {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} def _with_vertex_body_labels( - allowed_fields: tuple[str, ...], body: Mapping[str, object], litellm_params: Mapping[str, object] -) -> Mapping[str, object]: + allowed_fields: tuple[str, ...], + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> UpstreamRequest: resolved: Final = resolve_vertex_request_labels( allowed_fields, body.get("labels"), litellm_params, VERTEX_REQUEST_LABELS_MAX ) unlabelled: Final = {key: value for key, value in body.items() if key != "labels"} - return {**unlabelled, "labels": dict(resolved)} if resolved else unlabelled + labelled_body: Final = {**unlabelled, "labels": dict(resolved)} if resolved else unlabelled + return _without_vertex_labels_header(headers), labelled_body def _with_vertex_labels_header( @@ -2606,7 +2625,7 @@ def _with_vertex_labels_header( ) -> UpstreamRequest: caller_labels: Final = body.get("labels") if body is not None else None labels_header: Final = vertex_partner_labels_header(allowed_fields, caller_labels, headers, litellm_params) - unlabelled_headers: Final = {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} + unlabelled_headers: Final = _without_vertex_labels_header(headers) unlabelled_body: Final = None if body is None else {key: value for key, value in body.items() if key != "labels"} if labels_header is None: return unlabelled_headers, unlabelled_body @@ -2614,7 +2633,7 @@ def _with_vertex_labels_header( def _with_vertex_request_labels( - vertex_method: str | None, + placement: VertexLabelPlacement | None, headers: Mapping[str, str], body: Mapping[str, object] | None, litellm_params: Mapping[str, object], @@ -2622,11 +2641,15 @@ def _with_vertex_request_labels( allowed_fields: Final = vertex_request_label_fields() if not allowed_fields: return headers, body - if vertex_method in _VERTEX_HEADER_LABELLED_METHODS: - return _with_vertex_labels_header(allowed_fields, headers, body, litellm_params) - if body is not None and (vertex_method in _VERTEX_BODY_LABELLED_METHODS or "labels" in body): - return headers, _with_vertex_body_labels(allowed_fields, body, litellm_params) - return headers, body + match placement: + case "header": + return _with_vertex_labels_header(allowed_fields, headers, body, litellm_params) + case "body" if body is not None: + return _with_vertex_body_labels(allowed_fields, headers, body, litellm_params) + case "caller_body" if body is not None and "labels" in body: + return _with_vertex_body_labels(allowed_fields, headers, body, litellm_params) + case _: + return headers, body def get_vertex_pass_through_handler( @@ -3076,7 +3099,7 @@ async def _base_vertex_proxy_route( custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, upstream_request_filter=partial( - _with_vertex_request_labels, _vertex_labelled_method(request.method, encoded_endpoint) + _with_vertex_request_labels, _vertex_label_placement(request.method, encoded_endpoint) ), ) # dynamically construct pass-through endpoint based on incoming path diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index 0ebbbd2013f..8a3cf1fe638 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -309,6 +309,27 @@ def test_a_caller_set_claude_labels_header_cannot_forge_identity_and_is_replaced assert decoded_labels(undecodable) == IDENTITY +def test_an_undecodable_claude_labels_header_keeps_the_callers_labels(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + + sent: Final = claude_headers(metadata(**IDENTITY), {"labels": {"app": "bot"}}, {LABELS_HEADER: "%%"}) + + assert decoded_labels(sent) == {**IDENTITY, "app": "bot"} + + +@pytest.mark.parametrize( + ("setting", "expected"), + [("user_api_key_team_alias", {"user_api_key_team_alias": "ai-platform"}), ("", None)], + ids=["one-field", "empty"], +) +def test_a_single_field_written_as_text_turns_the_setting_on_for_that_field( + setting: str, expected: dict[str, str] | None, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + + assert gemini_labels(metadata(**IDENTITY), {}) == expected + + @pytest.mark.parametrize("surface", ALL_SURFACES) def test_identity_and_spend_logs_text_is_rewritten_to_labels_google_accepts( surface: LabelSource, monkeypatch: pytest.MonkeyPatch 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 325ade7dab1..4141609cbaa 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 @@ -588,3 +588,105 @@ def test_vertex_claude_pass_through_forwards_the_callers_labels_header_with_the_ assert response.status_code == 200 assert upstream.calls.last.request.headers.get_list("x-vertex-ai-labels") == [forged_header] + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_mistral_pass_through_moves_body_labels_into_the_labels_header( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + endpoint: Final = ( + "v1/projects/deployment-project/locations/global/publishers/mistralai/models/mistral-medium-3:rawPredict" + ) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={}) + ) + body: Final = {"model": "mistral-medium-3", "messages": [{"role": "user", "content": "hi"}]} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{endpoint}", + json={**body, "labels": _VERTEX_FORGED_LABELS}, + headers={"x-litellm-api-key": "sk-caller"}, + ) + + assert response.status_code == 200 + outbound: Final = upstream.calls.last.request + assert json.loads(base64.b64decode(outbound.headers["x-vertex-ai-labels"])) == { + **_VERTEX_IDENTITY_LABELS, + "env": "prod", + "app": "chat", + **_VERTEX_SPEND_LOGS_LABELS, + } + assert json.loads(outbound.content) == body + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize( + "endpoint", + [ + "v1/projects/deployment-project/locations/global/endpoints/123:rawPredict", + "v1/projects/deployment-project/locations/global/publishers/ai21/models/jamba-1.5-mini:streamRawPredict", + ], + ids=["custom-endpoint", "publisher-without-a-labels-header"], +) +def test_vertex_raw_predict_pass_through_leaves_other_publishers_as_the_caller_sent_them( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={}) + ) + body: Final = {"instances": [{"content": "hi"}], "labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.post(f"/vertex_ai/{endpoint}", json=body, headers={"x-litellm-api-key": "sk-caller"}) + + assert response.status_code == 200 + outbound: Final = upstream.calls.last.request + assert json.loads(outbound.content) == body + assert "x-vertex-ai-labels" not in outbound.headers + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("http_method", ["PATCH", "PUT"]) +def test_vertex_pass_through_leaves_labels_on_non_post_requests_as_the_caller_sent_them( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, http_method: str +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + endpoint: Final = "v1/projects/deployment-project/locations/global/endpoints/123" + upstream: Final = respx_mock.route(method=http_method, url=f"https://aiplatform.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={}) + ) + body: Final = {"labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.request( + http_method, f"/vertex_ai/{endpoint}", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert json.loads(upstream.calls.last.request.content) == body + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("action", ["generateContent", "countTokens"]) +def test_vertex_gemini_pass_through_drops_a_callers_labels_header_where_it_writes_body_labels( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:{action}").mock( + return_value=httpx.Response(200, json=_VERTEX_GEMINI_REPLY) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:{action}", + json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "labels": {"env": "prod"}}, + headers={"x-litellm-api-key": "sk-caller", "x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_FORGED_LABELS)}, + ) + + assert response.status_code == 200 + outbound: Final = upstream.calls.last.request + assert "x-vertex-ai-labels" not in outbound.headers + assert json.loads(outbound.content)["labels"] == { + **_VERTEX_IDENTITY_LABELS, + "env": "prod", + **_VERTEX_SPEND_LOGS_LABELS, + } From 807276953987b1b6cd8e5593ce0b9a25d0ae2263 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 10:35:59 -0800 Subject: [PATCH 09/13] fix(vertex): label Veo video, Gemini transcription, and pass-through predictLongRunning requests (cherry picked from commit 1cfcee2ee07beabd7fa88a6c67da9ab323b4e8dd) --- .../gemini_transcribe_transformation.py | 6 ++- litellm/llms/vertex_ai/common_utils.py | 30 +++++--------- litellm/llms/vertex_ai/request_labels.py | 10 ++--- .../llms/vertex_ai/videos/transformation.py | 7 ++++ .../llm_passthrough_endpoints.py | 2 +- ...tex_ai_gemini_transcribe_transformation.py | 36 +++++++++++++++++ .../test_vertex_video_transformation.py | 39 ++++++++++++++++++- .../test_llm_passthrough_endpoints.py | 3 +- 8 files changed, 102 insertions(+), 31 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index 23476e0e52e..09acc6384e1 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -20,7 +20,8 @@ from litellm.llms.vertex_ai.audio_transcription.transformation import ( validate_vertex_transcription_location, validate_vertex_transcription_project_id, ) -from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url, pop_vertex_request_labels +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.openai import ( AllMessageValues, @@ -169,7 +170,8 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB audioTranscriptionConfig=_audio_transcription_config(optional_params.get("language")) ), ) - return AudioTranscriptionRequestData(data=dict(request_body)) + labels: Final = pop_vertex_request_labels(None, dict(litellm_params)) if vertex_request_label_fields() else None + return AudioTranscriptionRequestData(data={**request_body, "labels": labels} if labels else dict(request_body)) def transform_audio_transcription_response( self, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 93f319a021d..812e1420997 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -20,6 +20,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.vertex_ai.request_labels import ( VERTEX_REQUEST_LABELS_MAX, resolve_vertex_request_labels, + string_labels, + vertex_metadata_sources, vertex_request_label_fields, ) from litellm.types.llms.openai import AllMessageValues @@ -171,6 +173,7 @@ def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> N _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_CALLER_LABELS: Final = TypeAdapter(object) def vertex_request_labels_from_litellm_params( @@ -182,23 +185,9 @@ def vertex_request_labels_from_litellm_params( using ``requester_metadata`` string key-value pairs (same convention as Gemini). ``metadata`` is tried first when both are present. """ - if not litellm_params: - return None - for key in ("metadata", "litellm_metadata"): - if key not in litellm_params: - continue - metadata = litellm_params[key] - if metadata is None or not isinstance(metadata, dict): - continue - if "requester_metadata" not in metadata: - continue - rm = metadata["requester_metadata"] - if not isinstance(rm, dict): - continue - labels = {k: v for k, v in rm.items() if isinstance(k, str) and isinstance(v, str)} - if labels: - return labels - return None + sources: Final = vertex_metadata_sources(litellm_params or {}) + candidates: Final = (dict(string_labels(source.get("requester_metadata"))) for source in sources) + return next((labels for labels in candidates if labels), None) def pop_vertex_caller_labels( @@ -209,11 +198,10 @@ def pop_vertex_caller_labels( Pop the caller's optional ``labels`` (Gemini-style), falling back to ``requester_metadata`` on ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``. """ - raw: Final = optional_params.pop("labels", None) if optional_params is not None else None - labels: Final = ( - {k: v for k, v in raw.items() if isinstance(k, str) and isinstance(v, str)} if isinstance(raw, dict) else None + raw: Final = ( + _CALLER_LABELS.validate_python(optional_params.pop("labels", None)) if optional_params is not None else None ) - return labels or vertex_request_labels_from_litellm_params(litellm_params) + return dict(string_labels(raw)) or vertex_request_labels_from_litellm_params(litellm_params) def pop_vertex_request_labels( diff --git a/litellm/llms/vertex_ai/request_labels.py b/litellm/llms/vertex_ai/request_labels.py index 605023754bc..b3b9b4b2c33 100644 --- a/litellm/llms/vertex_ai/request_labels.py +++ b/litellm/llms/vertex_ai/request_labels.py @@ -64,12 +64,12 @@ def _is_reserved(key: str) -> bool: return key.strip().lower().startswith(VERTEX_REQUEST_LABEL_IDENTITY_PREFIX) -def _metadata_sources(litellm_params: Mapping[str, object]) -> tuple[Mapping[object, object], ...]: +def vertex_metadata_sources(litellm_params: Mapping[str, object]) -> tuple[Mapping[object, object], ...]: candidates: Final = (litellm_params.get(name) for name in _METADATA_PARAM_NAMES) return tuple(_ANY_MAPPING.validate_python(source) for source in candidates if isinstance(source, Mapping)) -def _text_labels(source: object) -> tuple[Label, ...]: +def string_labels(source: object) -> tuple[Label, ...]: if not isinstance(source, Mapping): return () pairs: Final = _ANY_MAPPING.validate_python(source).items() @@ -96,7 +96,7 @@ def _identity_labels( def _spend_logs_labels(sources: tuple[Mapping[object, object], ...]) -> tuple[Label, ...]: pairs: Final = chain.from_iterable( - _text_labels(source.get(VERTEX_REQUEST_LABEL_CLIENT_FIELD)) for source in sources + string_labels(source.get(VERTEX_REQUEST_LABEL_CLIENT_FIELD)) for source in sources ) candidates: Final = (_normalized_label(key, value) for key, value in pairs) return tuple(label for label in candidates if label is not None) @@ -119,9 +119,9 @@ def resolve_vertex_request_labels( Identity comes first, then the caller's string labels minus the reserved prefix, then ``spend_logs_metadata`` pairs when allowed, each key keeping its first value, up to ``max_labels``. """ - sources: Final = _metadata_sources(litellm_params) + sources: Final = vertex_metadata_sources(litellm_params) spend_logs: Final = _spend_logs_labels(sources) if VERTEX_REQUEST_LABEL_CLIENT_FIELD in allowed_fields else () - client: Final = tuple(label for label in (*_text_labels(caller_labels), *spend_logs) if not _is_reserved(label[0])) + client: Final = tuple(label for label in (*string_labels(caller_labels), *spend_logs) if not _is_reserved(label[0])) return _first_occurrences((*_identity_labels(sources, allowed_fields), *client), max_labels) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index de368cea5ac..5bbeb2bf304 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -22,7 +22,9 @@ from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, get_vertex_base_url, + pop_vertex_request_labels, ) +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject @@ -368,6 +370,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Build instance with prompt instance_dict: Final[dict[str, object]] = {"prompt": prompt} params_copy: Final = video_create_optional_request_params.copy() + labels: Final = ( + pop_vertex_request_labels(params_copy, dict(litellm_params)) if vertex_request_label_fields() else None + ) # Check if user wants to provide full instance dict if "instances" in params_copy and isinstance(params_copy["instances"], dict): @@ -412,6 +417,8 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Only add parameters if there are any if vertex_params: request_data["parameters"] = vertex_params + if labels: + request_data["labels"] = labels # Append :predictLongRunning endpoint to api_base url: Final = f"{api_base}:predictLongRunning" diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 938305d9f2d..2b01ce29aae 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2577,7 +2577,7 @@ def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str] _VERTEX_BODY_LABELLED_METHODS: Final = frozenset( - {"generateContent", "streamGenerateContent", "predict", "batchPredictionJobs"} + {"generateContent", "streamGenerateContent", "predict", "predictLongRunning", "batchPredictionJobs"} ) _VERTEX_HEADER_LABELLED_METHODS: Final = frozenset({"rawPredict", "streamRawPredict"}) _VERTEX_HEADER_LABELLED_PUBLISHERS: Final = frozenset({"anthropic", "mistralai"}) diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index 532e013c3d7..584a0b42eef 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -1,6 +1,8 @@ import base64 import json import os +from types import MappingProxyType +from typing import Final, cast import httpx import pytest @@ -324,3 +326,37 @@ class TestOptionalParams: class TestModelCostEntry: REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType({**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) + + +def _transcription_request_body() -> dict[str, object]: + data: Final = ( + VertexGeminiAudioTranscriptionConfig() + .transform_audio_transcription_request( + model="gemini-3.5-transcribe-preview", + audio_file=AUDIO_BYTES, + optional_params={"language": "en"}, + litellm_params={"metadata": dict(_PROXY_METADATA)}, + ) + .data + ) + assert isinstance(data, dict) + return data + + +def test_gemini_transcription_carries_identity_labels_only_while_the_setting_is_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _transcription_request_body() + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _transcription_request_body() + + assert "labels" not in setting_off + assert {key: value for key, value in setting_on.items() if key != "labels"} == setting_off + assert setting_on["labels"] == {**_IDENTITY, "cost-center": "cc-1"} + assert list(cast(dict[str, str], setting_on["labels"])) == [*_IDENTITY, "cost-center"] diff --git a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py index 5c90d54ae90..3089fbd1537 100644 --- a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -5,8 +5,9 @@ Tests for Vertex AI (Veo) video generation transformation. import base64 import json from collections.abc import Mapping +from types import MappingProxyType from pathlib import Path -from typing import cast +from typing import Final, cast from unittest.mock import Mock, patch import httpx @@ -970,3 +971,39 @@ class TestImageAndParametersPassthrough: assert "parameters" not in data["parameters"] assert url.endswith(":predictLongRunning") + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType({**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +_FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) + + +def _veo_request_body() -> Mapping[str, object]: + body, _, _ = VertexAIVideoConfig().transform_video_create_request( + model="veo-3.1-generate-001", + prompt="a lighthouse at dusk", + api_base="https://vertex.test/models/veo-3.1-generate-001", + video_create_optional_request_params={"labels": dict(_FORGED), "durationSeconds": 4}, + litellm_params=GenericLiteLLMParams(metadata=dict(_PROXY_METADATA)), + headers={}, + ) + return cast(Mapping[str, object], body) + + +def test_veo_request_carries_identity_labels_only_while_the_setting_is_on(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _veo_request_body() + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _veo_request_body() + + assert setting_off == { + "instances": [{"prompt": "a lighthouse at dusk"}], + "parameters": {"labels": dict(_FORGED), "durationSeconds": 4}, + } + assert setting_on == { + "instances": [{"prompt": "a lighthouse at dusk"}], + "parameters": {"durationSeconds": 4}, + "labels": {**_IDENTITY, "app": "bot", "cost-center": "cc-1"}, + } + assert list(cast(Mapping[str, str], setting_on["labels"])) == [*_IDENTITY, "app", "cost-center"] 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 4141609cbaa..8e8bca317ec 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 @@ -431,9 +431,10 @@ def test_vertex_gemini_pass_through_forwards_the_body_untouched_with_the_setting "endpoint", [ f"{_VERTEX_GEMINI_MODEL_PATH}:predict", + f"{_VERTEX_GEMINI_MODEL_PATH}:predictLongRunning", "v1/projects/deployment-project/locations/global/batchPredictionJobs", ], - ids=["predict", "batch-prediction-job"], + ids=["predict", "predict-long-running", "batch-prediction-job"], ) def test_vertex_pass_through_adds_identity_labels_to_predict_and_batch_jobs( vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, endpoint: str From 23fea0ead9611c0d47b7c2b0801660bb9c4353dd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:26:05 -0800 Subject: [PATCH 10/13] fix(vertex): read the pass-through publisher from the end of the upstream path A project id named publishers shifted the first-match lookup, so a forged X-Vertex-AI-Labels header reached Claude or Mistral untouched (cherry picked from commit 5f4e9c8e8c31311ed985cb0e1254a5d1ee1baf8b) --- .../pass_through_endpoints/llm_passthrough_endpoints.py | 9 ++++++--- .../test_llm_passthrough_endpoints.py | 8 ++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2b01ce29aae..74c3606d436 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -18,7 +18,6 @@ import sys from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial -from itertools import pairwise from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast @@ -2586,7 +2585,11 @@ VertexLabelPlacement: TypeAlias = Literal["header", "body", "caller_body"] def _vertex_publisher(segments: tuple[str, ...]) -> str | None: - return next((value for name, value in pairwise(segments) if name == "publishers"), None) + match segments[-4:]: + case ("publishers", publisher, "models", _): + return publisher + case _: + return None def _vertex_label_placement(http_method: str, endpoint: str) -> VertexLabelPlacement | None: @@ -3099,7 +3102,7 @@ async def _base_vertex_proxy_route( custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, upstream_request_filter=partial( - _with_vertex_request_labels, _vertex_label_placement(request.method, encoded_endpoint) + _with_vertex_request_labels, _vertex_label_placement(request.method, updated_url.path) ), ) # 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 8e8bca317ec..f55f09d9f68 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 @@ -515,6 +515,14 @@ def _labels_header(labels: Mapping[str, str]) -> str: }, id="forwarded-header-on-caller-credentials", ), + pytest.param( + "v1/projects/publishers/locations/global/publishers/anthropic/models/claude-haiku-4-5:rawPredict", + { + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + "X-Vertex-AI-Labels": _labels_header(_VERTEX_FORGED_LABELS), + }, + id="project-id-named-publishers", + ), ], ) def test_vertex_claude_pass_through_replaces_a_forged_labels_header( From b510e84be8cd15e3fcf27a29ce82b8691e15307b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:18:24 -0800 Subject: [PATCH 11/13] fix(vertex): drop a caller labels header on body-labelled Vertex routes while identity labels are on (cherry picked from commit f4f85d36099406f809f32aca0376983c0abf5fe0) --- .../base_llm/google_genai/transformation.py | 3 + litellm/llms/custom_httpx/llm_http_handler.py | 4 +- .../gemini_transcribe_transformation.py | 4 +- .../vertex_and_google_ai_studio_gemini.py | 21 +++++-- .../vertex_ai/google_genai/transformation.py | 4 ++ .../vertex_gemini_transformation.py | 4 +- .../vertex_gemini_transformation.py | 6 +- .../vertex_imagen_transformation.py | 5 +- .../vertex_ai/interactions/transformation.py | 3 +- litellm/llms/vertex_ai/request_labels.py | 18 +++++- .../llms/vertex_ai/rerank/transformation.py | 5 +- .../llms/vertex_ai/videos/transformation.py | 6 +- ...test_vertex_request_label_identity_wire.py | 36 ++++++++++++ ...tex_ai_gemini_transcribe_transformation.py | 30 ++++++++++ ...est_vertex_ai_image_edit_transformation.py | 26 ++++++++- ...rtex_ai_image_generation_transformation.py | 52 +++++++++++++++++ ...t_vertex_ai_interactions_transformation.py | 14 +++++ .../test_vertex_ai_rerank_transformation.py | 56 +++++++++++++++++++ .../test_vertex_video_transformation.py | 26 +++++++++ 19 files changed, 300 insertions(+), 23 deletions(-) diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 51a6177cfc0..fbcbeee9025 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -82,6 +82,9 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): ) -> Mapping[str, object] | None: return request_fields + def transform_caller_headers(self, headers: Mapping[str, object]) -> Mapping[str, object]: + return headers + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ce4109e732c..35ab09fe47c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12116,7 +12116,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(generate_content_provider_config.transform_caller_headers(extra_headers)) # Get the request body from the provider config data: Final = generate_content_provider_config.transform_generate_content_request( @@ -12234,7 +12234,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(generate_content_provider_config.transform_caller_headers(extra_headers)) # Get the request body from the provider config data: Final = generate_content_provider_config.transform_generate_content_request( diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index 09acc6384e1..63d2c6dac22 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -21,7 +21,7 @@ from litellm.llms.vertex_ai.audio_transcription.transformation import ( validate_vertex_transcription_project_id, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url, pop_vertex_request_labels -from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields, without_caller_labels_header from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.openai import ( AllMessageValues, @@ -108,7 +108,7 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB custom_llm_provider="vertex_ai", ) return { - **headers, + **without_caller_labels_header(headers), "Authorization": f"Bearer {access_token}", "x-goog-user-project": project_id, "Content-Type": "application/json", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d87de7ecc50..b50adc6f0e0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -11,7 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import litellm from litellm import verbose_logger @@ -98,6 +98,7 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..request_labels import vertex_request_label_fields, without_caller_labels_header from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .long_decode import ( @@ -202,6 +203,9 @@ def _served_model_name(model_version: object) -> str | None: return model_version.split("@", 1)[0] +_CALLER_HEADERS: Final = TypeAdapter(dict[str, str]) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -3238,6 +3242,11 @@ class VertexLLM(VertexBase): api_base: str | None = None, ) -> ModelResponse | CustomStreamWrapper: stream: Final[bool | None] = optional_params.pop("stream", None) + caller_headers: Final = ( + dict(without_caller_labels_header(_CALLER_HEADERS.validate_python(extra_headers))) + if extra_headers is not None and custom_llm_provider != "gemini" and vertex_request_label_fields() + else extra_headers + ) transform_request_params: Final = { "gemini_api_key": gemini_api_key, @@ -3246,7 +3255,7 @@ class VertexLLM(VertexBase): "model": model, "client": client, "timeout": timeout, - "extra_headers": extra_headers, + "extra_headers": caller_headers, "optional_params": optional_params, "logging_obj": logging_obj, "custom_llm_provider": custom_llm_provider, @@ -3277,7 +3286,7 @@ class VertexLLM(VertexBase): vertex_credentials=vertex_credentials, gemini_api_key=gemini_api_key, custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, + extra_headers=caller_headers, ) ### ASYNC COMPLETION return self.async_completion( @@ -3300,7 +3309,7 @@ class VertexLLM(VertexBase): vertex_credentials=vertex_credentials, gemini_api_key=gemini_api_key, custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, + extra_headers=caller_headers, ) _auth_header, vertex_project = self._ensure_access_token( @@ -3326,7 +3335,7 @@ class VertexLLM(VertexBase): ) headers: Final = VertexGeminiConfig().validate_environment( api_key=auth_header, - headers=extra_headers, + headers=caller_headers, model=model, messages=messages, optional_params=optional_params, @@ -3353,7 +3362,7 @@ class VertexLLM(VertexBase): ) fresh_headers: Final = partial( - self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider + self._fresh_headers, headers, caller_headers, vertex_credentials, vertex_project, custom_llm_provider ) ## SYNC STREAMING CALL ## diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 9982763a846..b5ea0ca7c65 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -7,6 +7,7 @@ from typing import Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig from litellm.llms.vertex_ai.common_utils import vertex_top_level_request_fields +from litellm.llms.vertex_ai.request_labels import without_caller_labels_header from litellm.types.router import GenericLiteLLMParams @@ -47,6 +48,9 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): ) -> Mapping[str, object] | None: return vertex_top_level_request_fields(request_fields, litellm_params) + def transform_caller_headers(self, headers: Mapping[str, object]) -> Mapping[str, object]: + return without_caller_labels_header(headers) + def _camel_to_snake(self, camel_str: str) -> str: """Convert camelCase to snake_case""" import re diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 09c8596043d..700d19ace5c 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -14,7 +14,7 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, vertex_top_level_request_fields from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM -from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields +from litellm.llms.vertex_ai.request_labels import caller_labels_header_names, vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.vertex_ai import ( @@ -169,6 +169,8 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): litellm_params: GenericLiteLLMParams, headers: dict[str, str], ) -> tuple[dict[str, object], RequestFiles | None]: + for name in caller_labels_header_names(headers): + del headers[name] # rebind-ok: the handler posts this dict, so a caller-set labels header must go inline_parts: Final = self._prepare_inline_image_parts(image) if image else [] if not inline_parts: raise ValueError("Vertex AI Gemini image edit requires at least one image.") diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 98056e47aa1..039118a04e4 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -17,7 +17,7 @@ from litellm.llms.gemini.common_utils import ( from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, pop_vertex_request_labels from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM -from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields +from litellm.llms.vertex_ai.request_labels import caller_labels_header_names, vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( @@ -214,13 +214,15 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): prompt: str, optional_params: dict, litellm_params: dict, - headers: dict, + headers: dict[str, str], ) -> dict: """ Transform the image generation request to Gemini format Uses generateContent API with responseModalities: ["IMAGE"] """ + for name in caller_labels_header_names(headers): + del headers[name] # rebind-ok: the handler posts this dict, so a caller-set labels header must go # Prepare messages with the prompt contents: Final = [{"role": "user", "parts": [{"text": prompt}]}] diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 927c7b74e67..e0e88b6b051 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -14,6 +14,7 @@ from litellm.llms.vertex_ai.common_utils import ( pop_vertex_request_labels, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import caller_labels_header_names from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -187,13 +188,15 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): prompt: str, optional_params: dict, litellm_params: dict, - headers: dict, + headers: dict[str, str], ) -> dict: """ Transform the image generation request to Imagen format Uses predict API with instances and parameters """ + for name in caller_labels_header_names(headers): + del headers[name] # rebind-ok: the handler posts this dict, so a caller-set labels header must go # Default parameters default_params: Final = { "sampleCount": 1, diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index 7b600c4dc57..096d88cb959 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -5,6 +5,7 @@ from typing import Final from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig from litellm.llms.vertex_ai.common_utils import validate_vertex_location, vertex_top_level_request_fields +from litellm.llms.vertex_ai.request_labels import without_caller_labels_header from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.router import GenericLiteLLMParams @@ -94,7 +95,7 @@ class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): return { "Content-Type": "application/json", "Authorization": f"Bearer {access_token}", - **headers, + **without_caller_labels_header(headers), } def get_complete_url( diff --git a/litellm/llms/vertex_ai/request_labels.py b/litellm/llms/vertex_ai/request_labels.py index b3b9b4b2c33..7186f523c92 100644 --- a/litellm/llms/vertex_ai/request_labels.py +++ b/litellm/llms/vertex_ai/request_labels.py @@ -11,9 +11,9 @@ from __future__ import annotations import base64 import json import unicodedata -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from itertools import chain -from typing import Final, TypeAlias +from typing import Final, TypeAlias, TypeVar from pydantic import TypeAdapter @@ -33,6 +33,7 @@ _ANY_MAPPING: Final = TypeAdapter(Mapping[object, object]) _ANY_SEQUENCE: Final = TypeAdapter(tuple[object, ...]) Label: TypeAlias = tuple[str, str] +_HeaderValue: Final = TypeVar("_HeaderValue") def vertex_request_label_fields() -> tuple[str, ...]: @@ -129,6 +130,19 @@ def is_vertex_labels_header(name: str) -> bool: return name.lower() == VERTEX_LABELS_HEADER.lower() +def caller_labels_header_names(header_names: Iterable[str]) -> tuple[str, ...]: + """ + The ``X-Vertex-AI-Labels`` headers a Vertex call that carries ``labels`` in its body drops once + ``litellm.vertex_request_label_fields`` is set, since a caller-set header would bill beside the resolved labels. + """ + return tuple(filter(is_vertex_labels_header, header_names)) if vertex_request_label_fields() else () + + +def without_caller_labels_header(headers: Mapping[str, _HeaderValue]) -> Mapping[str, _HeaderValue]: + dropped: Final = caller_labels_header_names(headers) + return {name: value for name, value in headers.items() if name not in dropped} + + def _decoded_labels_header(encoded: object) -> object: if not isinstance(encoded, str): return None diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index ba0ed47bd5a..809d727e4a5 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -15,6 +15,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.request_labels import without_caller_labels_header from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( @@ -73,7 +74,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, api_key: str | None = None, optional_params: dict | None = None, @@ -105,7 +106,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): default_headers["Authorization"] = headers["Authorization"] # Merge other headers, overriding any default ones except Authorization - return {**default_headers, **headers} + return {**default_headers, **without_caller_labels_header(headers)} def transform_rerank_request( self, diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 5bbeb2bf304..8dba8150c9d 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -24,7 +24,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_base_url, pop_vertex_request_labels, ) -from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields +from litellm.llms.vertex_ai.request_labels import caller_labels_header_names, vertex_request_label_fields from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject @@ -345,7 +345,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): api_base: str, video_create_optional_request_params: dict, litellm_params: GenericLiteLLMParams, - headers: dict, + headers: dict[str, str], ) -> tuple[dict, RequestFiles, str]: """ Transform the video creation request for Veo API. @@ -367,6 +367,8 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): } } """ + for name in caller_labels_header_names(headers): + del headers[name] # rebind-ok: the handler posts this dict, so a caller-set labels header must go # Build instance with prompt instance_dict: Final[dict[str, object]] = {"prompt": prompt} params_copy: Final = video_create_optional_request_params.copy() diff --git a/tests/integration/providers/test_vertex_request_label_identity_wire.py b/tests/integration/providers/test_vertex_request_label_identity_wire.py index 2d4584cdb0d..9f5d374a0a2 100644 --- a/tests/integration/providers/test_vertex_request_label_identity_wire.py +++ b/tests/integration/providers/test_vertex_request_label_identity_wire.py @@ -214,6 +214,42 @@ def test_team_alias_outside_googles_label_rules_is_rewritten_to_a_compliant_labe assert labels == rig.uppercase_team_identity, labels + +_FORGED_LABELS_HEADER: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + + +def _gemini_route(rig: Rig, route: str, marker: str) -> tuple[str, dict[str, JsonValue]]: + caller_headers: Final[dict[str, JsonValue]] = {"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER} + if route == "generate_content": + return f"/v1beta/models/{rig.chat_model}:generateContent", { + "contents": [{"role": "user", "parts": [{"text": marker}]}], + "extra_headers": caller_headers, + } + return "/v1/chat/completions", { + "model": rig.chat_model, + "messages": [{"role": "user", "content": marker}], + "stream": route == "chat_stream", + "extra_headers": caller_headers, + } + + +@pytest.mark.parametrize("route", ["chat", "chat_stream", "generate_content"]) +def test_gemini_drops_a_forged_labels_header_and_bills_the_identity_in_the_body(rig: Rig, route: str) -> None: + marker: Final = f"forged-header-{route}-{uuid.uuid4().hex}" + path, body = _gemini_route(rig, route, marker) + response: Final = rig.proxy.request("POST", path, body, key=rig.key) + + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + request: Final = rig.request_sent(marker) + assert _LABELS_HEADER not in {name.lower() for name in request.headers}, request.headers + assert _JSON_OBJECT.validate_json(request.body).get("labels") == { + **rig.identity, + "cost-center": "cc-1234", + "owner": "platform_team", + } + + def test_embedding_labels_carry_the_caller_identity(rig: Rig) -> None: marker: Final = f"embed-{uuid.uuid4().hex}" response: Final = rig.proxy.request( diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index 584a0b42eef..c492ae70cc0 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -360,3 +360,33 @@ def test_gemini_transcription_carries_identity_labels_only_while_the_setting_is_ assert {key: value for key, value in setting_on.items() if key != "labels"} == setting_off assert setting_on["labels"] == {**_IDENTITY, "cost-center": "cc-1"} assert list(cast(dict[str, str], setting_on["labels"])) == [*_IDENTITY, "cost-center"] + + +class _TokenStubbedConfig(VertexGeminiAudioTranscriptionConfig): + def _ensure_access_token(self, credentials, project_id, custom_llm_provider): + return "fake-token", "resolved-project" + + +def _transcription_headers(caller_headers: dict[str, str]) -> dict[str, str]: + return _TokenStubbedConfig().validate_environment( + headers=caller_headers, + model="gemini-3.5-transcribe-preview", + messages=[], + optional_params={}, + litellm_params={"vertex_project": "resolved-project"}, + ) + + +def test_gemini_transcription_drops_a_caller_labels_header_only_while_the_setting_is_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forged: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + caller_headers: Final = {"X-Vertex-AI-Labels": forged, "x-trace-id": "t-1"} + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _transcription_headers(caller_headers) + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _transcription_headers(caller_headers) + + assert setting_off["X-Vertex-AI-Labels"] == forged + assert "X-Vertex-AI-Labels" not in setting_on + assert setting_on["x-trace-id"] == "t-1" diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 51838c114a1..9919feb89b9 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -648,7 +648,7 @@ _SENT_BODY: Final = TypeAdapter(Mapping[str, object]) _SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) -def _sent_gemini_edit_body() -> Mapping[str, object]: +def _sent_gemini_edit(extra_headers: Mapping[str, str] | None = None) -> httpx.Request: sent: Final[list[httpx.Request]] = [] def reply(request: httpx.Request) -> httpx.Response: @@ -664,9 +664,14 @@ def _sent_gemini_edit_body() -> Mapping[str, object]: api_base="https://vertex.test/edit", metadata=dict(_PROXY_METADATA), labels=dict(_FORGED), + extra_headers=dict(extra_headers or {}), client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(reply))), ) - return _SENT_BODY.validate_json(sent[0].content) + return sent[0] + + +def _sent_gemini_edit_body() -> Mapping[str, object]: + return _SENT_BODY.validate_json(_sent_gemini_edit().content) def test_vertex_gemini_edit_sends_identity_labels_only_while_the_setting_is_on(monkeypatch: pytest.MonkeyPatch): @@ -682,3 +687,20 @@ def test_vertex_gemini_edit_sends_identity_labels_only_while_the_setting_is_on(m ("app", "bot"), ("cost-center", "cc-1"), ] + + +_FORGED_LABELS_HEADER: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + + +def test_vertex_gemini_edit_drops_a_caller_labels_header_only_while_the_setting_is_on( + monkeypatch: pytest.MonkeyPatch, +): + caller_headers: Final = {"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER, "x-trace-id": "t-1"} + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_gemini_edit(caller_headers).headers + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _sent_gemini_edit(caller_headers).headers + + assert setting_off["x-vertex-ai-labels"] == _FORGED_LABELS_HEADER + assert "x-vertex-ai-labels" not in setting_on + assert setting_on["x-trace-id"] == "t-1" diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index dd6d0b1d228..32de7c15567 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -1,5 +1,7 @@ +import base64 import datetime import json +from collections.abc import Mapping from unittest.mock import MagicMock, patch from typing import Final @@ -1149,3 +1151,53 @@ def test_imagen_image_generation_validation_uses_supplied_vertex_token() -> None config._credentials_project_mapping.pop(cache_key, None) assert headers == {"Content-Type": "application/json", "Authorization": "Bearer test-access-token"} + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias") +_FORGED_LABELS_HEADER: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + + +_GEMINI_IMAGE_REPLY: Final = { + "candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}] +} +_IMAGEN_REPLY: Final = {"predictions": [{"bytesBase64Encoded": "aW1n", "mimeType": "image/png"}]} + + +def _sent_generation_headers(model: str, reply_json: Mapping[str, object]) -> httpx.Headers: + sent: Final[list[httpx.Request]] = [] # mutable-ok: MockTransport hands each request to a callback + + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json=reply_json) + + litellm.image_generation( + model=model, + prompt="a lighthouse", + api_base="https://vertex.invalid/generate", + metadata={"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}, + extra_headers={"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER, "x-trace-id": "t-1"}, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(reply))), + ) + return sent[0].headers + + +@pytest.mark.parametrize( + ("model", "reply_json"), + [ + ("vertex_ai/gemini-3.1-flash-image", _GEMINI_IMAGE_REPLY), + ("vertex_ai/imagen-4.0-fast-generate-001", _IMAGEN_REPLY), + ], +) +def test_vertex_image_generation_drops_a_caller_labels_header_only_while_the_setting_is_on( + model: str, + reply_json: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_generation_headers(model, reply_json) + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _sent_generation_headers(model, reply_json) + + assert setting_off["x-vertex-ai-labels"] == _FORGED_LABELS_HEADER + assert "x-vertex-ai-labels" not in setting_on + assert setting_on["x-trace-id"] == "t-1" diff --git a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py index 661f5200601..abe073bd390 100644 --- a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py +++ b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -327,3 +327,17 @@ async def test_create_sends_caller_labels_as_is_while_the_setting_is_off( assert with_labels["labels"] == dict(FORGED) assert "labels" not in without_labels + + +def test_create_drops_a_caller_labels_header_only_while_the_setting_is_on( + config: VertexAIInteractionsConfig, litellm_params: GenericLiteLLMParams, monkeypatch: pytest.MonkeyPatch +): + caller_headers: Final = {"x-vertex-ai-labels": "eyJ1c2VyX2FwaV9rZXlfdGVhbV9hbGlhcyI6ICJmb3JnZWQifQ==", "x-trace-id": "t-1"} + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = config.validate_environment(caller_headers, "gemini-omni-flash-preview", litellm_params) + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + setting_on: Final = config.validate_environment(caller_headers, "gemini-omni-flash-preview", litellm_params) + + assert setting_off["x-vertex-ai-labels"] == caller_headers["x-vertex-ai-labels"] + assert "x-vertex-ai-labels" not in setting_on + assert setting_on["x-trace-id"] == "t-1" diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 630b2e1eb34..e33e07bcffb 100644 --- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -3,13 +3,18 @@ Tests for Vertex AI rerank transformation functionality. Based on the test patterns from other rerank providers and the current Vertex AI implementation. """ +import base64 import json import os +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import responses +import respx +import litellm from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.types.rerank import RerankResponse @@ -646,3 +651,54 @@ class TestVertexAIRerankTransform: second_call = mock_ensure_access_token.call_args_list[1] assert first_call.kwargs["credentials"] == "path/to/credentials.json" assert second_call.kwargs["credentials"] == "path/to/credentials.json" + + +_FORGED_LABELS_HEADER: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + + +def _rank_request_headers(credential: str) -> httpx.Headers: + with ( + patch("litellm.disable_aiohttp_transport", True), + responses.RequestsMock(assert_all_requests_are_fired=False) as token_http, + respx.mock as rank_http, + ): + token_http.post( + "https://oauth2.googleapis.com/token", + json={"access_token": "rank-token", "expires_in": 3600, "token_type": "Bearer"}, + ) + rank_http.post(url__regex=r"https://discoveryengine\.googleapis\.com/.*:rank").respond( + 200, json={"records": [{"id": "0", "score": 0.9}]} + ) + litellm.rerank( + model="vertex_ai/semantic-ranker-default@latest", + query="billing labels", + documents=["labels show cost per team"], + vertex_credentials=credential, + vertex_project="rank-project", + metadata={"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}, + headers={"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER, "x-trace-id": "t-1"}, + caching=False, + num_retries=0, + ) + return rank_http.calls.last.request.headers + + +def test_rerank_drops_a_caller_labels_header_only_while_the_setting_is_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: + credential: Final = json.dumps( + { + "type": "authorized_user", + "client_id": "rank-principal", + "client_secret": "synthetic-client-secret", + "refresh_token": "synthetic-refresh-rank-principal", + } + ) + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _rank_request_headers(credential) + monkeypatch.setattr(litellm, "vertex_request_label_fields", ("user_api_key_alias", "user_api_key_team_alias")) + setting_on: Final = _rank_request_headers(credential) + + assert setting_off["x-vertex-ai-labels"] == _FORGED_LABELS_HEADER + assert "x-vertex-ai-labels" not in setting_on + assert setting_on["x-trace-id"] == "t-1" diff --git a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py index 3089fbd1537..99672d193ec 100644 --- a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -1007,3 +1007,29 @@ def test_veo_request_carries_identity_labels_only_while_the_setting_is_on(monkey "labels": {**_IDENTITY, "app": "bot", "cost-center": "cc-1"}, } assert list(cast(Mapping[str, str], setting_on["labels"])) == [*_IDENTITY, "app", "cost-center"] + + +_FORGED_LABELS_HEADER: Final = base64.b64encode(json.dumps({"user_api_key_team_alias": "forged"}).encode()).decode() + + +def _veo_request_headers() -> Mapping[str, str]: + headers: Final = {"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER, "x-trace-id": "t-1"} + VertexAIVideoConfig().transform_video_create_request( + model="veo-3.1-generate-001", + prompt="a lighthouse at dusk", + api_base="https://vertex.test/models/veo-3.1-generate-001", + video_create_optional_request_params={}, + litellm_params=GenericLiteLLMParams(metadata=dict(_PROXY_METADATA)), + headers=headers, + ) + return headers + + +def test_veo_request_drops_a_caller_labels_header_only_while_the_setting_is_on(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _veo_request_headers() + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _veo_request_headers() + + assert setting_off == {"X-Vertex-AI-Labels": _FORGED_LABELS_HEADER, "x-trace-id": "t-1"} + assert setting_on == {"x-trace-id": "t-1"} From 9594841e3736aacffc71ef83eac8d73d0111fe7a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:59:28 -0800 Subject: [PATCH 12/13] test(vertex): script the v1beta1 GenAI target in the labels header integration test (cherry picked from commit 14f56f569f6208f392a05752b5b4901e4a886aab) --- .../providers/test_vertex_request_label_identity_wire.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration/providers/test_vertex_request_label_identity_wire.py b/tests/integration/providers/test_vertex_request_label_identity_wire.py index 9f5d374a0a2..62f09d6adf8 100644 --- a/tests/integration/providers/test_vertex_request_label_identity_wire.py +++ b/tests/integration/providers/test_vertex_request_label_identity_wire.py @@ -23,6 +23,7 @@ _EMBED_BACKEND: Final = "text-embedding-005" _CLAUDE_BACKEND: Final = "claude-haiku-5-5" _MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models" _CHAT_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:generateContent" +_GENAI_TARGET: Final = _CHAT_TARGET.replace("/v1/", "/v1beta1/", 1) _STREAM_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:streamGenerateContent?alt=sse" _EMBED_TARGET: Final = f"{_MODELS_PATH}/{_EMBED_BACKEND}:predict" _CLAUDE_TARGET: Final = ( @@ -58,7 +59,7 @@ def _vertex_peer(request: Request) -> Reply: ) if request.target == _STREAM_TARGET: return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_GEMINI_REPLY)}\n\n".encode(),)) - if request.target == _CHAT_TARGET: + if request.target in (_CHAT_TARGET, _GENAI_TARGET): return Reply(body=json.dumps(_GEMINI_REPLY).encode()) if request.target == _EMBED_TARGET: return Reply(body=json.dumps(_EMBED_REPLY).encode()) From b01d7ef18d6f1fabd2be0aea8e1b65ab990fd41b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 16:38:53 -0800 Subject: [PATCH 13/13] test(vertex): type the transcription test's access token override (cherry picked from commit 63e6b48d7a003296f7c2cd9a5b43aa728a640707) --- .../test_vertex_ai_gemini_transcribe_transformation.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index c492ae70cc0..65b3b43232c 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -2,7 +2,7 @@ import base64 import json import os from types import MappingProxyType -from typing import Final, cast +from typing import Final, Literal, cast import httpx import pytest @@ -16,6 +16,7 @@ from litellm.llms.vertex_ai.audio_transcription.transformation import ( VertexAIAudioTranscriptionConfig, ) from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.utils import LlmProviders, TranscriptionUsageTokensObject from litellm.utils import ProviderConfigManager, get_optional_params_transcription @@ -363,7 +364,12 @@ def test_gemini_transcription_carries_identity_labels_only_while_the_setting_is_ class _TokenStubbedConfig(VertexGeminiAudioTranscriptionConfig): - def _ensure_access_token(self, credentials, project_id, custom_llm_provider): + def _ensure_access_token( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: return "fake-token", "resolved-project"