diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 29253359003..363e5596cb9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,5 +1,7 @@ # What is this? ## Handler file for calling claude-3 on vertex ai +import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -19,6 +21,19 @@ 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): + return None + string_labels: Final = { + key: value + for key, value in _LABELS.validate_python(labels).items() + if isinstance(key, str) and isinstance(value, str) + } + return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None class VertexAIError(Exception): @@ -102,6 +117,12 @@ class VertexAIAnthropicConfig(AnthropicConfig): litellm_params: dict, 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)) + if "labels" in params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + if labels_header is not None: + headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation data: Final = super().transform_request( model=model, messages=messages, diff --git a/tests/integration/providers/test_vertex_partner_labels_wire.py b/tests/integration/providers/test_vertex_partner_labels_wire.py new file mode 100644 index 00000000000..ebdf3f2022e --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_labels_wire.py @@ -0,0 +1,120 @@ +import base64 +import json +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-5-5" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_BACKEND}" +_TARGET: Final = f"{_MODEL_PATH}:rawPredict" +_STREAM_TARGET: Final = f"{_MODEL_PATH}:streamRawPredict?alt=sse" +_LABELS_HEADER: Final = "x-vertex-ai-labels" +_REPLY_TEXT: Final = "labels control" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_LABELS: Final[dict[str, JsonValue]] = {"team": "triage", "env": "prod"} + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_vertex_labels", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_REPLY, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, +) +_SSE: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _EVENTS) + + +def _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=_SSE) + if request.target == _TARGET: + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, api_base), + num_retries=0, + ) + + +def _chat(gateway: Gateway, wire: Wire, scenario: Scenario, stream: bool, **extra: JsonValue) -> Request: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Reply ok"}], + "max_tokens": 16, + "stream": stream, + **extra, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = tuple(request for request in wire.drain() if request.target in (_TARGET, _STREAM_TARGET)) + assert sent.target == (_STREAM_TARGET if stream else _TARGET), sent.target + return sent + + +def _decoded_labels(sent: Request) -> JsonValue: + return json.loads(base64.b64decode(sent.headers[_LABELS_HEADER])) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_regression_vertex_claude_labels_move_from_body_to_header(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, stream, labels=_LABELS) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _decoded_labels(sent) == _LABELS, sent.headers + + +def test_guard_vertex_claude_without_labels_sends_no_labels_header(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, False) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _LABELS_HEADER not in sent.headers, sent.headers + + +def test_regression_vertex_claude_caller_labels_header_wins_over_labels_param(gateway: Gateway) -> None: + own_header: Final = base64.b64encode(json.dumps({"team": "caller"}).encode()).decode() + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat( + gateway, wire, scenario, False, labels=_LABELS, extra_headers={"X-Vertex-AI-Labels": own_header} + ) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert sent.headers[_LABELS_HEADER] == own_header, sent.headers diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 21bda6f0345..1c432ec79da 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -1,4 +1,5 @@ +import base64 import copy import json from collections.abc import Mapping @@ -40,6 +41,33 @@ def test_vertex_ai_anthropic_thinking_param(model, expected_thinking): assert "thinking" not in supported_openai_params +@pytest.mark.parametrize( + ("labels", "headers", "expected_header"), + [ + ({"team": "triage", "env": "prod"}, {}, {"team": "triage", "env": "prod"}), + ({"team": "triage", "cost": 3}, {}, {"team": "triage"}), + ({"cost": 3}, {}, None), + ("team=triage", {}, None), + ({"team": "triage"}, {"x-vertex-ai-labels": "e30="}, {}), + ], + ids=["string_labels", "non_string_values_dropped", "no_string_values", "not_a_mapping", "caller_header_wins"], +) +def test_vertex_anthropic_labels_move_from_body_to_header( + labels: object, headers: dict[str, str], expected_header: dict[str, str] | None +) -> None: + body: Final = VertexAIAnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"max_tokens": 16, "anthropic_version": "vertex-2023-10-16", "labels": labels}, + litellm_params={}, + headers=headers, + ) + sent_labels: Final = next((value for name, value in headers.items() if name.lower() == "x-vertex-ai-labels"), None) + + assert "labels" not in body + assert (None if sent_labels is None else json.loads(base64.b64decode(sent_labels))) == expected_header + + @pytest.mark.parametrize( ("stream", "operation"), [(False, "rawPredict"), (True, "streamRawPredict")],