scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(vertex_ai): label Interactions, image edits, Mistral, and the /vertex_ai pass-through, and gate Gemini image labels on the setting

Browse source 
(cherry picked from commit 8c2a905e19)
This commit is contained in:
mateo-berri 2026-11-09 17:17:31 -08:00
parent 9fe1fe0b6a
commit 24ba38af06
18 changed files with 879 additions and 103 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/interactions/http_handler.py
View file

@ -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(

8
litellm/llms/base_llm/interactions/transformation.py
View file

@ -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,

54
litellm/llms/vertex_ai/common_utils.py
View file

@ -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"""

5
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -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))}

10
litellm/llms/vertex_ai/google_genai/transformation.py
View file

@ -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"""

14
litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py
View file

@ -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, [])

5
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
View file

@ -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}}

9
litellm/llms/vertex_ai/interactions/transformation.py
View file

@ -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,

28
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -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:

59
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -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:

54
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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):

135
tests/unit/google_genai/test_google_genai_main.py
View file

@ -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

48
tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
View file

@ -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"),
]

98
tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
View file

@ -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

24
tests/unit/llms/vertex_ai/test_request_labels.py
View file

@ -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)

112
tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py
View file

@ -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",
)

3
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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

306
tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py
View file

@ -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]