mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(vertex_ai): label Vertex requests with the caller's key, team, org, and project
Adds the opt-in litellm.vertex_request_label_fields setting, the Vertex twin of bedrock_request_metadata_fields. When set, Gemini, embeddings, Imagen and Gemini image generation, rerank, and Claude on Vertex (chat and /v1/messages) send the listed proxy identity fields and spend_logs_metadata as request labels, identity first, with the user_api_key_ prefix reserved so a caller cannot forge or evict them. Identity text is rewritten to fit Google's label rules. Gemini image generation now forwards labels at all.
This commit is contained in:
mateo-berri
parent
83ca14b392
commit
ae73db3855
9 changed files with 791 additions and 22 deletions
|
|
@ -216,6 +216,7 @@ overwrite_user_with_key_hash: bool = (
|
|||
bedrock_request_metadata_fields: Optional[Sequence[str]] = (
|
||||
None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata`
|
||||
)
|
||||
vertex_request_label_fields: Optional[Sequence[str]] = None
|
||||
store_audit_logs: bool | None = None
|
||||
skip_system_message_in_guardrail: bool = False
|
||||
skip_tool_message_in_guardrail: bool = False
|
||||
|
|
|
|||
|
|
@ -24,6 +24,11 @@ from litellm.llms.gemini.common_utils import (
|
|||
GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO,
|
||||
map_openai_size_to_gemini_image_config,
|
||||
)
|
||||
from litellm.llms.vertex_ai.request_labels import (
|
||||
VERTEX_REQUEST_LABELS_MAX,
|
||||
resolve_vertex_request_labels,
|
||||
vertex_request_label_fields,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
VERTEX_AI_PROVIDER_METADATA_FIELDS,
|
||||
|
|
@ -196,6 +201,9 @@ def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> N
|
|||
hidden_params.pop(field, None)
|
||||
|
||||
|
||||
_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def vertex_request_labels_from_litellm_params(
|
||||
litellm_params: dict | None,
|
||||
) -> dict[str, str] | None:
|
||||
|
|
@ -231,7 +239,8 @@ def pop_vertex_request_labels(
|
|||
"""
|
||||
Resolve labels from optional ``labels`` (Gemini-style) and/or
|
||||
``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
|
||||
(``requester_metadata``). Pops ``labels`` from optional_params when present.
|
||||
(``requester_metadata``), plus the proxy identity ``litellm.vertex_request_label_fields`` allows.
|
||||
Pops ``labels`` from optional_params when present.
|
||||
"""
|
||||
labels: dict[str, str] | None = None
|
||||
if optional_params is not None and "labels" in optional_params:
|
||||
|
|
@ -240,7 +249,13 @@ def pop_vertex_request_labels(
|
|||
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)
|
||||
return labels if labels else None
|
||||
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 dict(resolved) if resolved else None
|
||||
|
||||
|
||||
class VertexAIModelRoute(str, Enum):
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.llms.gemini.common_utils import (
|
|||
get_gemini_image_web_search_requests,
|
||||
map_gemini_image_tools_params,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size, pop_vertex_request_labels
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -255,12 +255,13 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
"generationConfig": generation_config,
|
||||
}
|
||||
|
||||
if tools := optional_params.get("tools"):
|
||||
request_body["tools"] = tools
|
||||
if tool_config := optional_params.get("toolConfig"):
|
||||
request_body["toolConfig"] = tool_config
|
||||
|
||||
return request_body
|
||||
params: Final = _JSON_OBJECT.validate_python(optional_params)
|
||||
optional_fields: Final = {
|
||||
"tools": params.get("tools"),
|
||||
"toolConfig": params.get("toolConfig"),
|
||||
"labels": pop_vertex_request_labels(optional_params, litellm_params),
|
||||
}
|
||||
return {**request_body, **{key: value for key, value in optional_fields.items() if value}}
|
||||
|
||||
def _transform_image_usage(self, usage: dict) -> ImageUsage:
|
||||
input_tokens_details: Final = ImageUsageInputTokensDetails(
|
||||
|
|
|
|||
158
litellm/llms/vertex_ai/request_labels.py
Normal file
158
litellm/llms/vertex_ai/request_labels.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""
|
||||
Vertex AI forwards request ``labels`` to Cloud Billing, so proxy identity becomes labels only when the operator sets
|
||||
``litellm.vertex_request_label_fields``. Identity resolves first and the ``user_api_key_`` prefix is reserved, so a
|
||||
caller can neither evict nor forge it. Identity and ``spend_logs_metadata`` text is lowercased, every character Google's
|
||||
label rules disallow becomes ``_``, and it is cut to 63 characters, so an alias like ``AI Platform`` still bills under a
|
||||
compliant ``ai_platform`` label. Caller labels pass through verbatim, as they do with the setting off.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import unicodedata
|
||||
from collections.abc import Mapping
|
||||
from itertools import chain
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
|
||||
VERTEX_REQUEST_LABELS_MAX: Final = 64
|
||||
VERTEX_PARTNER_REQUEST_LABELS_MAX: Final = 32
|
||||
VERTEX_REQUEST_LABEL_MAX_CHARS: Final = 63
|
||||
VERTEX_REQUEST_LABEL_IDENTITY_PREFIX: Final = "user_api_key_"
|
||||
VERTEX_REQUEST_LABEL_CLIENT_FIELD: Final = "spend_logs_metadata"
|
||||
VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels"
|
||||
|
||||
_METADATA_PARAM_NAMES: Final = ("metadata", "litellm_metadata")
|
||||
_KEY_START_CATEGORIES: Final = frozenset(("Ll", "Lo"))
|
||||
_LABEL_CHARACTER_CATEGORIES: Final = frozenset(("Ll", "Lo", "Nd", "Nl", "No"))
|
||||
_ANY_MAPPING: Final = TypeAdapter(Mapping[object, object])
|
||||
_ANY_SEQUENCE: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
Label: TypeAlias = tuple[str, str]
|
||||
|
||||
|
||||
def vertex_request_label_fields() -> tuple[str, ...]:
|
||||
configured: Final[object] = litellm.vertex_request_label_fields
|
||||
if not isinstance(configured, (list, tuple)):
|
||||
return ()
|
||||
return tuple(field for field in _ANY_SEQUENCE.validate_python(configured) if isinstance(field, str))
|
||||
|
||||
|
||||
def _is_label_character(character: str) -> bool:
|
||||
return character in "_-" or unicodedata.category(character) in _LABEL_CHARACTER_CATEGORIES
|
||||
|
||||
|
||||
def _normalized(text: str) -> str:
|
||||
compliant: Final = (character if _is_label_character(character) else "_" for character in text.lower())
|
||||
return "".join(compliant)[:VERTEX_REQUEST_LABEL_MAX_CHARS]
|
||||
|
||||
|
||||
def _normalized_label(key: str, value: str) -> Label | None:
|
||||
normalized_key: Final = _normalized(key)
|
||||
if not normalized_key or unicodedata.category(normalized_key[0]) not in _KEY_START_CATEGORIES:
|
||||
return None
|
||||
return normalized_key, _normalized(value)
|
||||
|
||||
|
||||
def _is_reserved(key: str) -> bool:
|
||||
return key.strip().lower().startswith(VERTEX_REQUEST_LABEL_IDENTITY_PREFIX)
|
||||
|
||||
|
||||
def _metadata_sources(litellm_params: Mapping[str, object]) -> tuple[Mapping[object, object], ...]:
|
||||
candidates: Final = (litellm_params.get(name) for name in _METADATA_PARAM_NAMES)
|
||||
return tuple(_ANY_MAPPING.validate_python(source) for source in candidates if isinstance(source, Mapping))
|
||||
|
||||
|
||||
def _text_labels(source: object) -> tuple[Label, ...]:
|
||||
if not isinstance(source, Mapping):
|
||||
return ()
|
||||
pairs: Final = _ANY_MAPPING.validate_python(source).items()
|
||||
return tuple((key, value) for key, value in pairs if isinstance(key, str) and isinstance(value, str))
|
||||
|
||||
|
||||
def _first_text(sources: tuple[Mapping[object, object], ...], field: str) -> str | None:
|
||||
return next((value for source in sources if isinstance(value := source.get(field), str)), None)
|
||||
|
||||
|
||||
def _identity_label(sources: tuple[Mapping[object, object], ...], field: str) -> Label | None:
|
||||
value: Final = _first_text(sources, field)
|
||||
if value is None or not _is_reserved(field):
|
||||
return None
|
||||
return _normalized_label(field, value)
|
||||
|
||||
|
||||
def _identity_labels(
|
||||
sources: tuple[Mapping[object, object], ...], allowed_fields: tuple[str, ...]
|
||||
) -> tuple[Label, ...]:
|
||||
candidates: Final = (_identity_label(sources, field) for field in allowed_fields)
|
||||
return tuple(label for label in candidates if label is not None)
|
||||
|
||||
|
||||
def _spend_logs_labels(sources: tuple[Mapping[object, object], ...]) -> tuple[Label, ...]:
|
||||
pairs: Final = chain.from_iterable(
|
||||
_text_labels(source.get(VERTEX_REQUEST_LABEL_CLIENT_FIELD)) for source in sources
|
||||
)
|
||||
candidates: Final = (_normalized_label(key, value) for key, value in pairs)
|
||||
return tuple(label for label in candidates if label is not None)
|
||||
|
||||
|
||||
def _first_occurrences(labels: tuple[Label, ...], max_labels: int) -> tuple[Label, ...]:
|
||||
first_value: Final = dict(reversed(labels))
|
||||
keys: Final = tuple(dict.fromkeys(key for key, _ in labels))
|
||||
return tuple((key, first_value[key]) for key in keys[:max_labels])
|
||||
|
||||
|
||||
def resolve_vertex_request_labels(
|
||||
allowed_fields: tuple[str, ...],
|
||||
caller_labels: object,
|
||||
litellm_params: Mapping[str, object],
|
||||
max_labels: int,
|
||||
) -> tuple[Label, ...]:
|
||||
"""
|
||||
The labels to send for a Vertex request once ``allowed_fields`` (``vertex_request_label_fields()``) is non-empty.
|
||||
Identity comes first, then the caller's string labels minus the reserved prefix, then ``spend_logs_metadata``
|
||||
pairs when allowed, each key keeping its first value, up to ``max_labels``.
|
||||
"""
|
||||
sources: Final = _metadata_sources(litellm_params)
|
||||
spend_logs: Final = _spend_logs_labels(sources) if VERTEX_REQUEST_LABEL_CLIENT_FIELD in allowed_fields else ()
|
||||
client: Final = tuple(label for label in (*_text_labels(caller_labels), *spend_logs) if not _is_reserved(label[0]))
|
||||
return _first_occurrences((*_identity_labels(sources, allowed_fields), *client), max_labels)
|
||||
|
||||
|
||||
def is_vertex_labels_header(name: str) -> bool:
|
||||
return name.lower() == VERTEX_LABELS_HEADER.lower()
|
||||
|
||||
|
||||
def _decoded_labels_header(encoded: object) -> object:
|
||||
if not isinstance(encoded, str):
|
||||
return None
|
||||
try:
|
||||
return _ANY_MAPPING.validate_json(base64.b64decode(encoded, validate=True))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def vertex_partner_labels_header(
|
||||
allowed_fields: tuple[str, ...],
|
||||
caller_labels: object,
|
||||
outbound_headers: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> str | None:
|
||||
"""
|
||||
The ``X-Vertex-AI-Labels`` value for a partner model call once ``allowed_fields`` is non-empty. A labels header the
|
||||
caller already set stands in for ``caller_labels``, as it does with the setting off, and the result replaces it.
|
||||
"""
|
||||
caller_header: Final = next(
|
||||
(value for name, value in outbound_headers.items() if is_vertex_labels_header(name)), None
|
||||
)
|
||||
resolved: Final = resolve_vertex_request_labels(
|
||||
allowed_fields,
|
||||
caller_labels if caller_header is None else _decoded_labels_header(caller_header),
|
||||
litellm_params,
|
||||
VERTEX_PARTNER_REQUEST_LABELS_MAX,
|
||||
)
|
||||
return base64.b64encode(json.dumps(dict(resolved)).encode()).decode() if resolved else None
|
||||
|
|
@ -15,9 +15,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
vertex_request_labels_from_litellm_params,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.rerank import (
|
||||
|
|
@ -163,7 +161,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
|
|||
# When return_documents is False, we want to ignore record details (return only IDs)
|
||||
request_data["ignoreRecordDetailsInResponse"] = not return_documents
|
||||
|
||||
user_labels: Final = vertex_request_labels_from_litellm_params(litellm_params)
|
||||
user_labels: Final = pop_vertex_request_labels(None, litellm_params)
|
||||
if user_labels:
|
||||
request_data["userLabels"] = user_labels
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,12 @@ from litellm.llms.anthropic.messages.native.transformation import (
|
|||
AnthropicMessagesConfig,
|
||||
messages_carry_output_config,
|
||||
)
|
||||
from litellm.llms.vertex_ai.request_labels import (
|
||||
VERTEX_LABELS_HEADER,
|
||||
is_vertex_labels_header,
|
||||
vertex_partner_labels_header,
|
||||
vertex_request_label_fields,
|
||||
)
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_HOSTED_TOOLS,
|
||||
|
|
@ -18,6 +24,18 @@ from ....vertex_llm_base import VertexBase
|
|||
from ..output_params_utils import sanitize_vertex_anthropic_output_params
|
||||
|
||||
|
||||
def _labels_header_replacement(
|
||||
request_headers: Mapping[str, str], litellm_params: Mapping[str, object]
|
||||
) -> tuple[tuple[str, ...], str | None]:
|
||||
allowed_fields: Final = vertex_request_label_fields()
|
||||
if not allowed_fields:
|
||||
return (), None
|
||||
return (
|
||||
tuple(filter(is_vertex_labels_header, request_headers)),
|
||||
vertex_partner_labels_header(allowed_fields, None, request_headers, litellm_params),
|
||||
)
|
||||
|
||||
|
||||
class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase):
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
|
|
@ -67,6 +85,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
project_id,
|
||||
vertex_ai_project,
|
||||
vertex_ai_location,
|
||||
litellm_params,
|
||||
)
|
||||
|
||||
async def avalidate_anthropic_messages_environment(
|
||||
|
|
@ -95,6 +114,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
project_id,
|
||||
vertex_ai_project,
|
||||
vertex_ai_location,
|
||||
litellm_params,
|
||||
)
|
||||
|
||||
def _finalize_vertex_messages_environment(
|
||||
|
|
@ -108,6 +128,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
project_id: str,
|
||||
vertex_ai_project: str | None,
|
||||
vertex_ai_location: str | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers
|
||||
# Router shallow-copies litellm_params, so headers may be shared between requests.
|
||||
request_headers: Final = dict(headers)
|
||||
|
|
@ -177,6 +198,11 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
if beta_values:
|
||||
request_headers["anthropic-beta"] = ",".join(beta_values)
|
||||
|
||||
replaced_headers, labels_header = _labels_header_replacement(request_headers, litellm_params)
|
||||
for name in replaced_headers:
|
||||
del request_headers[name]
|
||||
if labels_header is not None:
|
||||
request_headers[VERTEX_LABELS_HEADER] = labels_header
|
||||
return request_headers, api_base
|
||||
|
||||
def get_complete_url(
|
||||
|
|
|
|||
|
|
@ -11,6 +11,12 @@ from pydantic import TypeAdapter
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.llms.vertex_ai.request_labels import (
|
||||
VERTEX_LABELS_HEADER,
|
||||
is_vertex_labels_header,
|
||||
vertex_partner_labels_header,
|
||||
vertex_request_label_fields,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
|
@ -21,12 +27,11 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
_NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels"
|
||||
_LABELS: Final = TypeAdapter(Mapping[object, object])
|
||||
|
||||
|
||||
def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None:
|
||||
if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names):
|
||||
if not isinstance(labels, dict) or any(map(is_vertex_labels_header, header_names)):
|
||||
return None
|
||||
string_labels: Final = {
|
||||
key: value
|
||||
|
|
@ -118,11 +123,27 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
headers: dict,
|
||||
) -> dict:
|
||||
params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params)
|
||||
labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers))
|
||||
outbound_headers: Final = _NON_DEFAULT_PARAMS.validate_python(headers)
|
||||
allowed_fields: Final = vertex_request_label_fields()
|
||||
labels_header, replaced_headers = (
|
||||
(
|
||||
vertex_partner_labels_header(
|
||||
allowed_fields,
|
||||
params.get("labels"),
|
||||
outbound_headers,
|
||||
_NON_DEFAULT_PARAMS.validate_python(litellm_params),
|
||||
),
|
||||
tuple(filter(is_vertex_labels_header, outbound_headers)),
|
||||
)
|
||||
if allowed_fields
|
||||
else (_vertex_labels_header(params.get("labels"), outbound_headers), ())
|
||||
)
|
||||
if "labels" in params:
|
||||
del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body
|
||||
for name in replaced_headers:
|
||||
del headers[name] # rebind-ok: the resolved identity labels replace a caller-set labels header
|
||||
if labels_header is not None:
|
||||
headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation
|
||||
headers[VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation
|
||||
data: Final = super().transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
@ -158,11 +179,8 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
|
||||
extra_headers: Final = optional_params.get("extra_headers") or {}
|
||||
anthropic_beta_value: Final = extra_headers.get("anthropic-beta", "")
|
||||
if isinstance(anthropic_beta_value, str) and anthropic_beta_value:
|
||||
for beta in anthropic_beta_value.split(","):
|
||||
beta = beta.strip()
|
||||
if beta:
|
||||
beta_set.add(beta)
|
||||
if isinstance(anthropic_beta_value, str):
|
||||
beta_set.update(filter(None, (beta.strip() for beta in anthropic_beta_value.split(","))))
|
||||
elif isinstance(anthropic_beta_value, list):
|
||||
beta_set.update(anthropic_beta_value)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,261 @@
|
|||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_PROJECT: Final = "scripted-project"
|
||||
_LOCATION: Final = "us-central1"
|
||||
_CHAT_BACKEND: Final = "gemini-3.7-flash"
|
||||
_EMBED_BACKEND: Final = "text-embedding-005"
|
||||
_CLAUDE_BACKEND: Final = "claude-haiku-5-5"
|
||||
_MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models"
|
||||
_CHAT_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:generateContent"
|
||||
_STREAM_TARGET: Final = f"{_MODELS_PATH}/{_CHAT_BACKEND}:streamGenerateContent?alt=sse"
|
||||
_EMBED_TARGET: Final = f"{_MODELS_PATH}/{_EMBED_BACKEND}:predict"
|
||||
_CLAUDE_TARGET: Final = (
|
||||
f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_CLAUDE_BACKEND}:rawPredict"
|
||||
)
|
||||
_LABELS_HEADER: Final = "x-vertex-ai-labels"
|
||||
_REPLY_TEXT: Final = "labels ok"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_GEMINI_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": _REPLY_TEXT}]}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5},
|
||||
"modelVersion": _CHAT_BACKEND,
|
||||
}
|
||||
_EMBED_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"predictions": [{"embeddings": {"values": [0.5, 0.25], "statistics": {"token_count": 2}}}]
|
||||
}
|
||||
_CLAUDE_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"id": "msg_labels",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": _CLAUDE_BACKEND,
|
||||
"content": [{"type": "text", "text": _REPLY_TEXT}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
def _vertex_peer(request: Request) -> Reply:
|
||||
if request.target == "/_oauth/token":
|
||||
return Reply(
|
||||
body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode()
|
||||
)
|
||||
if request.target == _STREAM_TARGET:
|
||||
return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_GEMINI_REPLY)}\n\n".encode(),))
|
||||
if request.target == _CHAT_TARGET:
|
||||
return Reply(body=json.dumps(_GEMINI_REPLY).encode())
|
||||
if request.target == _EMBED_TARGET:
|
||||
return Reply(body=json.dumps(_EMBED_REPLY).encode())
|
||||
if request.target == _CLAUDE_TARGET:
|
||||
return Reply(body=json.dumps(_CLAUDE_REPLY).encode())
|
||||
return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode())
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, wire: Wire, backend: str) -> str:
|
||||
return scenario.model(
|
||||
model=f"vertex_ai/{backend}",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
vertex_project=_PROJECT,
|
||||
vertex_location=_LOCATION,
|
||||
vertex_credentials=service_account_json(_PROJECT, wire.url),
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
chat_model: str
|
||||
embed_model: str
|
||||
claude_model: str
|
||||
key: str
|
||||
identity: Mapping[str, str]
|
||||
uppercase_team_key: str
|
||||
uppercase_team_identity: Mapping[str, str]
|
||||
|
||||
def request_sent(self, marker: str) -> Request:
|
||||
(request,) = tuple(request for request in self.wire.drain() if marker.encode() in request.body)
|
||||
return request
|
||||
|
||||
def labels_sent(self, marker: str) -> JsonValue:
|
||||
return _JSON_OBJECT.validate_json(self.request_sent(marker).body).get("labels")
|
||||
|
||||
def header_labels_sent(self, marker: str) -> JsonValue:
|
||||
headers: Final = self.request_sent(marker).headers
|
||||
(encoded,) = tuple(value for name, value in headers.items() if name.lower() == _LABELS_HEADER)
|
||||
return _JSON_OBJECT.validate_json(base64.b64decode(encoded))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
suffix: Final = uuid.uuid4().hex[:12]
|
||||
directory: Final = tmp_path_factory.mktemp("vertex-request-labels")
|
||||
base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config: Final = {
|
||||
**base_config,
|
||||
"litellm_settings": {
|
||||
**base_config.get("litellm_settings", {}),
|
||||
"vertex_request_label_fields": [
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_alias",
|
||||
"user_api_key_org_alias",
|
||||
"user_api_key_project_alias",
|
||||
"spend_logs_metadata",
|
||||
],
|
||||
},
|
||||
}
|
||||
config_path: Final = directory / "vertex-request-labels.yaml"
|
||||
config_path.write_text(yaml.safe_dump(config))
|
||||
with (
|
||||
wire_server(_vertex_peer) as wire,
|
||||
owned_proxy(session_gateway, directory, {}, config=config_path) as proxy,
|
||||
proxy.scenario() as scenario,
|
||||
):
|
||||
chat_model: Final = _deployment(scenario, wire, _CHAT_BACKEND)
|
||||
embed_model: Final = _deployment(scenario, wire, _EMBED_BACKEND)
|
||||
claude_model: Final = _deployment(scenario, wire, _CLAUDE_BACKEND)
|
||||
org_alias: Final = f"ml-org-{suffix}"
|
||||
team_alias: Final = f"ai-platform-{suffix}"
|
||||
project_alias: Final = f"support-bot-{suffix}"
|
||||
key_alias: Final = f"chargeback-{suffix}"
|
||||
uppercase_key_alias: Final = f"support-{suffix}"
|
||||
team: Final = scenario.team(
|
||||
team_alias=team_alias,
|
||||
organization_id=scenario.organization(
|
||||
organization_alias=org_alias, models=[chat_model, embed_model, claude_model]
|
||||
),
|
||||
)
|
||||
yield Rig(
|
||||
proxy=proxy,
|
||||
wire=wire,
|
||||
chat_model=chat_model,
|
||||
embed_model=embed_model,
|
||||
claude_model=claude_model,
|
||||
key=scenario.key(
|
||||
key_alias=key_alias,
|
||||
team_id=team,
|
||||
project_id=scenario.project(team, project_alias=project_alias),
|
||||
metadata={"spend_logs_metadata": {"cost-center": "cc-1234", "Owner": "Platform Team"}},
|
||||
),
|
||||
identity={
|
||||
"user_api_key_alias": key_alias,
|
||||
"user_api_key_team_alias": team_alias,
|
||||
"user_api_key_org_alias": org_alias,
|
||||
"user_api_key_project_alias": project_alias,
|
||||
},
|
||||
uppercase_team_key=scenario.key(
|
||||
key_alias=uppercase_key_alias,
|
||||
team_id=scenario.team(team_alias=f"AI Platform {suffix}"),
|
||||
),
|
||||
uppercase_team_identity={
|
||||
"user_api_key_alias": uppercase_key_alias,
|
||||
"user_api_key_team_alias": f"ai_platform_{suffix}",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _chat(rig: Rig, key: str, marker: str, stream: bool = False, **extra: JsonValue) -> JsonValue:
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": rig.chat_model, "messages": [{"role": "user", "content": marker}], "stream": stream, **extra},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _REPLY_TEXT in response.text, response.text
|
||||
return rig.labels_sent(marker)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
|
||||
def test_chat_labels_carry_the_key_team_org_and_project_identity_with_no_client_labels(rig: Rig, stream: bool) -> None:
|
||||
labels: Final = _chat(rig, rig.key, f"identity-{uuid.uuid4().hex}", stream)
|
||||
|
||||
assert labels == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"}, labels
|
||||
|
||||
|
||||
def test_caller_cannot_forge_identity_and_its_labels_rank_above_spend_logs_metadata(rig: Rig) -> None:
|
||||
labels: Final = _chat(
|
||||
rig,
|
||||
rig.key,
|
||||
f"forged-{uuid.uuid4().hex}",
|
||||
labels={"user_api_key_team_alias": "forged", "cost-center": "caller", "ticket": "t-42"},
|
||||
)
|
||||
|
||||
assert labels == {**rig.identity, "cost-center": "caller", "ticket": "t-42", "owner": "platform_team"}, labels
|
||||
|
||||
|
||||
def test_request_metadata_labels_still_reach_vertex_under_the_identity(rig: Rig) -> None:
|
||||
labels: Final = _chat(rig, rig.key, f"requester-{uuid.uuid4().hex}", metadata={"env": "prod"})
|
||||
|
||||
assert labels == {**rig.identity, "env": "prod", "cost-center": "cc-1234", "owner": "platform_team"}, labels
|
||||
|
||||
|
||||
def test_team_alias_outside_googles_label_rules_is_rewritten_to_a_compliant_label(rig: Rig) -> None:
|
||||
labels: Final = _chat(rig, rig.uppercase_team_key, f"uppercase-{uuid.uuid4().hex}")
|
||||
|
||||
assert labels == rig.uppercase_team_identity, labels
|
||||
|
||||
|
||||
def test_embedding_labels_carry_the_caller_identity(rig: Rig) -> None:
|
||||
marker: Final = f"embed-{uuid.uuid4().hex}"
|
||||
response: Final = rig.proxy.request(
|
||||
"POST", "/v1/embeddings", {"model": rig.embed_model, "input": [marker]}, key=rig.key
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert rig.labels_sent(marker) == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"}
|
||||
|
||||
|
||||
def test_claude_chat_labels_header_carries_the_identity_and_drops_a_forged_label(rig: Rig) -> None:
|
||||
marker: Final = f"claude-chat-{uuid.uuid4().hex}"
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": rig.claude_model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"labels": {"user_api_key_team_alias": "forged", "ticket": "t-42"},
|
||||
},
|
||||
key=rig.key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _REPLY_TEXT in response.text, response.text
|
||||
assert rig.header_labels_sent(marker) == {
|
||||
**rig.identity,
|
||||
"ticket": "t-42",
|
||||
"cost-center": "cc-1234",
|
||||
"owner": "platform_team",
|
||||
}
|
||||
|
||||
|
||||
def test_claude_messages_labels_header_carries_the_identity(rig: Rig) -> None:
|
||||
marker: Final = f"claude-messages-{uuid.uuid4().hex}"
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{"model": rig.claude_model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]},
|
||||
key=rig.key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _REPLY_TEXT in response.text, response.text
|
||||
assert rig.header_labels_sent(marker) == {**rig.identity, "cost-center": "cc-1234", "owner": "platform_team"}
|
||||
291
tests/unit/llms/vertex_ai/test_request_labels.py
Normal file
291
tests/unit/llms/vertex_ai/test_request_labels.py
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.gemini.transformation import transform_request_body
|
||||
from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig
|
||||
from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import VertexAIImagenImageGenerationConfig
|
||||
from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
|
||||
from litellm.llms.vertex_ai.vertex_embeddings.transformation import VertexAITextEmbeddingConfig
|
||||
|
||||
FIELDS: Final = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"]
|
||||
IDENTITY: Final = {"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}
|
||||
LABELS_HEADER: Final = "x-vertex-ai-labels"
|
||||
|
||||
LabelSource = Callable[[dict[str, JsonValue], dict[str, JsonValue]], object]
|
||||
|
||||
|
||||
def gemini_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object:
|
||||
body: Final = transform_request_body(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
model="gemini-3.8-flash",
|
||||
optional_params=dict(optional_params),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params=dict(litellm_params),
|
||||
cached_content=None,
|
||||
)
|
||||
return body.get("labels")
|
||||
|
||||
|
||||
def embedding_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object:
|
||||
request: Final = VertexAITextEmbeddingConfig().transform_openai_request_to_vertex_embedding_request(
|
||||
input=["hi"], optional_params=dict(optional_params), model="text-embedding-005", litellm_params=litellm_params
|
||||
)
|
||||
return request.get("labels")
|
||||
|
||||
|
||||
def imagen_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object:
|
||||
request: Final = VertexAIImagenImageGenerationConfig().transform_image_generation_request(
|
||||
model="imagen-4.0-generate-001",
|
||||
prompt="a cat",
|
||||
optional_params=dict(optional_params),
|
||||
litellm_params=dict(litellm_params),
|
||||
headers={},
|
||||
)
|
||||
return request.get("labels")
|
||||
|
||||
|
||||
def gemini_image_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object:
|
||||
request: Final = VertexAIGeminiImageGenerationConfig().transform_image_generation_request(
|
||||
model="gemini-3.1-flash-image",
|
||||
prompt="a cat",
|
||||
optional_params=dict(optional_params),
|
||||
litellm_params=dict(litellm_params),
|
||||
headers={},
|
||||
)
|
||||
return request.get("labels")
|
||||
|
||||
|
||||
def rerank_labels(litellm_params: dict[str, JsonValue], _: dict[str, JsonValue]) -> object:
|
||||
request: Final = VertexAIRerankConfig().transform_rerank_request(
|
||||
model="semantic-ranker-default@latest",
|
||||
optional_rerank_params={"query": "q", "documents": ["a", "b"]},
|
||||
headers={},
|
||||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
return request.get("userLabels")
|
||||
|
||||
|
||||
def encoded(labels: Mapping[str, str]) -> str:
|
||||
return base64.b64encode(json.dumps(dict(labels)).encode()).decode()
|
||||
|
||||
|
||||
def claude_headers(
|
||||
litellm_params: dict[str, JsonValue],
|
||||
optional_params: dict[str, JsonValue],
|
||||
headers: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> Mapping[str, str]:
|
||||
outbound: Final = dict(headers)
|
||||
VertexAIAnthropicConfig().transform_request(
|
||||
model="claude-haiku-5-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params=dict(optional_params),
|
||||
litellm_params=dict(litellm_params),
|
||||
headers=outbound,
|
||||
)
|
||||
return outbound
|
||||
|
||||
|
||||
def labels_header_names(headers: Mapping[str, str]) -> list[str]:
|
||||
return [name for name in headers if name.lower() == LABELS_HEADER]
|
||||
|
||||
|
||||
def decoded_labels(headers: Mapping[str, str]) -> object:
|
||||
names: Final = labels_header_names(headers)
|
||||
return json.loads(base64.b64decode(headers[names[0]])) if names else None
|
||||
|
||||
|
||||
def claude_labels(litellm_params: dict[str, JsonValue], optional_params: dict[str, JsonValue]) -> object:
|
||||
return decoded_labels(claude_headers(litellm_params, optional_params))
|
||||
|
||||
|
||||
LABELLED_SURFACES: Final = [gemini_labels, embedding_labels, imagen_labels, gemini_image_labels]
|
||||
REQUESTER_METADATA_SURFACES: Final = [*LABELLED_SURFACES, rerank_labels]
|
||||
ALL_SURFACES: Final = [*REQUESTER_METADATA_SURFACES, claude_labels]
|
||||
CALLER_LABEL_SURFACES: Final = [*LABELLED_SURFACES, claude_labels]
|
||||
|
||||
|
||||
def metadata(key: str = "metadata", **fields: JsonValue) -> dict[str, JsonValue]:
|
||||
return {key: dict(fields)}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [None, []])
|
||||
@pytest.mark.parametrize("surface", REQUESTER_METADATA_SURFACES)
|
||||
def test_setting_off_sends_only_the_callers_labels(
|
||||
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={"user_api_key_team_alias": "caller-chosen"},
|
||||
**IDENTITY,
|
||||
)
|
||||
|
||||
assert surface(params, {}) == {"user_api_key_team_alias": "caller-chosen"}
|
||||
assert surface(metadata(**IDENTITY), {}) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [None, []])
|
||||
def test_setting_off_leaves_claude_labels_as_the_caller_sent_them(setting: object, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", setting)
|
||||
params: Final = metadata(spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY)
|
||||
caller_header: Final = {"X-Vertex-AI-Labels": encoded({"user_api_key_team_alias": "Caller Chosen"})}
|
||||
|
||||
assert claude_labels(params, {"labels": {"user_api_key_team_alias": "Caller Chosen"}}) == {
|
||||
"user_api_key_team_alias": "Caller Chosen"
|
||||
}
|
||||
assert claude_labels(params, {}) is None
|
||||
assert claude_headers(params, {"labels": {"app": "bot"}}, caller_header) == caller_header
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
@pytest.mark.parametrize("surface", ALL_SURFACES)
|
||||
def test_identity_and_spend_logs_metadata_label_every_vertex_surface(
|
||||
metadata_key: str, surface: LabelSource, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
params: Final = metadata(metadata_key, spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY)
|
||||
|
||||
assert surface(params, {}) == {**IDENTITY, "cost-center": "cc-1"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field_order", [FIELDS, list(reversed(FIELDS))])
|
||||
@pytest.mark.parametrize(
|
||||
("surface", "cap"),
|
||||
[
|
||||
(gemini_labels, 64),
|
||||
(embedding_labels, 64),
|
||||
(imagen_labels, 64),
|
||||
(gemini_image_labels, 64),
|
||||
(claude_labels, 32),
|
||||
],
|
||||
)
|
||||
def test_identity_survives_a_caller_filling_every_label_slot(
|
||||
field_order: list[str], surface: LabelSource, cap: int, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""64 labels per call for Google models, 32 for partner models,
|
||||
https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls (2026-10-28)"""
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", field_order)
|
||||
caller_labels: Final = {f"caller-{index:02d}": "v" for index in range(cap)}
|
||||
|
||||
labels: Final = surface(metadata(**IDENTITY), {"labels": caller_labels})
|
||||
|
||||
assert isinstance(labels, dict)
|
||||
assert len(labels) == cap
|
||||
assert {key: labels[key] for key in IDENTITY} == IDENTITY
|
||||
assert all(labels[key] == "v" for key in tuple(caller_labels)[: cap - len(IDENTITY)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", CALLER_LABEL_SURFACES)
|
||||
def test_caller_cannot_write_or_override_the_identity_prefix_in_any_case(
|
||||
surface: LabelSource, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
forged: Final = {
|
||||
"user_api_key_team_alias": "forged",
|
||||
"User_API_Key_Team_Alias": "forged",
|
||||
" user_api_key_user_id": "forged",
|
||||
"app": "bot",
|
||||
}
|
||||
|
||||
assert surface(metadata(**IDENTITY), {"labels": forged}) == {**IDENTITY, "app": "bot"}
|
||||
assert surface(metadata(user_api_key_alias="chargeback-key"), {"labels": forged}) == {
|
||||
"user_api_key_alias": "chargeback-key",
|
||||
"app": "bot",
|
||||
}
|
||||
assert surface(metadata(spend_logs_metadata={"USER_API_KEY_TEAM_ALIAS": "forged"}), {}) is None
|
||||
|
||||
|
||||
def test_a_caller_set_claude_labels_header_cannot_forge_identity_and_is_replaced(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
caller_header: Final = {"x-Vertex-AI-labels": encoded({"user_api_key_team_alias": "forged", "app": "bot"})}
|
||||
|
||||
sent: Final = claude_headers(metadata(**IDENTITY), {"labels": {"ignored": "x"}}, caller_header)
|
||||
undecodable: Final = claude_headers(metadata(**IDENTITY), {}, {LABELS_HEADER: "%%"})
|
||||
|
||||
assert len(labels_header_names(sent)) == 1
|
||||
assert decoded_labels(sent) == {**IDENTITY, "app": "bot"}
|
||||
assert len(labels_header_names(undecodable)) == 1
|
||||
assert decoded_labels(undecodable) == IDENTITY
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ALL_SURFACES)
|
||||
def test_identity_and_spend_logs_text_is_rewritten_to_labels_google_accepts(
|
||||
surface: LabelSource, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Key and value rules from https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls (2026-10-28)"""
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
spend_logs: Final = {
|
||||
"Owner": "Platform Team",
|
||||
"1st": "x",
|
||||
"_private": "x",
|
||||
"": "x",
|
||||
"dotted": "a.b",
|
||||
"too-long": "V" * 64,
|
||||
"K" * 64: "",
|
||||
"équipe": "Données-2",
|
||||
"部門": "営業",
|
||||
"numerals": "ⅻ½",
|
||||
}
|
||||
params: Final = metadata(
|
||||
user_api_key_alias="Chargeback Key",
|
||||
user_api_key_team_alias="AI Platform " + "x" * 60,
|
||||
spend_logs_metadata=spend_logs,
|
||||
)
|
||||
|
||||
assert surface(params, {}) == {
|
||||
"user_api_key_alias": "chargeback_key",
|
||||
"user_api_key_team_alias": ("ai_platform_" + "x" * 60)[:63],
|
||||
"owner": "platform_team",
|
||||
"dotted": "a_b",
|
||||
"too-long": "v" * 63,
|
||||
"k" * 63: "",
|
||||
"équipe": "données-2",
|
||||
"部門": "営業",
|
||||
"numerals": "ⅻ½",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", CALLER_LABEL_SURFACES)
|
||||
def test_caller_labels_keep_todays_pass_through_and_win_over_spend_logs_metadata(
|
||||
surface: LabelSource, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
params: Final = metadata(spend_logs_metadata={"cost-center": "cc-admin", "env": "prod"}, **IDENTITY)
|
||||
|
||||
assert surface(params, {"labels": {"cost-center": "cc-caller", "Mixed": "Case", 1: "dropped"}}) == {
|
||||
**IDENTITY,
|
||||
"cost-center": "cc-caller",
|
||||
"Mixed": "Case",
|
||||
"env": "prod",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ALL_SURFACES)
|
||||
def test_only_allow_listed_identity_fields_and_opted_in_spend_logs_metadata_become_labels(
|
||||
surface: LabelSource, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", ["user_api_key_team_alias", "trace_id"])
|
||||
params: Final = metadata(trace_id="t-1", spend_logs_metadata={"cost-center": "cc-1"}, **IDENTITY)
|
||||
|
||||
assert surface(params, {}) == {"user_api_key_team_alias": "ai-platform"}
|
||||
|
||||
|
||||
def test_metadata_wins_over_litellm_metadata_for_the_same_identity_field(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "vertex_request_label_fields", FIELDS)
|
||||
params: Final[Mapping[str, JsonValue]] = {
|
||||
"metadata": {"user_api_key_team_alias": "from-metadata"},
|
||||
"litellm_metadata": {"user_api_key_team_alias": "from-litellm-metadata", "user_api_key_alias": "key-a"},
|
||||
}
|
||||
|
||||
assert gemini_labels(dict(params), {}) == {
|
||||
"user_api_key_team_alias": "from-metadata",
|
||||
"user_api_key_alias": "key-a",
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue