From a7254fede6b95e00ac84f5ad46e8711595030434 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 --- 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 | 42 ++++++ ..._ai_multimodal_embedding_transformation.py | 76 +++++++++- .../llms/vertex_ai/test_request_labels.py | 64 +++++++- 14 files changed, 489 insertions(+), 15 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 e7da6586b97..b601922f739 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12343,9 +12343,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( @@ -12458,9 +12461,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 7c4238e3498..3ca288f17f2 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 @@ -20,6 +20,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 ( @@ -81,6 +86,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)) + + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): super().__init__(*args, **kwargs) @@ -97,6 +110,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( @@ -125,6 +139,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 99ef750b5fc..791fc04d1fb 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -158,6 +158,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 @@ -176,12 +177,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 c282ae22ea6..a79939d6d2c 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 f1176f777c7..f3f9c830163 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, @@ -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/unit/batches/test_main.py b/tests/unit/batches/test_main.py index fb97e0e4919..8a04f0cc403 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -1204,3 +1204,102 @@ async def test_delete_batch_output_file(respx_mock: respx.MockRouter, monkeypatc assert delete_route.call_count == 1 assert delete_response.id == "file-output123" 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 diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 5cac61e727e..4b47f4d7331 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -4,10 +4,12 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping +from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter @@ -382,3 +384,138 @@ 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"}) +_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 91a2adb9d19..ff5bdcecfb6 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,8 @@ from __future__ import annotations import asyncio import json +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -259,6 +261,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(