From 24ba38af06fdbe03bc0491eadae4f1ec39103155 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:17:31 -0800 Subject: [PATCH] fix(vertex_ai): label Interactions, image edits, Mistral, and the /vertex_ai pass-through, and gate Gemini image labels on the setting (cherry picked from commit 8c2a905e19263152a07a08f52dc022a7a5fd5314) --- litellm/interactions/http_handler.py | 10 +- .../base_llm/interactions/transformation.py | 8 + litellm/llms/vertex_ai/common_utils.py | 54 +++- .../llms/vertex_ai/gemini/transformation.py | 5 +- .../vertex_ai/google_genai/transformation.py | 10 +- .../vertex_gemini_transformation.py | 14 +- .../vertex_gemini_transformation.py | 5 +- .../vertex_ai/interactions/transformation.py | 9 +- .../vertex_ai_partner_models/main.py | 28 +- .../llm_passthrough_endpoints.py | 59 ++++ .../pass_through_endpoints.py | 54 +++- .../google_genai/test_google_genai_main.py | 135 +++++--- ...est_vertex_ai_image_edit_transformation.py | 48 ++- ...t_vertex_ai_interactions_transformation.py | 98 ++++++ .../llms/vertex_ai/test_request_labels.py | 24 +- ...i_partner_models_mistral_transformation.py | 112 +++++++ .../test_llm_pass_through_endpoints.py | 3 + .../test_llm_passthrough_endpoints.py | 306 ++++++++++++++++++ 18 files changed, 879 insertions(+), 103 deletions(-) diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 09aa656e49b..5cb1931259d 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -161,8 +161,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( @@ -256,8 +257,9 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): headers=headers, ) - if extra_body: - data.update(extra_body) + request_fields: Final = interactions_api_config.transform_top_level_request_fields(extra_body, litellm_params) + if request_fields: + data.update(request_fields) # Logging logging_obj.pre_call( diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index b7bbf7bde73..ff69a21d7f0 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -11,6 +11,7 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -118,6 +119,13 @@ class BaseInteractionsAPIConfig(ABC): raise ValueError("api_base is required") return api_base + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return request_fields + @abstractmethod def transform_request( self, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 364dd31605a..93f319a021d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -28,6 +28,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages @@ -173,7 +174,7 @@ _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) def vertex_request_labels_from_litellm_params( - litellm_params: dict | None, + litellm_params: Mapping[str, object] | None, ) -> dict[str, str] | None: """ Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``: @@ -194,38 +195,57 @@ def vertex_request_labels_from_litellm_params( rm = metadata["requester_metadata"] if not isinstance(rm, dict): continue - labels = {k: v for k, v in rm.items() if isinstance(v, str)} + labels = {k: v for k, v in rm.items() if isinstance(k, str) and isinstance(v, str)} if labels: return labels return None +def pop_vertex_caller_labels( + optional_params: dict | None, + litellm_params: Mapping[str, object] | None, +) -> dict[str, str] | None: + """ + Pop the caller's optional ``labels`` (Gemini-style), falling back to ``requester_metadata`` + on ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``. + """ + raw: Final = optional_params.pop("labels", None) if optional_params is not None else None + labels: Final = ( + {k: v for k, v in raw.items() if isinstance(k, str) and isinstance(v, str)} if isinstance(raw, dict) else None + ) + return labels or vertex_request_labels_from_litellm_params(litellm_params) + + def pop_vertex_request_labels( optional_params: dict | None, litellm_params: dict | None, ) -> dict[str, str] | None: """ - Resolve labels from optional ``labels`` (Gemini-style) and/or - ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]`` - (``requester_metadata``), plus the proxy identity ``litellm.vertex_request_label_fields`` allows. - Pops ``labels`` from optional_params when present. + The caller's labels from ``pop_vertex_caller_labels``, plus the proxy identity + ``litellm.vertex_request_label_fields`` allows. """ - labels: dict[str, str] | None = None - if optional_params is not None and "labels" in optional_params: - raw: Final = optional_params.pop("labels") - if isinstance(raw, dict): - labels = {k: v for k, v in raw.items() if isinstance(v, str)} - if not labels: - labels = vertex_request_labels_from_litellm_params(litellm_params) + request_params: Final = _LITELLM_PARAMS.validate_python(litellm_params or {}) + labels: Final = pop_vertex_caller_labels(optional_params, request_params) allowed_fields: Final = vertex_request_label_fields() if not allowed_fields: - return labels if labels else None - resolved: Final = resolve_vertex_request_labels( - allowed_fields, labels, _LITELLM_PARAMS.validate_python(litellm_params or {}), VERTEX_REQUEST_LABELS_MAX - ) + return labels + resolved: Final = resolve_vertex_request_labels(allowed_fields, labels, request_params, VERTEX_REQUEST_LABELS_MAX) return dict(resolved) if resolved else None +def vertex_top_level_request_fields( + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, +) -> Mapping[str, object] | None: + """The caller's top-level body fields with ``labels`` replaced by ``pop_vertex_request_labels`` while the setting is on.""" + if not vertex_request_label_fields(): + return request_fields + fields: Final = request_fields or {} + unlabelled: Final = {name: value for name, value in fields.items() if name != "labels"} + labels: Final = pop_vertex_request_labels({"labels": fields.get("labels")}, dict(litellm_params)) + return {**unlabelled, "labels": labels} if labels else unlabelled + + class VertexAIModelRoute(str, Enum): """Enum for Vertex AI model routing""" diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 702187e1c12..dfde8150b34 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -29,7 +29,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels +from litellm.llms.vertex_ai.common_utils import pop_vertex_caller_labels, pop_vertex_request_labels from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.types.files import ( get_file_mime_type_for_file_type, @@ -1186,8 +1186,9 @@ def transform_request_body( optional_params.pop("output_config", None) config_fields: Final = GenerationConfig.__annotations__.keys() + request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) - labels: Final = pop_vertex_request_labels(optional_params, litellm_params) + labels: Final = pop_vertex_caller_labels(optional_params, request_params) filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))} diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 78fe644a872..9982763a846 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -6,8 +6,7 @@ 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.llms.vertex_ai.common_utils import vertex_top_level_request_fields from litellm.types.router import GenericLiteLLMParams @@ -46,12 +45,7 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): 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 + return vertex_top_level_request_fields(request_fields, litellm_params) def _camel_to_snake(self, camel_str: str) -> str: """Convert camelCase to snake_case""" diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 8e12add0316..09c8596043d 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -12,8 +12,9 @@ import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, vertex_top_level_request_fields from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.vertex_ai import ( @@ -190,7 +191,16 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value } - request_body: Final[dict[str, object]] = {"contents": contents, "generationConfig": generation_config} + label_fields: Final = ( + vertex_top_level_request_fields({"labels": litellm_params.get("labels")}, litellm_params) + if vertex_request_label_fields() + else None + ) + request_body: Final[dict[str, object]] = { + "contents": contents, + "generationConfig": generation_config, + **(label_fields or {}), + } payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 0655de9fc37..98056e47aa1 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.gemini.common_utils import ( from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, pop_vertex_request_labels from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.request_labels import vertex_request_label_fields from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( @@ -257,7 +258,9 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): optional_fields: Final = { "tools": params.get("tools"), "toolConfig": params.get("toolConfig"), - "labels": pop_vertex_request_labels(optional_params, litellm_params), + "labels": pop_vertex_request_labels(optional_params, litellm_params) + if vertex_request_label_fields() + else None, } return {**request_body, **{key: value for key, value in optional_fields.items() if value}} diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index 36965fe666f..7b600c4dc57 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -4,7 +4,7 @@ from typing import Final from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig -from litellm.llms.vertex_ai.common_utils import validate_vertex_location +from litellm.llms.vertex_ai.common_utils import validate_vertex_location, vertex_top_level_request_fields from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.router import GenericLiteLLMParams @@ -111,6 +111,13 @@ class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): collection_url: Final = self._target(api_base, params).collection_url return f"{collection_url}?alt=sse" if stream else collection_url + def transform_top_level_request_fields( + self, + request_fields: Mapping[str, object] | None, + litellm_params: GenericLiteLLMParams, + ) -> Mapping[str, object] | None: + return vertex_top_level_request_fields(request_fields, litellm_params) + def _interaction_by_id_request( self, interaction_id: str, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index c855a073648..d2459b4afa1 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -1,14 +1,20 @@ # What is this? ## API Handler for calling Vertex AI Partner Models -from collections.abc import Callable +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm import LlmProviders +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.types.llms.vertex_ai import VertexPartnerProvider from litellm.utils import ModelResponse @@ -16,6 +22,7 @@ from ...custom_httpx.llm_http_handler import BaseLLMHTTPHandler from ..vertex_llm_base import VertexBase base_llm_http_handler = BaseLLMHTTPHandler() +_PARAMS: Final = TypeAdapter(Mapping[str, object]) class VertexAIError(Exception): @@ -158,6 +165,21 @@ class VertexAIPartnerModels(VertexBase): model=model, ) + label_fields: Final = vertex_request_label_fields() if partner == VertexPartnerProvider.mistralai else () + labels_header: Final = ( + vertex_partner_labels_header( + label_fields, + _PARAMS.validate_python(optional_params).get("labels"), + _PARAMS.validate_python(headers or {}), + _PARAMS.validate_python(litellm_params), + ) + if label_fields + else None + ) + if label_fields and "labels" in optional_params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + vertex_labels_headers: Final = {} if labels_header is None else {VERTEX_LABELS_HEADER: labels_header} + if "codestral" in model or "mistral" in model: model = model.split("@")[0] @@ -179,6 +201,7 @@ class VertexAIPartnerModels(VertexBase): logger_fn=logger_fn, timeout=timeout, encoding=encoding, + headers=vertex_labels_headers, ) elif "claude" in model: # Build a new dict so we never mutate the shared deployment extra_headers object. @@ -249,6 +272,9 @@ class VertexAIPartnerModels(VertexBase): encoding=encoding, custom_llm_provider="vertex_ai", custom_endpoint=True, + headers={"Content-Type": "application/json", **vertex_labels_headers} + if vertex_labels_headers + else None, ) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 978882ef3f1..22ecbcd0f74 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -68,6 +68,14 @@ from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request +from litellm.llms.vertex_ai.request_labels import ( + VERTEX_LABELS_HEADER, + VERTEX_REQUEST_LABELS_MAX, + is_vertex_labels_header, + resolve_vertex_request_labels, + vertex_partner_labels_header, + vertex_request_label_fields, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * @@ -109,6 +117,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + UpstreamRequest, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -2566,6 +2575,53 @@ def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str] return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) +_VERTEX_BODY_LABELLED_METHODS: Final = frozenset( + {"generateContent", "streamGenerateContent", "predict", "batchPredictionJobs"} +) +_VERTEX_HEADER_LABELLED_METHODS: Final = frozenset({"rawPredict", "streamRawPredict"}) + + +def _vertex_labelled_method(http_method: str, endpoint: str) -> str | None: + if http_method != "POST": + return None + resource: Final = endpoint.rstrip("/").rsplit("/", 1)[-1] + return resource.partition(":")[2] or resource + + +def _with_vertex_body_labels( + allowed_fields: tuple[str, ...], body: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_vertex_request_labels( + allowed_fields, body.get("labels"), litellm_params, VERTEX_REQUEST_LABELS_MAX + ) + unlabelled: Final = {key: value for key, value in body.items() if key != "labels"} + return {**unlabelled, "labels": dict(resolved)} if resolved else unlabelled + + +def _with_vertex_labels_header( + allowed_fields: tuple[str, ...], headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + labels_header: Final = vertex_partner_labels_header(allowed_fields, None, headers, litellm_params) + unlabelled: Final = {name: value for name, value in headers.items() if not is_vertex_labels_header(name)} + return unlabelled if labels_header is None else {**unlabelled, VERTEX_LABELS_HEADER: labels_header} + + +def _with_vertex_request_labels( + vertex_method: str | None, + headers: Mapping[str, str], + body: Mapping[str, object] | None, + litellm_params: Mapping[str, object], +) -> UpstreamRequest: + allowed_fields: Final = vertex_request_label_fields() + if not allowed_fields: + return headers, body + if body is not None and (vertex_method in _VERTEX_BODY_LABELLED_METHODS or "labels" in body): + return headers, _with_vertex_body_labels(allowed_fields, body, litellm_params) + if vertex_method in _VERTEX_HEADER_LABELLED_METHODS: + return _with_vertex_labels_header(allowed_fields, headers, litellm_params), body + return headers, body + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], ) -> BaseVertexAIPassThroughHandler: @@ -3012,6 +3068,9 @@ async def _base_vertex_proxy_route( target=target, custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, + upstream_request_filter=partial( + _with_vertex_request_labels, _vertex_labelled_method(request.method, encoded_endpoint) + ), ) # dynamically construct pass-through endpoint based on incoming path try: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a41e2ea068..aaac97360c8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -465,7 +465,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): url: httpx.URL, headers: dict, requested_query_params: dict | None = None, - _parsed_body: dict | None = None, + _parsed_body: Mapping[str, object] | None = None, forward_multipart: bool = False, ) -> httpx.Response: """ @@ -1087,6 +1087,28 @@ from litellm.passthrough.timeout_utils import ( resolve_pass_through_request_timeout, ) +UpstreamRequest: TypeAlias = tuple[Mapping[str, str], Mapping[str, object] | None] +UpstreamRequestFilter: TypeAlias = Callable[ + [Mapping[str, str], Mapping[str, object] | None, Mapping[str, object]], UpstreamRequest +] +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def _unfiltered_request( + headers: Mapping[str, str], body: Mapping[str, object] | None, litellm_params: Mapping[str, object] +) -> UpstreamRequest: + return headers, body + + +def _filtered_upstream_request( + upstream_request_filter: UpstreamRequestFilter, + headers: Mapping[str, str], + body: Mapping[str, object] | None, + litellm_params: Mapping[str, object], +) -> tuple[dict[str, str], Mapping[str, object] | None]: # mutable-ok: the relay helpers take a dict of headers + filtered_headers, filtered_body = upstream_request_filter(headers, body, litellm_params) + return dict(filtered_headers), filtered_body + async def pass_through_request( request: Request, @@ -1103,6 +1125,7 @@ async def pass_through_request( custom_llm_provider: str | None = None, guardrails_config: dict | None = None, timeout: float | None = None, + upstream_request_filter: UpstreamRequestFilter = _unfiltered_request, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1123,6 +1146,8 @@ async def pass_through_request( guardrails_config: Optional field - guardrails configuration for passthrough endpoint timeout: Optional per-endpoint timeout in seconds. Falls back to general_settings.pass_through_request_timeout, then 600s. + upstream_request_filter: Optional last rewrite of the headers and JSON body sent upstream, given the + litellm_params the request is logged with """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1461,13 +1486,20 @@ async def pass_through_request( else: logging_url = str(url) + "?" + requested_query_params_str + outgoing_body: Final = cast("Mapping[str, object] | None", _parsed_body) # cast-ok: json + outbound_headers, upstream_body = _filtered_upstream_request( + upstream_request_filter, + upstream_headers, + outgoing_body, + _LITELLM_PARAMS.validate_python(kwargs["litellm_params"]), + ) logging_obj.pre_call( input=[{"role": "user", "content": safe_dumps(_parsed_body)}], api_key="", additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": upstream_headers, + "headers": outbound_headers, }, ) stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body( @@ -1484,7 +1516,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=upstream_headers, + headers=outbound_headers, requested_query_params=requested_query_params, stream=True, ) @@ -1496,7 +1528,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=upstream_headers, + headers=outbound_headers, content=state_raw_body, ) if state_raw_body is not None @@ -1504,8 +1536,8 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=upstream_headers, - json=_parsed_body, + headers=outbound_headers, + json=upstream_body, ) ) @@ -1582,7 +1614,7 @@ async def pass_through_request( raw_body_request: Final = async_client.build_request( request.method, url, - headers=upstream_headers, + headers=outbound_headers, params=requested_query_params, content=state_raw_body, ) @@ -1592,9 +1624,9 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=upstream_headers, + headers=outbound_headers, requested_query_params=requested_query_params, - _parsed_body=_parsed_body, + _parsed_body=upstream_body, forward_multipart=is_multipart, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -2107,6 +2139,7 @@ def create_pass_through_route( guardrails: dict[str, object] | None = None, config_file_path: str | None = None, timeout: float | None = None, + upstream_request_filter: UpstreamRequestFilter = _unfiltered_request, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2246,6 +2279,7 @@ def create_pass_through_route( custom_llm_provider=custom_llm_provider, guardrails_config=cast(dict | None, param_guardrails), timeout=cast(float | None, param_timeout), + upstream_request_filter=upstream_request_filter, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 8a03c76f408..fd36b2fcab8 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -6,15 +6,16 @@ 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 +from typing import Final, Literal import pytest from pydantic import TypeAdapter - - - import litellm +from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams @pytest.mark.asyncio @@ -298,11 +299,10 @@ _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", +_HANDLER_CALLS: Final = pytest.mark.parametrize( + ("is_async", "stream"), + [(False, False), (True, False), (False, True), (True, True)], + ids=["sync", "async", "sync-stream", "async-stream"], ) _GEMINI_REPLY: Final = { "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], @@ -312,32 +312,67 @@ _SENT_BODY: Final = TypeAdapter(Mapping[str, object]) _SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) +class _SignedInVertexGoogleGenAIConfig(VertexAIGoogleGenAIConfig): + def _ensure_access_token( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: + return "token", "proj" + + async def _ensure_access_token_async( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: + return "token", "proj" + + async def _generate_content_body( - entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object], model: str + config: GoogleGenAIConfig, + litellm_params: GenericLiteLLMParams, + extra_body: Mapping[str, object], + *, + is_async: bool, + stream: bool, ) -> Mapping[str, object]: - from unittest.mock import AsyncMock, MagicMock + import datetime + from unittest.mock import MagicMock import httpx - from litellm import google_genai + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler 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, + call: Final = BaseLLMHTTPHandler().generate_content_handler( + model="gemini-3.8-flash", contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + generate_content_provider_config=config, + generate_content_config_dict={}, + tools=None, + custom_llm_provider=config.custom_llm_provider, + litellm_params=litellm_params, + logging_obj=Logging( + model="gemini-3.8-flash", + messages=[], + stream=stream, + call_type="agenerate_content" if is_async else "generate_content", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="labels-test", + function_id="labels-test", + ), + extra_body=dict(extra_body), + _is_async=is_async, client=client, - metadata=dict(_PROXY_METADATA), - **request_fields, + stream=stream, ) if is_async: await call @@ -347,81 +382,77 @@ async def _generate_content_body( async def _vertex_generate_content_body( - entrypoint: str, monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object] + extra_body: Mapping[str, object], *, is_async: bool, stream: bool ) -> Mapping[str, object]: return await _generate_content_body( - entrypoint, - monkeypatch, - {**request_fields, "vertex_location": "us-central1"}, - "vertex_ai/gemini-3.8-flash", + _SignedInVertexGoogleGenAIConfig(), + GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), vertex_location="us-central1"), + extra_body, + is_async=is_async, + stream=stream, ) @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"], -) +@_HANDLER_CALLS async def test_vertex_generate_content_sends_identity_labels_a_caller_cannot_forge( - entrypoint: str, request_fields: Mapping[str, object], monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + body: Final = await _vertex_generate_content_body({"labels": dict(_FORGED)}, is_async=is_async, stream=stream) labels: Final = _SENT_LABELS.validate_python(body["labels"]) assert list(labels.items()) == [*_IDENTITY.items(), ("app", "bot"), ("cost-center", "cc-1")] @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS async def test_vertex_generate_content_without_caller_labels_still_carries_identity( - entrypoint: str, monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, {}) + body: Final = await _vertex_generate_content_body({}, is_async=is_async, stream=stream) labels: Final = _SENT_LABELS.validate_python(body["labels"]) assert list(labels.items()) == [*_IDENTITY.items(), ("team", "search"), ("cost-center", "cc-1")] @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS @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"], + "extra_body, expected_labels", + [({}, "absent"), ({"labels": dict(_FORGED)}, _FORGED), ({"labels": "not-a-mapping"}, "not-a-mapping")], + ids=["no-labels", "labels", "non-mapping-labels"], ) async def test_vertex_generate_content_with_the_setting_off_sends_only_the_callers_labels( - entrypoint: str, - request_fields: Mapping[str, object], + is_async: bool, + stream: bool, + extra_body: Mapping[str, object], expected_labels: object, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "vertex_request_label_fields", None) - body: Final = await _vertex_generate_content_body(entrypoint, monkeypatch, request_fields) + body: Final = await _vertex_generate_content_body(extra_body, is_async=is_async, stream=stream) assert body.get("labels", "absent") == expected_labels @pytest.mark.asyncio -@pytest.mark.parametrize("entrypoint", _GENAI_ENTRYPOINTS) +@_HANDLER_CALLS async def test_ai_studio_generate_content_keeps_the_callers_labels_with_the_setting_on( - entrypoint: str, monkeypatch: pytest.MonkeyPatch + is_async: bool, stream: bool, 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" + GoogleGenAIConfig(), + GenericLiteLLMParams(metadata=dict(_PROXY_METADATA), api_key="test-key"), + {"labels": dict(_FORGED)}, + is_async=is_async, + stream=stream, ) assert body["labels"] == _FORGED diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 9b2f8a40e0e..51838c114a1 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -1,14 +1,16 @@ import base64 import json import os +from collections.abc import Mapping from datetime import datetime from io import BytesIO +from types import MappingProxyType from typing import Dict, Final from unittest.mock import MagicMock, patch import httpx import pytest -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging @@ -636,3 +638,47 @@ def test_imagen_image_edit_maps_malformed_predictions_to_a_connection_error_with 500, False, ) + + +_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +_PROXY_METADATA: Final = MappingProxyType({**_IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +_FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +_SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +_SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +def _sent_gemini_edit_body() -> Mapping[str, object]: + sent: Final[list[httpx.Request]] = [] + + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, json={"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}]} + ) + + litellm.image_edit( + model="vertex_ai/gemini-3.1-flash-image", + image=BytesIO(b"\x89PNG\r\n\x1a\nimage"), + prompt="mountains", + api_base="https://vertex.test/edit", + metadata=dict(_PROXY_METADATA), + labels=dict(_FORGED), + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(reply))), + ) + return _SENT_BODY.validate_json(sent[0].content) + + +def test_vertex_gemini_edit_sends_identity_labels_only_while_the_setting_is_on(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_gemini_edit_body() + monkeypatch.setattr(litellm, "vertex_request_label_fields", _LABEL_FIELDS) + setting_on: Final = _sent_gemini_edit_body() + + assert sorted(setting_off) == ["contents", "generationConfig"] + assert {key: value for key, value in setting_on.items() if key != "labels"} == setting_off + assert list(_SENT_LABELS.validate_python(setting_on["labels"]).items()) == [ + *_IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] diff --git a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py index 3364b1b3872..661f5200601 100644 --- a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py +++ b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -1,7 +1,18 @@ +import json +from collections.abc import Mapping +from datetime import datetime +from types import MappingProxyType +from typing import Final + +import httpx import pytest +from pydantic import TypeAdapter import litellm +from litellm.interactions.http_handler import InteractionsHTTPHandler from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig, ) @@ -229,3 +240,90 @@ class TestInteractionByIdRequests: assert url == f"{GLOBAL_BASE}/abc123:cancel" assert request_body == {} + + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +async def _created_interaction_body( + config: VertexAIInteractionsConfig, extra_body: Mapping[str, object], *, is_async: bool +) -> Mapping[str, object]: + sent: Final[list[httpx.Request]] = [] # mutable-ok: MockTransport hands each request to a callback + + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"id": "interaction-1", "status": "completed", "outputs": []}) + + transport: Final = httpx.MockTransport(reply) + litellm_params: Final = GenericLiteLLMParams(vertex_project="test-proj", metadata=dict(PROXY_METADATA)) + logging_obj: Final = Logging( + model="gemini-omni-flash-preview", + messages=[], + stream=False, + call_type="acreate_interaction" if is_async else "create_interaction", + start_time=datetime(2026, 1, 1), + litellm_call_id="labels-test", + function_id="labels-test", + ) + handler: Final = InteractionsHTTPHandler() + if is_async: + await handler.async_create_interaction( + interactions_api_config=config, + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + logging_obj=logging_obj, + model="gemini-omni-flash-preview", + input="hi", + extra_body=extra_body, + client=AsyncHTTPHandler(transport=transport), + ) + else: + handler.create_interaction( + interactions_api_config=config, + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + logging_obj=logging_obj, + model="gemini-omni-flash-preview", + input="hi", + extra_body=extra_body, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + return SENT_BODY.validate_python(json.loads(sent[0].content)) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_create_replaces_forged_extra_body_labels_with_identity_while_the_setting_is_on( + config: VertexAIInteractionsConfig, monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + + body: Final = await _created_interaction_body(config, {"labels": dict(FORGED)}, is_async=is_async) + + assert list(SENT_LABELS.validate_python(body["labels"]).items()) == [ + *IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_create_sends_caller_labels_as_is_while_the_setting_is_off( + config: VertexAIInteractionsConfig, monkeypatch: pytest.MonkeyPatch, is_async: bool +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + + with_labels: Final = await _created_interaction_body(config, {"labels": dict(FORGED)}, is_async=is_async) + without_labels: Final = await _created_interaction_body(config, {}, is_async=is_async) + + assert with_labels["labels"] == dict(FORGED) + assert "labels" not in without_labels diff --git a/tests/unit/llms/vertex_ai/test_request_labels.py b/tests/unit/llms/vertex_ai/test_request_labels.py index cbe990dceb7..0ebbbd2013f 100644 --- a/tests/unit/llms/vertex_ai/test_request_labels.py +++ b/tests/unit/llms/vertex_ai/test_request_labels.py @@ -131,11 +131,12 @@ def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[st return decoded_labels(claude_headers(litellm_params, optional_params)) -LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels] +LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels] REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels] -SETTING_GATED_SURFACES: Final = [multimodal_embedding_labels, google_genai_labels] +SETTING_GATED_SURFACES: Final = [gemini_image_labels, multimodal_embedding_labels, google_genai_labels] ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels, *SETTING_GATED_SURFACES] CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels, *SETTING_GATED_SURFACES] +BODY_SURFACES_ADDED_BY_THE_SETTING: Final = [gemini_image_labels, multimodal_embedding_labels] def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]: @@ -172,13 +173,16 @@ def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: objec @pytest.mark.parametrize("setting", [None, []]) -def test_setting_off_sends_no_multimodal_embedding_labels(setting: object, monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("surface", BODY_SURFACES_ADDED_BY_THE_SETTING) +def test_setting_off_sends_no_labels_where_none_were_sent_before( + setting: object, surface: LabelSource, monkeypatch: pytest.MonkeyPatch +): monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) params: Final = metadata( spend_logs_metadata={"cost-center": "cc-1"}, requester_metadata={"team": "search"}, **IDENTITY ) - assert multimodal_embedding_labels(params, {"labels": {"app": "bot"}}) is None + assert surface(params, {"labels": {"app": "bot"}}) is None @pytest.mark.parametrize("setting", [None, []]) @@ -273,6 +277,18 @@ def test_gemini_extra_body_labels_cannot_replace_or_drop_identity( assert gemini_labels(metadata(**IDENTITY), {"extra_body": {"labels": extra_labels}}) == expected +def test_gemini_extra_body_labels_outrank_spend_logs_metadata_at_the_cap(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS) + spend_logs: Final = {f"spend-{index:02d}": "v" for index in range(64)} + caller_labels: Final = {"labels": {"env": "prod"}, "extra_body": {"labels": {"app": "bot"}}} + + labels: Final = gemini_labels(metadata(**IDENTITY, spend_logs_metadata=spend_logs), caller_labels) + + assert isinstance(labels, dict) + assert len(labels) == 64 + assert {key: labels[key] for key in (*IDENTITY, "env", "app")} == {**IDENTITY, "env": "prod", "app": "bot"} + + @pytest.mark.parametrize("setting", [None, []]) def test_setting_off_merges_gemini_extra_body_labels_as_sent(setting: object, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "vertex_request_label_fields", setting) diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index 15df8e47af3..4a754325be5 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,15 @@ +import base64 +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +import httpx import pytest +from pydantic import TypeAdapter import litellm +import litellm.main +from litellm.llms.custom_httpx.http_handler import HTTPHandler @pytest.fixture @@ -30,3 +39,106 @@ def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model drop_params=True, ) assert "reasoning_effort" not in dropped + + +LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) +PROXY_METADATA: Final = MappingProxyType({**IDENTITY, "spend_logs_metadata": {"cost-center": "cc-1"}}) +FORGED: Final = MappingProxyType({"user_api_key_alias": "forged", "app": "bot"}) +CHAT_REPLY: Final = MappingProxyType( + { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": "mistral-large-2411", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } +) +SENT_BODY: Final = TypeAdapter(Mapping[str, object]) +SENT_LABELS: Final = TypeAdapter(Mapping[str, str]) + + +def _transport(sent: list[httpx.Request]) -> httpx.MockTransport: + def reply(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json=dict(CHAT_REPLY)) + + return httpx.MockTransport(reply) + + +def _send_mistral_chat(transport: httpx.MockTransport, monkeypatch: pytest.MonkeyPatch) -> None: + litellm.completion( + model="vertex_ai/mistral-large-2411", + messages=[{"role": "user", "content": "hi"}], + vertex_location="us-central1", + metadata=dict(PROXY_METADATA), + labels=dict(FORGED), + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + +def _send_codestral_fim(transport: httpx.MockTransport, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=transport))) + monkeypatch.setitem( + litellm.known_tokenizer_config, + "codestral-2", + {"status": "success", "tokenizer": {"chat_template": "{{ messages[0]['content'] }}"}}, + ) + litellm.text_completion( + model="vertex_ai/codestral-2", + prompt="def add(", + vertex_location="us-central1", + metadata=dict(PROXY_METADATA), + labels=dict(FORGED), + ) + + +MISTRAL_CALLS: Final = pytest.mark.parametrize( + "send, todays_body", + [ + ( + _send_mistral_chat, + {"model": "mistral-large-2411", "messages": [{"role": "user", "content": "hi"}], "stream": False}, + ), + (_send_codestral_fim, {"model": "codestral-2", "prompt": "def add("}), + ], + ids=["chat", "codestral-fim"], +) + + +def _sent_mistral_request( + send: Callable[[httpx.MockTransport, pytest.MonkeyPatch], None], monkeypatch: pytest.MonkeyPatch +) -> httpx.Request: + monkeypatch.setattr( + litellm.main.vertex_partner_models_chat_completion, "_ensure_access_token", lambda **_: ("token", "proj") + ) + sent: Final[list[httpx.Request]] = [] + send(_transport(sent), monkeypatch) + assert [request.url.path.rsplit(":", 1)[-1] for request in sent] == ["rawPredict"] + return sent[0] + + +@MISTRAL_CALLS +def test_mistral_labels_header_carries_identity_only_while_the_setting_is_on( + send: Callable[[httpx.MockTransport, pytest.MonkeyPatch], None], + todays_body: Mapping[str, object], + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "vertex_request_label_fields", None) + setting_off: Final = _sent_mistral_request(send, monkeypatch) + monkeypatch.setattr(litellm, "vertex_request_label_fields", LABEL_FIELDS) + setting_on: Final = _sent_mistral_request(send, monkeypatch) + + assert "x-vertex-ai-labels" not in setting_off.headers + assert SENT_BODY.validate_json(setting_off.content) == {**todays_body, "labels": dict(FORGED)} + assert list(SENT_LABELS.validate_json(base64.b64decode(setting_on.headers["x-vertex-ai-labels"])).items()) == [ + *IDENTITY.items(), + ("app", "bot"), + ("cost-center", "cc-1"), + ] + assert SENT_BODY.validate_json(setting_on.content) == todays_body + assert (setting_on.headers["content-type"], setting_on.headers["authorization"]) == ( + "application/json", + "Bearer token", + ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..e5e792a832c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -369,6 +369,7 @@ class TestVertexAIPassThroughHandler: target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.asyncio @@ -468,6 +469,7 @@ class TestVertexAIPassThroughHandler: target=f"https://aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.parametrize( @@ -563,6 +565,7 @@ class TestVertexAIPassThroughHandler: target=f"https://{default_location}-aiplatform.googleapis.com/v1/projects/{default_project}/locations/{default_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {default_credentials}"}, is_streaming_request=False, + upstream_request_filter=mock.ANY, ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index 98ea2269252..98269df8fa2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -1,10 +1,17 @@ import asyncio +import base64 import importlib import json +from collections.abc import Iterator, Mapping from typing import Final +import httpx import pytest +import responses import respx +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi.testclient import TestClient import litellm from litellm import Router @@ -257,3 +264,302 @@ async def test_gemini_pass_through_endpoint_forwards_key_and_body( assert outbound.url.path == f"/{endpoint}" assert outbound.url.params["key"] == "gemini-test-key" assert json.loads(outbound.content) == json.loads(body) + + +_VERTEX_DEPLOYMENT_PATH: Final = ( + "v1/projects/deployment-project/locations/global/publishers/anthropic/models/claude-sonnet-4-5:rawPredict" +) +_VERTEX_CALLER_PATH: Final = ( + "v1/projects/caller-project/locations/global/publishers/anthropic/models/claude-haiku-4-5:rawPredict" +) +_VERTEX_DEPLOYMENT_PROJECT: Final = "deployment-project" +_VERTEX_TOKEN_URI: Final = "https://oauth2.googleapis.com/token" +_VERTEX_MINTED_TOKEN: Final = "minted-vertex-token" +_VERTEX_CALLER_GOOGLE_TOKEN: Final = "caller-google-token" +_VERTEX_CLAUDE_BODY: Final = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} +_VERTEX_CLAUDE_REPLY: Final = {"id": "msg_vertex", "type": "message", "content": []} + + +def _vertex_service_account_json() -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _VERTEX_DEPLOYMENT_PROJECT, + "private_key_id": "test", + "private_key": private_key, + "client_email": f"test@{_VERTEX_DEPLOYMENT_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "token_uri": _VERTEX_TOKEN_URI, + } + ) + + +@pytest.fixture +def vertex_proxy( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, isolate_passthrough_endpoint_router_state: None +) -> Iterator[TestClient]: + from litellm.proxy import proxy_server + + router: Final = Router( + model_list=[ + { + "model_name": "claude-sonnet-4-5", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-4-5", + "vertex_project": _VERTEX_DEPLOYMENT_PROJECT, + "vertex_location": "global", + "vertex_credentials": _vertex_service_account_json(), + "use_in_pass_through": True, + }, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + with responses.RequestsMock(assert_all_requests_are_fired=False) as google_oauth: + google_oauth.post( + _VERTEX_TOKEN_URI, + json={"access_token": _VERTEX_MINTED_TOKEN, "expires_in": 3600, "token_type": "Bearer"}, + ) + yield TestClient(proxy_server.app) + + +_VERTEX_GEMINI_MODEL_PATH: Final = "v1/projects/deployment-project/locations/global/publishers/google/models/gemini-2.5-flash" +_VERTEX_LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") +_VERTEX_IDENTITY_LABELS: Final = {"user_api_key_alias": "search-bot", "user_api_key_team_alias": "search-team"} +_VERTEX_SPEND_LOGS_LABELS: Final = {"cost_center": "cc-42"} +_VERTEX_FORGED_LABELS: Final = {"user_api_key_team_alias": "other-team", "env": "prod", "app": "chat"} +_VERTEX_GEMINI_REPLY: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]} + + +@pytest.fixture +def search_team_caller(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def authenticate(request: object, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key=api_key, + key_alias="search-bot", + team_id="team-search", + team_alias="search-team", + metadata={"spend_logs_metadata": _VERTEX_SPEND_LOGS_LABELS}, + ) + + monkeypatch.setattr(proxy_server, "user_custom_auth", authenticate) + + +def _gemini_upstream(respx_mock: "respx.MockRouter", action: str) -> respx.Route: + reply: Final = ( + httpx.Response(200, json=_VERTEX_GEMINI_REPLY) + if action == "generateContent" + else httpx.Response( + 200, text=f"data: {json.dumps(_VERTEX_GEMINI_REPLY)}\n\n", headers={"content-type": "text/event-stream"} + ) + ) + return respx_mock.post(url__startswith=f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:").mock( + return_value=reply + ) + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("action", ["generateContent", "streamGenerateContent"]) +@pytest.mark.parametrize( + ("caller_labels", "expected_labels"), + [ + pytest.param( + _VERTEX_FORGED_LABELS, + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}, + id="forged-team-replaced-caller-labels-kept", + ), + pytest.param(None, {**_VERTEX_IDENTITY_LABELS, **_VERTEX_SPEND_LOGS_LABELS}, id="identity-added"), + ], +) +def test_vertex_gemini_pass_through_bills_the_callers_own_identity( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + action: str, + caller_labels: Mapping[str, str] | None, + expected_labels: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = _gemini_upstream(respx_mock, action) + body: Final = { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + **({} if caller_labels is None else {"labels": caller_labels}), + } + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:{action}", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + sent: Final = json.loads(upstream.calls.last.request.content) + assert list(sent["labels"].items()) == list(expected_labels.items()) + assert {key: value for key, value in sent.items() if key != "labels"} == { + key: value for key, value in body.items() if key != "labels" + } + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize("action", ["generateContent", "streamGenerateContent"]) +def test_vertex_gemini_pass_through_forwards_the_body_untouched_with_the_setting_off( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", action: str +) -> None: + upstream: Final = _gemini_upstream(respx_mock, action) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:{action}", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert list(json.loads(upstream.calls.last.request.content).items()) == list(body.items()) + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize( + "endpoint", + [ + f"{_VERTEX_GEMINI_MODEL_PATH}:predict", + "v1/projects/deployment-project/locations/global/batchPredictionJobs", + ], + ids=["predict", "batch-prediction-job"], +) +def test_vertex_pass_through_adds_identity_labels_to_predict_and_batch_jobs( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{endpoint}").mock( + return_value=httpx.Response(200, json={}) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{endpoint}", json={"instances": [{"content": "hi"}]}, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert json.loads(upstream.calls.last.request.content) == { + "instances": [{"content": "hi"}], + "labels": {**_VERTEX_IDENTITY_LABELS, **_VERTEX_SPEND_LOGS_LABELS}, + } + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_pass_through_replaces_forged_labels_on_any_endpoint( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:countTokens").mock( + return_value=httpx.Response(200, json={"totalTokens": 1}) + ) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "labels": _VERTEX_FORGED_LABELS} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:countTokens", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert list(json.loads(upstream.calls.last.request.content)["labels"].items()) == list( + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}.items() + ) + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_pass_through_adds_no_labels_where_the_caller_sent_none_outside_billed_methods( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter", monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_MODEL_PATH}:countTokens").mock( + return_value=httpx.Response(200, json={"totalTokens": 1}) + ) + body: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_GEMINI_MODEL_PATH}:countTokens", json=body, headers={"x-litellm-api-key": "sk-caller"} + ) + + assert response.status_code == 200 + assert json.loads(upstream.calls.last.request.content) == body + assert "x-vertex-ai-labels" not in upstream.calls.last.request.headers + + +def _labels_header(labels: Mapping[str, str]) -> str: + return base64.b64encode(json.dumps(dict(labels)).encode()).decode() + + +@pytest.mark.usefixtures("search_team_caller") +@pytest.mark.parametrize( + ("model_path", "caller_headers"), + [ + pytest.param( + _VERTEX_DEPLOYMENT_PATH, + {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_FORGED_LABELS)}, + id="x-pass-header-on-deployment-credentials", + ), + pytest.param( + _VERTEX_CALLER_PATH, + { + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + "X-Vertex-AI-Labels": _labels_header(_VERTEX_FORGED_LABELS), + }, + id="forwarded-header-on-caller-credentials", + ), + ], +) +def test_vertex_claude_pass_through_replaces_a_forged_labels_header( + vertex_proxy: TestClient, + respx_mock: "respx.MockRouter", + monkeypatch: pytest.MonkeyPatch, + model_path: str, + caller_headers: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "vertex_request_label_fields", list(_VERTEX_LABEL_FIELDS)) + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{model_path}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{model_path}", json=_VERTEX_CLAUDE_BODY, headers={"x-litellm-api-key": "sk-caller", **caller_headers} + ) + + assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY) + outbound: Final = upstream.calls.last.request + sent_headers: Final = outbound.headers.get_list("x-vertex-ai-labels") + assert len(sent_headers) == 1 + assert list(json.loads(base64.b64decode(sent_headers[0])).items()) == list( + {**_VERTEX_IDENTITY_LABELS, "env": "prod", "app": "chat", **_VERTEX_SPEND_LOGS_LABELS}.items() + ) + assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY + + +@pytest.mark.usefixtures("search_team_caller") +def test_vertex_claude_pass_through_forwards_the_callers_labels_header_with_the_setting_off( + vertex_proxy: TestClient, respx_mock: "respx.MockRouter" +) -> None: + upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_CALLER_PATH}").mock( + return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY) + ) + forged_header: Final = _labels_header(_VERTEX_FORGED_LABELS) + + response: Final = vertex_proxy.post( + f"/vertex_ai/{_VERTEX_CALLER_PATH}", + json=_VERTEX_CLAUDE_BODY, + headers={ + "x-litellm-api-key": "sk-caller", + "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}", + "X-Vertex-AI-Labels": forged_header, + }, + ) + + assert response.status_code == 200 + assert upstream.calls.last.request.headers.get_list("x-vertex-ai-labels") == [forged_header]