From ae73db38557c271ad39052e434da96dd4eaaab6b 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] 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. --- 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 +- .../anthropic/messages/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, 791 insertions(+), 22 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 8dfe5adca1c..c409c95a75f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -216,6 +216,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 0fa80cb2cad..95b5a7523fd 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -24,6 +24,11 @@ from litellm.llms.gemini.common_utils import ( GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO, map_openai_size_to_gemini_image_config, ) +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, @@ -196,6 +201,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: @@ -231,7 +239,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: @@ -240,7 +249,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 e1049ef8ece..9e04de1cecb 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -14,7 +14,7 @@ from litellm.llms.gemini.common_utils import ( get_gemini_image_web_search_requests, map_gemini_image_tools_params, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size, 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 @@ -255,12 +255,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_usage(self, usage: dict) -> ImageUsage: input_tokens_details: Final = ImageUsageInputTokensDetails( 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 398471a7c9b..ede6a56ab1f 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -15,9 +15,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict 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 ( @@ -163,7 +161,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/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index 6d8eb228786..440bfba0f08 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -6,6 +6,12 @@ from litellm.llms.anthropic.messages.native.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, @@ -18,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: @@ -67,6 +85,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert project_id, vertex_ai_project, vertex_ai_location, + litellm_params, ) async def avalidate_anthropic_messages_environment( @@ -95,6 +114,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert project_id, vertex_ai_project, vertex_ai_location, + litellm_params, ) def _finalize_vertex_messages_environment( @@ -108,6 +128,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert project_id: str, vertex_ai_project: str | None, vertex_ai_location: str | None, + litellm_params: Mapping[str, object], ) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers # Router shallow-copies litellm_params, so headers may be shared between requests. request_headers: Final = dict(headers) @@ -177,6 +198,11 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if beta_values: request_headers["anthropic-beta"] = ",".join(beta_values) + replaced_headers, labels_header = _labels_header_replacement(request_headers, litellm_params) + for name in replaced_headers: + del request_headers[name] + if labels_header is not None: + request_headers[VERTEX_LABELS_HEADER] = labels_header return request_headers, api_base def get_complete_url( 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", + }