diff --git a/litellm/__init__.py b/litellm/__init__.py index dbe31c4a104..acae4fc70e6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 212205a9d7b..e7ae005d0bb 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -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( diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 09aa656e49b..5cb1931259d 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -161,8 +161,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( @@ -256,8 +257,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 0c7054326b7..03e5b8f6769 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -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, diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index 39ca1d8f7d7..b3b73c24aac 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -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, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c4a9c3fec2..aacc754f0c0 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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) diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index 09a09ef96dc..0d9b172e116 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 58f589fea61..f1477e122cb 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Coroutine, Sequence +from collections.abc import Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import urlparse @@ -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( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a1843924c43..f7fb5290993 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -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( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 48bf955f707..0013e678931 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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): diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 6e801ed5756..765ccff23f5 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -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 diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 69e2f7758bd..e462e6a1488 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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 ## diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 563a87b862d..b5ea0ca7c65 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -2,9 +2,12 @@ Transformation for Calling Google models in their native format. """ +from collections.abc import Mapping from typing import Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.common_utils import 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 diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 8e9bef97085..30ab535828c 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -12,8 +12,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, []) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 252054d7330..8a6ff5e61bd 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index daf87b1ec58..4362c721d61 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index d81818d021d..dc2dba2661f 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py index 8af05b3ef32..18dd942a4b3 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py @@ -1,7 +1,9 @@ import json +from collections.abc import Mapping from typing import Final, Literal import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -19,6 +21,7 @@ from litellm.types.utils import EmbeddingResponse from .transformation import VertexAIMultimodalEmbeddingConfig vertex_multimodal_embedding_handler: Final = VertexAIMultimodalEmbeddingConfig() +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) class VertexMultimodalEmbedding(VertexLLM): @@ -84,7 +87,7 @@ class VertexMultimodalEmbedding(VertexLLM): sync_handler = client request_data: Final = vertex_multimodal_embedding_handler.transform_embedding_request( - model, input, optional_params, headers + model, input, optional_params, headers, _LITELLM_PARAMS.validate_python(litellm_params) ) headers = vertex_multimodal_embedding_handler.validate_environment( diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index 7c276b43dc7..d2f8603d4cc 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final, cast from httpx import Headers, Response @@ -22,7 +23,8 @@ from litellm.types.utils import ( from litellm.utils import count_characters, is_base64_encoded from ...base_llm.embedding.transformation import BaseEmbeddingConfig, 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, diff --git a/litellm/llms/vertex_ai/request_labels.py b/litellm/llms/vertex_ai/request_labels.py new file mode 100644 index 00000000000..7186f523c92 --- /dev/null +++ b/litellm/llms/vertex_ai/request_labels.py @@ -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 diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 398471a7c9b..95d91f4774a 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -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 diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index 816301c6637..03976e54a56 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -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( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 3d1690fca33..a6b1a2b844a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index e905b97d29f..c67a5b72b53 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -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: diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 2b88861b79b..e7ca3cafa7a 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -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" diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 782b5698f3c..f7ecda99d9d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a6f2fe1c309..6e21ab904a5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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, ) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 72a17714181..968fb11e44a 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum from typing import 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): diff --git a/tests/integration/providers/test_vertex_request_label_identity_wire.py b/tests/integration/providers/test_vertex_request_label_identity_wire.py new file mode 100644 index 00000000000..62f09d6adf8 --- /dev/null +++ b/tests/integration/providers/test_vertex_request_label_identity_wire.py @@ -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"} diff --git a/tests/unit/batches/test_main.py b/tests/unit/batches/test_main.py index b12ea53bce3..4c869250f88 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -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( { diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 5cac61e727e..ae350c98075 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -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 diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index 532e013c3d7..65b3b43232c 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -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" diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 2c246e065a8..97d4a88f16d 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio import json +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -262,6 +263,46 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +@pytest.mark.parametrize( + "create_data, model", + [ + (CREATE_DATA, "publishers/google/models/gemini-1.5-flash-001"), + (ENDPOINT_CREATE_DATA, TUNED_MODEL_RESOURCE), + ], + ids=["publisher-model", "fine-tuned-endpoint"], +) +def test_create_batch_labels_the_job_from_litellm_metadata_never_metadata( + create_data: Mapping[str, str], model: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", ["user_api_key_alias", "user_api_key_team_alias"]) + h: Final = _make_handler() + client: Final = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}.get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()), + ): + h.create_batch( + _is_async=False, + create_batch_data=create_data, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + litellm_params={ + "metadata": {"user_api_key_alias": "forged"}, + "litellm_metadata": {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}, + }, + ) + + sent: Final = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == model + assert sent["labels"] == {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} + + @pytest.mark.parametrize( "api_base, expected", [ diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index e311313c136..5af84bd2dcd 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -1,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" diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index e1eadb89328..99bc2330fd5 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -1,5 +1,7 @@ +import base64 import datetime import json +from collections.abc import Mapping from unittest.mock import MagicMock, patch from typing import Final @@ -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" diff --git a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py index bebc7d8e021..e4a847367d2 100644 --- a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py +++ b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -1,7 +1,18 @@ +import json +from collections.abc import Mapping +from datetime import datetime +from types import MappingProxyType +from typing import Final + +import httpx import pytest +from pydantic import TypeAdapter import litellm +from litellm.interactions.http_handler import InteractionsHTTPHandler from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig, ) @@ -229,3 +240,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" diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index 123d4128679..fd808f4de20 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -1,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) diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 8abd012f9d8..0e6420442ab 100644 --- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -3,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" diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py new file mode 100644 index 00000000000..8a3cf1fe638 --- /dev/null +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -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", + } diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index 15df8e47af3..a6412032dde 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,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, + } diff --git a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py index ecf131f8215..a243f325534 100644 --- a/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -5,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"} diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index e95945cfd76..da7578e71bc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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 diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index 2e196d6a8c2..179a843c347 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -1,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, + }