From dd2eb2764a846848fe22df283f4e2168f68b372b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:10:58 -0700 Subject: [PATCH] feat(vertex_ai): send the opt-in identity labels on multimodal embeddings, the Google GenAI route, and batches With litellm_settings.vertex_request_label_fields on, the multimodal embedding predict body, the native generateContent and streamGenerateContent body, and the BatchPredictionJob now carry the key's identity labels the same way chat does Caller labels on GenAI are resolved after the extra_body merge, so neither a top-level labels field nor extra_body.labels can write the reserved prefix. Batch labels come only from the proxy's litellm_metadata, never the caller's OpenAI batch metadata. With the setting off every route sends the same body as before (cherry picked from commit a7254fede6b95e00ac84f5ad46e8711595030434) --- litellm/batches/main.py | 1 + .../base_llm/google_genai/transformation.py | 7 + litellm/llms/custom_httpx/llm_http_handler.py | 14 +- litellm/llms/vertex_ai/batches/handler.py | 17 ++- .../llms/vertex_ai/batches/transformation.py | 6 +- .../vertex_ai/google_genai/transformation.py | 15 ++ .../embedding_handler.py | 5 +- .../multimodal_embeddings/transformation.py | 14 +- litellm/types/llms/vertex_ai.py | 5 +- tests/unit/batches/test_main.py | 99 +++++++++++++ .../google_genai/test_google_genai_main.py | 139 ++++++++++++++++++ .../llms/vertex_ai/batches/test_handler.py | 41 ++++++ ..._ai_multimodal_embedding_transformation.py | 76 +++++++++- .../llms/vertex_ai/test_request_labels.py | 64 +++++++- 14 files changed, 489 insertions(+), 14 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9b5630d81ad..6b722db922d 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -359,6 +359,7 @@ def create_batch( max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, custom_endpoint=optional_params.get("custom_endpoint"), + litellm_params=litellm_params, ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 232d9d8eb42..51a6177cfc0 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -75,6 +75,13 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): """ return ("safetySettings", "toolConfig", "cachedContent", "labels", "continuationToken") + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return request_fields + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6737905192e..ce4109e732c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12126,9 +12126,12 @@ class BaseLLMHTTPHandler: generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) + request_fields: Final = generate_content_provider_config.transform_top_level_request_fields( + extra_body, litellm_params + ) - if extra_body: - data.update(extra_body) + if request_fields: + data.update(request_fields) ## LOGGING logging_obj.pre_call( @@ -12241,9 +12244,12 @@ class BaseLLMHTTPHandler: generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) + request_fields: Final = generate_content_provider_config.transform_top_level_request_fields( + extra_body, litellm_params + ) - if extra_body: - data.update(extra_body) + if request_fields: + data.update(request_fields) ## LOGGING logging_obj.pre_call( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index a039347e824..3689b51d9dc 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Coroutine, Sequence +from collections.abc import Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import urlparse @@ -21,6 +21,11 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_REQUEST_LABELS_MAX, + resolve_vertex_request_labels, + vertex_request_label_fields, +) from litellm.llms.vertex_ai.vertex_llm_base import graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -91,6 +96,14 @@ def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatus ) +def _proxy_identity_labels(litellm_params: Mapping[str, object]) -> Mapping[str, str]: + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return {} + proxy_metadata: Final = {"litellm_metadata": litellm_params.get("litellm_metadata")} + return dict(resolve_vertex_request_labels(allowed_fields, None, proxy_metadata, VERTEX_REQUEST_LABELS_MAX)) + + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): super().__init__(*args, **kwargs) @@ -107,6 +120,7 @@ class VertexAIBatchPrediction(VertexLLM): timeout: float | httpx.Timeout, max_retries: int | None, custom_endpoint: bool | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: if custom_endpoint: raise VertexAIError( @@ -135,6 +149,7 @@ class VertexAIBatchPrediction(VertexLLM): request=create_batch_data, vertex_project=vertex_project or project_id, vertex_location=vertex_location or "us-central1", + labels=_proxy_identity_labels(litellm_params or {}), ) ) vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 2460278294c..66f12d5b636 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -205,6 +205,7 @@ class VertexAIBatchTransformation: request: CreateBatchRequest, vertex_project: str | None = None, vertex_location: str | None = None, + labels: Mapping[str, str] | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -223,12 +224,15 @@ class VertexAIBatchTransformation: predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), ) - return VertexAIBatchPredictionJob( + job: Final = VertexAIBatchPredictionJob( inputConfig=input_config, outputConfig=output_config, model=model, displayName=request_display_name, ) + if not labels: + return job + return {**job, "labels": dict(labels)} @classmethod def transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 563a87b862d..78fe644a872 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -2,9 +2,12 @@ Transformation for Calling Google models in their native format. """ +from collections.abc import Mapping from typing import Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.types.router import GenericLiteLLMParams @@ -38,6 +41,18 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): return default_headers + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + if not vertex_request_label_fields(): + return request_fields + fields: Final = request_fields or {} + unlabelled: Final = {name: value for name, value in fields.items() if name != "labels"} + labels: Final = pop_vertex_request_labels({"labels": fields.get("labels")}, dict(litellm_params)) + return {**unlabelled, "labels": labels} if labels else unlabelled + def _camel_to_snake(self, camel_str: str) -> str: """Convert camelCase to snake_case""" import re diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py index 8af05b3ef32..18dd942a4b3 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py @@ -1,7 +1,9 @@ import json +from collections.abc import Mapping from typing import Final, Literal import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -19,6 +21,7 @@ from litellm.types.utils import EmbeddingResponse from .transformation import VertexAIMultimodalEmbeddingConfig vertex_multimodal_embedding_handler: Final = VertexAIMultimodalEmbeddingConfig() +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) class VertexMultimodalEmbedding(VertexLLM): @@ -84,7 +87,7 @@ class VertexMultimodalEmbedding(VertexLLM): sync_handler = client request_data: Final = vertex_multimodal_embedding_handler.transform_embedding_request( - model, input, optional_params, headers + model, input, optional_params, headers, _LITELLM_PARAMS.validate_python(litellm_params) ) headers = vertex_multimodal_embedding_handler.validate_environment( diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index da77fa338db..61028ee982d 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final, cast from httpx import Headers, Response @@ -22,7 +23,8 @@ from litellm.types.utils import ( from litellm.utils import count_characters, is_base64_encoded from ...base_llm.embedding.transformation import BaseEmbeddingConfig -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, pop_vertex_request_labels +from ..request_labels import vertex_request_label_fields class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): @@ -169,8 +171,14 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): input: AllEmbeddingInputValues, optional_params: dict, headers: dict, - ) -> dict: + litellm_params: Mapping[str, object] | None = None, + ) -> dict[str, object]: optional_params = optional_params or {} + labels: Final = ( + pop_vertex_request_labels(optional_params, dict(litellm_params or {})) + if vertex_request_label_fields() + else None + ) request_data: Final = VertexMultimodalEmbeddingRequest(instances=[]) @@ -192,7 +200,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if "outputDimensionality" in optional_params: request_data["parameters"] = {"dimension": optional_params["outputDimensionality"]} - return cast(dict, request_data) + return {**request_data, "labels": labels} if labels else {**request_data} def transform_embedding_response( self, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 0480d6b1f3e..691c2c37e0e 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + NotRequired, ReadOnly, Required, TypedDict, @@ -581,6 +582,7 @@ class Instance(TypedDict, total=False): class VertexMultimodalEmbeddingRequest(TypedDict, total=False): instances: Required[list[Instance]] parameters: dict + labels: ReadOnly[Mapping[str, str]] class VideoEmbedding(TypedDict): @@ -728,6 +730,7 @@ class VertexAIBatchPredictionJob(TypedDict): model: str inputConfig: InputConfig outputConfig: OutputConfig + labels: NotRequired[ReadOnly[Mapping[str, str]]] class VertexBatchPredictionResponse(TypedDict, total=False): diff --git a/tests/unit/batches/test_main.py b/tests/unit/batches/test_main.py index 7106f8c7bc2..8bb434e1442 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -1229,3 +1229,102 @@ def test_retrieve__gemini_routes_to_base_http_handler_with_gemini_config(seams): forwarded = seams.base_http.retrieve_batch.call_args.kwargs assert type(forwarded["provider_config"]).__name__ == "GeminiBatchesConfig" assert forwarded["batch_id"] == "la6" + + +VERTEX_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +VERTEX_IDENTITY: Final = MappingProxyType( + {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} +) +VERTEX_PROXY_METADATA: Final = MappingProxyType({**VERTEX_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +VERTEX_BATCH_FILE: Final = "gs://bucket/publishers/google/models/gemini-3.8-flash/file-uuid" +VERTEX_SENT_JOB: Final = TypeAdapter(Mapping[str, object]) +VERTEX_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) +VERTEX_BATCH_JOBS_URL: Final = ( + "https://us-central1-aiplatform.googleapis.com/v1/projects/proj/locations/us-central1/batchPredictionJobs" +) + + +async def _vertex_batch_job_sent( + monkeypatch: pytest.MonkeyPatch, is_async: bool, batch_metadata: Mapping[str, object] +) -> Mapping[str, object]: + import litellm.llms.vertex_ai.batches.handler as vertex_batches + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + reply: Final = MagicMock( + side_effect=lambda _request: httpx.Response( + 200, + json={ + "name": "projects/proj/locations/us-central1/batchPredictionJobs/123", + "state": "JOB_STATE_PENDING", + "createTime": "2026-10-29T00:00:00Z", + }, + ) + ) + transport: Final = httpx.MockTransport(reply) + handler: Final = vertex_batches.VertexAIBatchPrediction(gcs_bucket_name="bucket") + monkeypatch.setattr(handler, "_ensure_access_token", lambda **_: ("token", "proj")) + monkeypatch.setattr(bm, "vertex_ai_batches_instance", handler) + monkeypatch.setattr( + vertex_batches, "get_httpx_client", lambda **_: HTTPHandler(client=httpx.Client(transport=transport)) + ) + monkeypatch.setattr(vertex_batches, "get_async_httpx_client", lambda **_: AsyncHTTPHandler(transport=transport)) + batch: Final = { + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + "input_file_id": VERTEX_BATCH_FILE, + "custom_llm_provider": "vertex_ai", + "vertex_location": "us-central1", + **batch_metadata, + } + if is_async: + await bm.acreate_batch(**batch) + else: + bm.create_batch(**batch) + + sent: Final[httpx.Request] = reply.call_args.args[0] + assert str(sent.url) == VERTEX_BATCH_JOBS_URL + return VERTEX_SENT_JOB.validate_json(sent.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_job_carries_proxy_identity_never_the_callers_batch_metadata( + monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", VERTEX_LABEL_FIELDS) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, + is_async, + { + "metadata": {"user_api_key_alias": "forged", "user_api_key_team_alias": "forged", "app": "bot"}, + "litellm_metadata": dict(VERTEX_PROXY_METADATA), + }, + ) + + labels: Final = VERTEX_SENT_LABELS.validate_python(job["labels"]) + assert list(labels.items()) == [*VERTEX_IDENTITY.items(), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_metadata_alone_never_labels_the_job(monkeypatch: pytest.MonkeyPatch, is_async: bool): + monkeypatch.setattr(litellm, "vertex_request_label_fields", VERTEX_LABEL_FIELDS) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, is_async, {"metadata": {**VERTEX_IDENTITY, "spend_logs_metadata": "cc-1"}} + ) + + assert "labels" not in job + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_vertex_batch_job_with_the_setting_off_is_unlabelled(monkeypatch: pytest.MonkeyPatch, is_async: bool): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + job: Final = await _vertex_batch_job_sent( + monkeypatch, is_async, {"metadata": {"app": "bot"}, "litellm_metadata": dict(VERTEX_PROXY_METADATA)} + ) + + assert "labels" not in job diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 21f42ca2dff..8a03c76f408 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -4,8 +4,12 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final import pytest +from pydantic import TypeAdapter @@ -286,3 +290,138 @@ def test_vertex_generate_content_routes_by_deployment_base_model(base_model, use ) assert (setup.generate_content_provider_config is not None) is uses_native_api + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType( + {**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}, "requester_metadata": {"team": "search"}} +) +_FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +_GENAI_ENTRYPOINTS: Final = ( + "generate_content", + "agenerate_content", + "generate_content_stream", + "agenerate_content_stream", +) +_GEMINI_REPLY: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, +} +_SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def _generate_content_body( + entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object], model: str +) -> Mapping[str, object]: + from unittest.mock import AsyncMock, MagicMock + + import httpx + + from litellm import google_genai + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig + + reply: Final = MagicMock(side_effect=lambda _request: httpx.Response(200, json=_GEMINI_REPLY)) + vertex_token: Final = ("token", "proj") + monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token", MagicMock(return_value=vertex_token)) + monkeypatch.setattr(VertexAIGoogleGenAIConfig, "_ensure_access_token_async", AsyncMock(return_value=vertex_token)) + is_async: Final = entrypoint.startswith("a") + transport: Final = httpx.MockTransport(reply) + client: Final = ( + AsyncHTTPHandler(transport=transport) if is_async else HTTPHandler(client=httpx.Client(transport=transport)) + ) + call: Final = getattr(google_genai, entrypoint)( + model=model, + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + client=client, + metadata=dict(_PROXY_METADATA), + **request_fields, + ) + if is_async: + await call + + sent: Final[httpx.Request] = reply.call_args.args[0] + return _SENT_BODY.validate_json(sent.content) + + +async def _vertex_generate_content_body( + entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object] +) -> Mapping[str, object]: + return await _generate_content_body( + entrypoint, + monkeypatch, + {**request_fields, "vertex_location": "us-central1"}, + "vertex_ai/gemini-3.8-flash", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@pytest.mark.parametrize( + "request_fields", + [{"labels": dict(_FORGED)}, {"extra_body": {"labels": dict(_FORGED)}}], + ids=["top-level-labels", "extra-body-labels"], +) +async def test_vertex_generate_content_sends_identity_labels_a_caller_cannot_forge( + entrypoint: str, request_fields: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + + labels: Final = _SENT_LABELS.validate_python(body["labels"]) + assert list(labels.items()) == [*_IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +async def test_vertex_generate_content_without_caller_labels_still_carries_identity( + entrypoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, {}) + + labels: Final = _SENT_LABELS.validate_python(body["labels"]) + assert list(labels.items()) == [*_IDENTITY.items(), ("team", "search"), ("cost-center", "cc-1")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@pytest.mark.parametrize( + "request_fields, expected_labels", + [ + ({}, "absent"), + ({"labels": dict(_FORGED)}, _FORGED), + ({"extra_body": {"labels": dict(_FORGED)}}, _FORGED), + ({"labels": "not-a-mapping"}, "not-a-mapping"), + ], + ids=["no-labels", "top-level-labels", "extra-body-labels", "non-mapping-labels"], +) +async def test_vertex_generate_content_with_the_setting_off_sends_only_the_callers_labels( + entrypoint: str, + request_fields: Mapping[str, object], + expected_labels: object, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + + assert body.get("labels", "absent") == expected_labels + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +async def test_ai_studio_generate_content_keeps_the_callers_labels_with_the_setting_on( + entrypoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + + body: Final = await _generate_content_body( + entrypoint, monkeypatch, {"labels": dict(_FORGED), "api_key": "test-key"}, "gemini/gemini-3.8-flash" + ) + + assert body["labels"] == _FORGED diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 2c246e065a8..97d4a88f16d 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio import json +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -262,6 +263,46 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +@pytest.mark.parametrize( + "create_data, model", + [ + (CREATE_DATA, "publishers/google/models/gemini-1.5-flash-001"), + (ENDPOINT_CREATE_DATA, TUNED_MODEL_RESOURCE), + ], + ids=["publisher-model", "fine-tuned-endpoint"], +) +def test_create_batch_labels_the_job_from_litellm_metadata_never_metadata( + create_data: Mapping[str, str], model: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", ["user_api_key_alias", "user_api_key_team_alias"]) + h: Final = _make_handler() + client: Final = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}.get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()), + ): + h.create_batch( + _is_async=False, + create_batch_data=create_data, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + litellm_params={ + "metadata": {"user_api_key_alias": "forged"}, + "litellm_metadata": {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}, + }, + ) + + sent: Final = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == model + assert sent["labels"] == {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} + + @pytest.mark.parametrize( "api_base, expected", [ diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index edb6e889814..c790de4bc47 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -1,10 +1,16 @@ import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import TypeAdapter - +import litellm +import litellm.main +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.multimodal_embeddings.transformation import ( VertexAIMultimodalEmbeddingConfig, ) @@ -174,3 +180,71 @@ class TestVertexMultimodalEmbedding: headers={}, ) assert "parameters" not in request + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +TODAYS_INSTANCES: Final = [{"text": "a red apple"}] +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def multimodal_embedding_body( + is_async: bool, monkeypatch: pytest.MonkeyPatch, request_labels: Mapping[str, object] +) -> Mapping[str, object]: + reply: Final = MagicMock( + side_effect=lambda _request: httpx.Response(200, json={"predictions": [{"textEmbedding": [0.1, 0.2]}]}) + ) + monkeypatch.setattr(litellm.main.vertex_multimodal_embedding, "_ensure_access_token", lambda **_: ("token", "proj")) + transport: Final = httpx.MockTransport(reply) + embed: Final = { + "model": "vertex_ai/multimodalembedding@001", + "input": "a red apple", + "vertex_location": "us-central1", + "metadata": dict(PROXY_METADATA), + **request_labels, + } + if is_async: + await litellm.aembedding(**embed, client=AsyncHTTPHandler(transport=transport)) + else: + litellm.embedding(**embed, client=HTTPHandler(client=httpx.Client(transport=transport))) + + sent: Final[httpx.Request] = reply.call_args.args[0] + return SENT_BODY.validate_json(sent.content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("request_labels", [{}, {"labels": dict(FORGED)}], ids=["no-caller-labels", "caller-labels"]) +async def test_setting_off_sends_todays_multimodal_embedding_body( + is_async: bool, request_labels: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + assert await multimodal_embedding_body(is_async, monkeypatch, request_labels) == {"instances": TODAYS_INSTANCES} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "request_labels, expected_labels", + [ + ({}, [*IDENTITY.items(), ("cost-center", "cc-1")]), + ({"labels": dict(FORGED)}, [*IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")]), + ], + ids=["no-caller-labels", "forged-caller-labels"], +) +async def test_multimodal_embedding_sends_identity_labels_a_caller_cannot_forge( + is_async: bool, + request_labels: Mapping[str, object], + expected_labels: Sequence[tuple[str, str]], + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + + body: Final = await multimodal_embedding_body(is_async, monkeypatch, request_labels) + + assert body["instances"] == TODAYS_INSTANCES + assert list(SENT_LABELS.validate_python(body["labels"]).items()) == list(expected_labels) diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index 6721da0003b..c91694a3c5a 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -9,11 +9,14 @@ from pydantic import JsonValue import litellm from litellm.llms.vertex_ai.gemini.transformation import transform_request_body +from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import VertexAIImagenImageGenerationConfig +from litellm.llms.vertex_ai.multimodal_embeddings.transformation import VertexAIMultimodalEmbeddingConfig from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.llms.vertex_ai.vertex_embeddings.transformation import VertexAITextEmbeddingConfig +from litellm.types.router import GenericLiteLLMParams FIELDS: Final = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] IDENTITY: Final = {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"} @@ -63,6 +66,28 @@ def gemini_image_labels(litellm_params: dict[str, JsonValue], optional_params: d return request.get("labels") +def multimodal_embedding_labels( + litellm_params: Mapping[str, JsonValue], optional_params: Mapping[str, JsonValue] +) -> object: + request: Final = VertexAIMultimodalEmbeddingConfig().transform_embedding_request( + model="multimodalembedding@001", + input="a red apple", + optional_params=dict(optional_params), + headers={}, + litellm_params=litellm_params, + ) + return request.get("labels") + + +def google_genai_labels( + litellm_params: Mapping[str, JsonValue], optional_params: Mapping[str, JsonValue] +) -> object: + fields: Final = VertexAIGoogleGenAIConfig().transform_top_level_request_fields( + optional_params or None, GenericLiteLLMParams(**litellm_params) + ) + return (fields or {}).get("labels") + + def rerank_labels(litellm_params: dict[str, JsonValue], _: dict[str, JsonValue]) -> object: request: Final = VertexAIRerankConfig().transform_rerank_request( model="semantic-ranker-default@latest", @@ -108,8 +133,9 @@ def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[st LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels] REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels] -ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels] -CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels] +SETTING_GATED_SURFACES: Final = [multimodal_embedding_labels, google_genai_labels] +ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels, *SETTING_GATED_SURFACES] +CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels, *SETTING_GATED_SURFACES] def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]: @@ -145,6 +171,38 @@ def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: objec assert claude_headers(params, {"labels": {"app": "bot"}}, caller_header) == caller_header +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_sends_no_multimodal_embedding_labels(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = metadata( + spend_logs_metadata={"cost-center": "cc-1"}, requester_metadata={"team": "search"}, **IDENTITY + ) + + assert multimodal_embedding_labels(params, {"labels": {"app": "bot"}}) is None + + +@pytest.mark.parametrize("setting", [None, []]) +def test_setting_off_forwards_google_genai_fields_untouched(setting: object, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) + params: Final = GenericLiteLLMParams(**metadata(spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY)) + fields: Final = {"labels": {"user_api_key_team_alias": "Caller Chosen", "app": "bot"}, "cachedContent": "c-1"} + + assert VertexAIGoogleGenAIConfig().transform_top_level_request_fields(fields, params) is fields + assert VertexAIGoogleGenAIConfig().transform_top_level_request_fields(None, params) is None + + +def test_google_genai_keeps_the_other_top_level_fields_and_drops_an_empty_labels_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + config: Final = VertexAIGoogleGenAIConfig() + fields: Final = {"labels": {"user_api_key_alias": "forged"}, "cachedContent": "c-1"} + + assert config.transform_top_level_request_fields(fields, GenericLiteLLMParams()) == {"cachedContent": "c-1"} + assert config.transform_top_level_request_fields(fields, GenericLiteLLMParams(**metadata(**IDENTITY))) == { + "labels": IDENTITY, + "cachedContent": "c-1", + } + + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("surface", ALL_SURFACES) def test_identity_and_spend_logs_metadata_label_every_vertex_surface( @@ -165,6 +223,8 @@ def test_identity_and_spend_logs_metadata_label_every_vertex_surface( (imagen_labels, 64), (gemini_image_labels, 64), (claude_labels, 32), + (multimodal_embedding_labels, 64), + (google_genai_labels, 64), ], ) def test_identity_survives_a_caller_filling_every_label_slot(