fix(vertex_ai): send Vertex Claude labels as X-Vertex-AI-Labels header instead of the request body (#48753)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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,

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

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