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(