scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(vertex_ai): label Vertex requests with the caller's key, team, org, and project (#50732)

Browse source 
* 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.

* 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

* fix(vertex_ai): keep identity labels when a Gemini caller sends extra_body labels

* fix(vertex_ai): label Interactions, image edits, Mistral, and the /vertex_ai pass-through, and gate Gemini image labels on the setting

* fix(vertex_ai): move body labels into the labels header on Vertex rawPredict pass-through calls

* 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

* 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

* fix(vertex): label Veo video, Gemini transcription, and pass-through predictLongRunning requests

* 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

* fix(vertex): drop a caller labels header on body-labelled Vertex routes while identity labels are on

* test(vertex): script the v1beta1 GenAI target in the labels header integration test

* test(vertex): type the transcription test's access token override

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
zoveki 2026-11-13 21:39:01 -08:00 • committed by GitHub
parent c6aa2a3b2a
commit 8973e2dc4f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
43 changed files with 2649 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -214,6 +214,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

1
litellm/batches/main.py
View file

@ -365,6 +365,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(

10
litellm/interactions/http_handler.py
View file

@ -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(

10
litellm/llms/base_llm/google_genai/transformation.py
View file

@ -75,6 +75,16 @@ 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
def transform_caller_headers(self, headers: Mapping[str, object]) -> Mapping[str, object]:
return headers
@abstractmethod
def map_generate_content_optional_params(
self,

7
litellm/llms/base_llm/interactions/transformation.py
View file

@ -120,6 +120,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,

18
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -12570,7 +12570,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(
@ -12580,9 +12580,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)
log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers)
log_request(data)
@ -12674,7 +12677,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(
@ -12684,9 +12687,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)
log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers)
log_request(data)

8
litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py
View file

@ -22,7 +22,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, without_caller_labels_header
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.openai import (
AllMessageValues,
@ -110,7 +111,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",
@ -172,7 +173,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,

17
litellm/llms/vertex_ai/batches/handler.py
View file

@ -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
@ -22,6 +22,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 (
@ -83,6 +88,14 @@ def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str
return response.json()
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))
def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatusError:
body: Final = mask_sensitive_info(error.response.text)
return MaskedHTTPStatusError(
@ -108,6 +121,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(
@ -136,6 +150,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(

6
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -212,6 +212,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
@ -230,12 +231,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(

81
litellm/llms/vertex_ai/common_utils.py
View file

@ -24,12 +24,20 @@ 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,
string_labels,
vertex_metadata_sources,
vertex_request_label_fields,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import (
VERTEX_AI_PROVIDER_METADATA_FIELDS,
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
@ -200,8 +208,12 @@ 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])
_CALLER_LABELS: Final = TypeAdapter(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``:
@ -209,23 +221,23 @@ 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(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(
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 = (
_CALLER_LABELS.validate_python(optional_params.pop("labels", None)) if optional_params is not None else None
)
return dict(string_labels(raw)) or vertex_request_labels_from_litellm_params(litellm_params)
def pop_vertex_request_labels(
@ -233,18 +245,29 @@ def pop_vertex_request_labels(
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``). 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)
return labels if labels else None
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
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):

11
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -35,7 +35,8 @@ from litellm.litellm_core_utils.prompt_templates.file_id_utils 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,
get_file_type_from_extension,
@ -1175,14 +1176,14 @@ 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)
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
drops_top_k: Final = "top_k" in optional_params and not VertexGeminiConfig.supports_sampling_params(model)
if drops_top_k:
request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature
VertexGeminiConfig.reject_unsupported_sampling_param(
model=model,
param="top_k",
@ -1241,6 +1242,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

17
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -115,6 +115,7 @@ from ..context_caching.rejected_cache_retry import (
post_with_fresh_cache_async,
post_with_fresh_cache_sync,
)
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 (
@ -274,6 +275,7 @@ _GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType(
_SPEECH_CONFIG_ITEMS: Final = TypeAdapter(list[object])
_CALLER_HEADERS: Final = TypeAdapter(dict[str, str])
def _normalize_gemini_speech_config_item(value: Mapping[str, object]) -> Mapping[str, object]:
@ -3888,6 +3890,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,
@ -3896,7 +3903,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,
@ -3927,7 +3934,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(
@ -3950,7 +3957,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(
@ -3976,7 +3983,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,
@ -4001,7 +4008,7 @@ class VertexLLM(VertexBase):
log_request(data)
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 ##

13
litellm/llms/vertex_ai/google_genai/transformation.py
View file

@ -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 vertex_top_level_request_fields
from litellm.llms.vertex_ai.request_labels import without_caller_labels_header
from litellm.types.router import GenericLiteLLMParams
@ -38,6 +41,16 @@ 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:
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

20
litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py
View file

@ -12,8 +12,13 @@ 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, map_vertex_image_size
from litellm.llms.vertex_ai.common_utils import (
get_vertex_base_url,
map_vertex_image_size,
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 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 (
@ -168,6 +173,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 litellm.BadRequestError(
@ -194,7 +201,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, [])

22
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
View file

@ -15,8 +15,9 @@ 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, 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.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 +215,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}]}]
@ -254,12 +257,15 @@ 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)
if vertex_request_label_fields()
else None,
}
return {**request_body, **{key: value for key, value in optional_fields.items() if value}}
def transform_image_generation_response(
self,

5
litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py
View file

@ -15,6 +15,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,
@ -177,13 +178,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,

12
litellm/llms/vertex_ai/interactions/transformation.py
View file

@ -4,7 +4,8 @@ from typing import Final
from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base
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.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(
@ -111,6 +112,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,

5
litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py
View file

@ -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(

14
litellm/llms/vertex_ai/multimodal_embeddings/transformation.py
View file

@ -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, extra_body_fields
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=[])
@ -195,7 +203,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig):
elif extra_body:
request_data["parameters"] = dict(extra_body)
return cast(dict, request_data)
return {**request_data, "labels": labels} if labels else {**request_data}
def transform_embedding_response(
self,

176
litellm/llms/vertex_ai/request_labels.py Normal file
View file

@ -0,0 +1,176 @@
"""
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 Iterable, Mapping
from itertools import chain
from typing import Final, TypeAlias, TypeVar
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]
_HeaderValue: Final = TypeVar("_HeaderValue")
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))
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 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 string_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(
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)
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 = 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 (*string_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 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
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 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 decoded_header is None else decoded_header,
litellm_params,
VERTEX_PARTNER_REQUEST_LABELS_MAX,
)
return base64.b64encode(json.dumps(dict(resolved)).encode()).decode() if resolved else None

11
litellm/llms/vertex_ai/rerank/transformation.py
View file

@ -15,9 +15,8 @@ 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.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 (
@ -87,7 +86,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,
@ -119,7 +118,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,
@ -163,7 +162,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

26
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py
View file

@ -14,6 +14,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,
@ -64,6 +70,18 @@ def _without_version_header(headers: Mapping[str, str]) -> dict[str, str]: # mu
return {name: value for name, value in headers.items() if name.lower() != _VERSION_HEADER_VERTEX_READS_FROM_BODY}
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:
@ -113,6 +131,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
project_id,
vertex_ai_project,
vertex_ai_location,
litellm_params,
)
async def avalidate_anthropic_messages_environment(
@ -141,6 +160,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
project_id,
vertex_ai_project,
vertex_ai_location,
litellm_params,
)
def _finalize_vertex_messages_environment(
@ -154,6 +174,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)
@ -217,6 +238,11 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
if all_beta_values:
request_headers["anthropic-beta"] = ",".join(sorted(all_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(

29
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py
View file

@ -13,6 +13,12 @@ from litellm.anthropic_beta_headers_manager import get_header_only_betas
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
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.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
@ -25,12 +31,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
@ -148,11 +153,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,

41
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -6,10 +6,16 @@ from types import MappingProxyType
from typing import Final
import httpx
from pydantic import TypeAdapter
import litellm
from litellm import LlmProviders
from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls
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
@ -17,6 +23,16 @@ 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])
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):
@ -159,6 +175,27 @@ class VertexAIPartnerModels(VertexBase):
model=model,
)
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,
_caller_labels(params),
_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
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:
model = model.split("@")[0]
@ -180,6 +217,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.
@ -250,6 +288,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:

11
litellm/llms/vertex_ai/videos/transformation.py
View file

@ -25,7 +25,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 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
@ -350,7 +352,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.
@ -372,9 +374,14 @@ 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()
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):
@ -419,6 +426,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"

92
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -83,6 +83,14 @@ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError
from litellm.llms.openai.common_utils import is_openai_agents_api_path, openai_api_resource
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.llms.vllm.common_utils import endpoint_under_api_base
from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
@ -148,6 +156,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
STORED_PROMPT_PROBE_MODEL,
HttpPassThroughEndpointHelpers,
UpstreamRequest,
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
@ -3008,6 +3017,86 @@ 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", "predictLongRunning", "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_publisher(segments: tuple[str, ...]) -> str | None:
match segments[-4:]:
case ("publishers", publisher, "models", _):
return publisher
case _:
return None
def _vertex_label_placement(http_method: str, endpoint: str) -> VertexLabelPlacement | None:
if http_method != "POST":
return None
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, ...],
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"}
labelled_body: Final = {**unlabelled, "labels": dict(resolved)} if resolved else unlabelled
return _without_vertex_labels_header(headers), labelled_body
def _with_vertex_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 = _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
return {**unlabelled_headers, VERTEX_LABELS_HEADER: labels_header}, unlabelled_body
def _with_vertex_request_labels(
placement: VertexLabelPlacement | 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
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(
call_type: Literal["discovery", "aiplatform"],
) -> BaseVertexAIPassThroughHandler:
@ -3447,6 +3536,9 @@ async def _base_vertex_proxy_route(
upstream_headers_filter=partial(
_without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers)
),
upstream_request_filter=partial(
_with_vertex_request_labels, _vertex_label_placement(request.method, updated_url.path)
),
) # dynamically construct pass-through endpoint based on incoming path
try:

46
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -1283,6 +1283,28 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
return headers
UpstreamRequest: TypeAlias = tuple[Mapping[str, str], Mapping[str, object] | None]
UpstreamRequestFilter: TypeAlias = Callable[
[Mapping[str, str], Mapping[str, object] | None, Mapping[str, object]], UpstreamRequest
]
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
def _caller_query_source(
keep_target_query: bool, query_params: Mapping[str, object] | None, request: Request
) -> Mapping[str, object]:
@ -1411,6 +1433,7 @@ async def pass_through_request(
guardrails_config: dict | None = None,
timeout: float | None = None,
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
upstream_request_filter: UpstreamRequestFilter = _unfiltered_request,
operation: GenAIOperation | None = None,
keep_target_query: bool = False,
):
@ -1434,6 +1457,8 @@ async def pass_through_request(
timeout: Optional per-endpoint timeout in seconds. Falls back to
general_settings.pass_through_request_timeout, then 600s.
upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion
upstream_request_filter: Optional last rewrite of the headers and JSON body sent upstream, given the
litellm_params the request is logged with
operation: Optional GenAI operation the route reports, unless the upstream names its own
keep_target_query: Send the target URL's query exactly as given, adding none of the caller's query params
"""
@ -1795,14 +1820,19 @@ async def pass_through_request(
logging_url = str(url) + "?" + requested_query_params_str
outgoing_body: Final = cast("Mapping[str, object] | None", _parsed_body) # cast-ok: json
upstream_body: Final = _with_anthropic_user_id(outgoing_body, client_user_id)
outbound_headers, upstream_body = _filtered_upstream_request(
upstream_request_filter,
upstream_headers,
_with_anthropic_user_id(outgoing_body, client_user_id),
_METADATA_ADAPTER.validate_python(kwargs["litellm_params"]),
)
logging_obj.pre_call(
input=[{"role": "user", "content": safe_dumps(_parsed_body)}],
api_key="",
additional_args={
"complete_input_dict": upstream_body,
"api_base": str(logging_url),
"headers": redact_credential_headers(upstream_headers),
"headers": redact_credential_headers(outbound_headers),
},
)
stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body(
@ -1819,7 +1849,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,
)
@ -1831,7 +1861,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
@ -1839,7 +1869,7 @@ async def pass_through_request(
request.method,
url,
params=requested_query_params,
headers=upstream_headers,
headers=outbound_headers,
json=upstream_body,
)
)
@ -1927,7 +1957,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,
)
@ -1937,7 +1967,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,
_parsed_body=upstream_body,
forward_multipart=is_multipart,
@ -2482,6 +2512,7 @@ def create_pass_through_route(
config_file_path: str | None = None,
timeout: float | None = None,
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
upstream_request_filter: UpstreamRequestFilter = _unfiltered_request,
operation: GenAIOperation | None = None,
keep_target_query: bool = False,
):
@ -2632,6 +2663,7 @@ def create_pass_through_route(
guardrails_config=cast(dict | None, param_guardrails),
timeout=cast(float | None, param_timeout),
upstream_headers_filter=upstream_headers_filter,
upstream_request_filter=upstream_request_filter,
operation=operation,
keep_target_query=keep_target_query,
)

5
litellm/types/llms/vertex_ai.py
View file

@ -1,8 +1,9 @@
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from enum import Enum
from typing import Final, Literal, Protocol
from typing_extensions import (
NotRequired,
ReadOnly,
Required,
TypedDict,
@ -591,6 +592,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):
@ -738,6 +740,7 @@ class VertexAIBatchPredictionJob(TypedDict):
model: str
inputConfig: InputConfig
outputConfig: OutputConfig
labels: NotRequired[ReadOnly[Mapping[str, str]]]
class VertexBatchPredictionResponse(TypedDict, total=False):

298
tests/integration/providers/test_vertex_request_label_identity_wire.py Normal file
View file

@ -0,0 +1,298 @@
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"
_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 = (
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 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())
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
_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(
"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"}

99
tests/unit/batches/test_main.py
View file

@ -1204,6 +1204,105 @@ async def test_delete_batch_output_file(respx_mock: respx.MockRouter, monkeypatc
assert delete_response.deleted is True
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
VLLM_ROUTER_API_BASE: Final = "http://vllm-router.test/v1"
VLLM_ROUTER_BATCH: Final = MappingProxyType(
{

178
tests/unit/google_genai/test_google_genai_main.py
View file

@ -4,15 +4,18 @@ Test to verify the Google GenAI generate_content adapter functionality
"""
import json
from collections.abc import Iterator
from typing import Final
from collections.abc import Iterator, Mapping
from types import MappingProxyType
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
@ -382,3 +385,168 @@ async def test_vertex_generate_content_spend_follows_the_deployment_location(
recorded: Final = await _logged_vertex_generate_content_cost(location, stream=stream)
assert recorded == pytest.approx(expected), (location, stream)
_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"})
_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"}],
"usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3},
}
_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(
config: GoogleGenAIConfig,
litellm_params: GenericLiteLLMParams,
extra_body: Mapping[str, object],
*,
is_async: bool,
stream: bool,
) -> Mapping[str, object]:
import datetime
from unittest.mock import MagicMock
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
reply: Final = MagicMock(side_effect=lambda _request: httpx.Response(200, json=_GEMINI_REPLY))
transport: Final = httpx.MockTransport(reply)
client: Final = (
AsyncHTTPHandler(transport=transport) if is_async else HTTPHandler(client=httpx.Client(transport=transport))
)
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,
stream=stream,
)
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(
extra_body: Mapping[str, object], *, is_async: bool, stream: bool
) -> Mapping[str, object]:
return await _generate_content_body(
_SignedInVertexGoogleGenAIConfig(),
GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), vertex_location="us-central1"),
extra_body,
is_async=is_async,
stream=stream,
)
@pytest.mark.asyncio
@_HANDLER_CALLS
async def test_vertex_generate_content_sends_identity_labels_a_caller_cannot_forge(
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({"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
@_HANDLER_CALLS
async def test_vertex_generate_content_without_caller_labels_still_carries_identity(
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({}, 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
@_HANDLER_CALLS
@pytest.mark.parametrize(
"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(
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(extra_body, is_async=is_async, stream=stream)
assert body.get("labels", "absent") == expected_labels
@pytest.mark.asyncio
@_HANDLER_CALLS
async def test_ai_studio_generate_content_keeps_the_callers_labels_with_the_setting_on(
is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS)
body: Final = await _generate_content_body(
GoogleGenAIConfig(),
GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), api_key="test-key"),
{"labels": dict(_FORGED)},
is_async=is_async,
stream=stream,
)
assert body["labels"] == _FORGED

72
tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py
View file

@ -1,6 +1,8 @@
import base64
import json
import os
from types import MappingProxyType
from typing import Final, Literal, cast
import httpx
import pytest
@ -14,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
@ -324,3 +327,72 @@ 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"]
class _TokenStubbedConfig(VertexGeminiAudioTranscriptionConfig):
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"
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"

41
tests/unit/llms/vertex_ai/batches/test_handler.py
View file

@ -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",
[

68
tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
View file

@ -1,9 +1,11 @@
import base64
import json
import os
from collections.abc import Mapping
from datetime import datetime
from io import BytesIO, StringIO
from pathlib import Path
from types import MappingProxyType
from typing import Dict, Final
from unittest.mock import MagicMock, patch
@ -758,3 +760,69 @@ 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(extra_headers: Mapping[str, str] | None = None) -> httpx.Request:
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),
extra_headers=dict(extra_headers or {}),
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(reply))),
)
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):
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"),
]
_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"

52
tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
View file

@ -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
@ -1240,3 +1242,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"

112
tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
View file

@ -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,104 @@ 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
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"

76
tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py
View file

@ -1,4 +1,7 @@
import json
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import google.auth
@ -6,10 +9,11 @@ import google.auth.credentials
import httpx
import pytest
import respx
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,
)
@ -206,3 +210,71 @@ def test_extra_body_fields_reach_the_multimodal_predict_body(
"instances": [{"text": "Hello world"}],
"parameters": {"dimension": 128},
}
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, 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]}]})
)
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,
}
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)
@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, 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, request_labels)
assert body["instances"] == TODAYS_INSTANCES
assert list(SENT_LABELS.validate_python(body["labels"]).items()) == list(expected_labels)

53
tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
View file

@ -3,6 +3,7 @@ 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 datetime import datetime, timedelta
@ -732,3 +733,55 @@ async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens(
await rank(credentials[0], "first-project", "initial", 3)
clock.return_value = now + timedelta(hours=2)
await rank(credentials[0], "first-project", "renewed", 4)
_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
@pytest.mark.usefixtures("clear_vertex_rerank_cache")
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"

405
tests/unit/llms/vertex_ai/test_request_labels.py Normal file
View file

@ -0,0 +1,405 @@
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.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"}
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 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",
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]
REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_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]:
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("setting", [None, []])
@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 surface(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(
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),
(multimodal_embedding_labels, 64),
(google_genai_labels, 64),
],
)
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
@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
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)
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"})}
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
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
):
"""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",
}

138
tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py
View file

@ -1,6 +1,16 @@
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
from pydantic import TypeAdapter
import litellm
import litellm.main
from litellm.llms.custom_httpx.http_handler import HTTPHandler
@pytest.fixture
@ -30,3 +40,131 @@ 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)))
litellm.text_completion(
model="vertex_ai/codestral-2",
prompt="def add(",
roles={"user": {"pre_message": "", "post_message": ""}},
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:
sent: Final[list[httpx.Request]] = []
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]
@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",
)
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,
}

63
tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py
View file

@ -5,6 +5,7 @@ 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 Final, cast
from unittest.mock import Mock, patch
@ -985,3 +986,65 @@ def test_first_class_video_fields_are_rejected_or_dropped(field: str) -> None:
with pytest.raises(UnsupportedParamsError, match=field):
config.map_openai_params({field: 1}, "video-model", False)
assert config.map_openai_params({field: 1}, "video-model", True) == config.map_openai_params({}, "video-model", True)
_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"]
_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"}

3
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -344,6 +344,7 @@ class TestVertexAIPassThroughHandler:
custom_headers={"Authorization": f"Bearer {test_token}"},
is_streaming_request=False,
upstream_headers_filter=mock.ANY,
upstream_request_filter=mock.ANY,
)
@pytest.mark.asyncio
@ -436,6 +437,7 @@ class TestVertexAIPassThroughHandler:
custom_headers={"Authorization": f"Bearer {test_token}"},
is_streaming_request=False,
upstream_headers_filter=mock.ANY,
upstream_request_filter=mock.ANY,
)
@pytest.mark.parametrize(
@ -528,6 +530,7 @@ class TestVertexAIPassThroughHandler:
custom_headers={"Authorization": f"Bearer {default_credentials}"},
is_streaming_request=False,
upstream_headers_filter=mock.ANY,
upstream_request_filter=mock.ANY,
)
@pytest.mark.asyncio

372
tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py
View file

@ -1,5 +1,6 @@
import base64
import json
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from typing import Final
import httpx
@ -525,3 +526,372 @@ def test_vertex_pass_through_forwards_a_request_without_betas_despite_a_malforme
assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY)
assert "anthropic-beta" not in upstream.calls.last.request.headers
_VERTEX_GEMINI_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",
f"{_VERTEX_GEMINI_MODEL_PATH}:predictLongRunning",
"v1/projects/deployment-project/locations/global/batchPredictionJobs",
],
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
) -> 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",
),
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(
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_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"
) -> 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]
@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,
}