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.
(cherry picked from commit ae73db3855)
This commit is contained in:
mateo-berri
parent
60ce0a0611
commit
c84e2e731c
9 changed files with 790 additions and 23 deletions
|
|
@ -237,6 +237,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
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs
|
||||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
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,
|
||||
|
|
@ -164,6 +169,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:
|
||||
|
|
@ -199,7 +207,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:
|
||||
|
|
@ -208,7 +217,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):
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ from litellm.llms.gemini.common_utils import (
|
|||
map_gemini_image_tools_params,
|
||||
)
|
||||
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, 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
|
||||
|
|
@ -253,12 +253,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_generation_response(
|
||||
self,
|
||||
|
|
|
|||
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
|
||||
|
|
@ -14,9 +14,7 @@ import httpx
|
|||
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 (
|
||||
|
|
@ -151,7 +149,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
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
|
|
@ -5,6 +6,12 @@ from litellm.llms.anthropic.pass_through.messages.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,
|
||||
|
|
@ -17,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:
|
||||
|
|
@ -118,7 +137,12 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
if beta_values:
|
||||
headers["anthropic-beta"] = ",".join(beta_values)
|
||||
|
||||
return headers, api_base
|
||||
replaced_headers, labels_header = _labels_header_replacement(headers, litellm_params)
|
||||
kept_headers: Final = {name: value for name, value in headers.items() if name not in replaced_headers}
|
||||
labelled_headers: Final = (
|
||||
kept_headers if labels_header is None else {**kept_headers, VERTEX_LABELS_HEADER: labels_header}
|
||||
)
|
||||
return labelled_headers, api_base
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -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