From 3a15725832fd56f59f8dc83a9108fdd84d8e0947 Mon Sep 17 00:00:00 2001 From: Abhishek Acharya Date: Sun, 25 Oct 2026 02:55:03 +1100 Subject: [PATCH] fix(vertex_ai): send caller headers to Vertex Gemma endpoints VertexAIGemmaModels.completion accepted headers (extra_headers included) but never passed them on, and the sync and async paths built a fixed header dict with only Authorization and Content-Type. Pass the headers through and merge them into the request headers, with the Vertex access token kept as Authorization, like the Model Garden and partner model routes. --- .../vertex_ai/vertex_gemma_models/main.py | 1 + .../vertex_gemma_models/transformation.py | 21 ++-- .../test_vertex_gemma_header_forwarding.py | 116 ++++++++++++++++++ 3 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_header_forwarding.py diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 35cb3929198..6cdb6fe2c3c 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -131,6 +131,7 @@ class VertexAIGemmaModels(VertexBase): timeout=timeout, encoding=encoding, custom_llm_provider="vertex_ai", + headers=headers, ) except Exception as e: diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index ec2fd16f93e..48ac41787f7 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -8,7 +8,7 @@ Handles the custom request/response format: The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI-compatible format. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, cast import httpx @@ -267,6 +267,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): timeout: float | httpx.Timeout | None = None, encoding=None, custom_llm_provider: str = "vertex_ai", + headers: Mapping[str, str] | None = None, ): """ Make completion request to Vertex Gemma endpoint. @@ -287,6 +288,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client=async_client, timeout=timeout, encoding=encoding, + headers=headers, ) else: sync_client = client if isinstance(client, (HTTPHandler, httpx.Client)) else None @@ -303,6 +305,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client=sync_client, timeout=timeout, encoding=encoding, + headers=headers, ) def _sync_completion( @@ -319,6 +322,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: HTTPHandler | httpx.Client | None = None, timeout: float | httpx.Timeout | None = None, encoding: "Tokenizer | None" = None, + headers: Mapping[str, str] | None = None, ): """Synchronous completion request""" from litellm.utils import convert_to_model_response_object @@ -336,9 +340,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Set up headers - headers: Final = { - "Authorization": f"Bearer {api_key}", + request_headers: Final = { "Content-Type": "application/json", + **(headers or {}), + "Authorization": f"Bearer {api_key}", } # Log the request @@ -355,7 +360,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): response: Final = self._sync_post( client=client, api_base=api_base, - headers=headers, + headers=request_headers, request_data=request_data, timeout=timeout, ) @@ -414,6 +419,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: AsyncHTTPHandler | httpx.AsyncClient | None = None, timeout: float | httpx.Timeout | None = None, encoding: "Tokenizer | None" = None, + headers: Mapping[str, str] | None = None, ): """Asynchronous completion request""" from litellm.utils import convert_to_model_response_object @@ -431,9 +437,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Set up headers - headers: Final = { - "Authorization": f"Bearer {api_key}", + request_headers: Final = { "Content-Type": "application/json", + **(headers or {}), + "Authorization": f"Bearer {api_key}", } # Log the request @@ -450,7 +457,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): response: Final = await self._async_post( client=client, api_base=api_base, - headers=headers, + headers=request_headers, request_data=request_data, timeout=timeout, ) diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_header_forwarding.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_header_forwarding.py new file mode 100644 index 00000000000..aa404cf46f9 --- /dev/null +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_header_forwarding.py @@ -0,0 +1,116 @@ +"""Unit tests verifying caller headers forwarding for Vertex AI Gemma models.""" + +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels + +USER_API_BASE = "https://32277599999999999.us-central1-10582012152.prediction.vertexai.goog/v1/projects/PROJECT_ID/locations/us-central1/endpoints/ENDPOINT_ID:predict" + +GEMMA_RESPONSE = { + "deployedModelId": "1207280419999999999", + "model": "projects/993702345710/locations/us-central1/models/gemma-3-12b-it-1222199011122", + "modelDisplayName": "gemma-3-12b-it-1222199011122", + "modelVersionId": "1", + "predictions": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello from Gemma!", + "role": "assistant", + }, + } + ], + "created": 1739818800, + "id": "chatcmpl-test-gemma", + "model": "gemma-3-12b-it-1222199011122", + "object": "chat.completion", + "usage": { + "completion_tokens": 5, + "prompt_tokens": 10, + "total_tokens": 15, + }, + }, +} + + +@pytest.fixture +def gemma_chat_route(respx_mock, monkeypatch): + """The Gemma route with vertexai import, access token minting and network mocked.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + vertexai_stub = MagicMock() + monkeypatch.setitem(sys.modules, "vertexai", vertexai_stub) + monkeypatch.setitem(sys.modules, "vertexai.preview", vertexai_stub.preview) + monkeypatch.setattr( + VertexAIGemmaModels, + "_ensure_access_token", + lambda self, **kwargs: ("fake-gemma-token", "PROJECT_ID"), + ) + return respx_mock.post(USER_API_BASE) + + +def _gemma_kwargs(**overrides): + return { + "model": "vertex_ai/gemma/gemma-3-12b-it-1222199011122", + "messages": [{"role": "user", "content": "hello"}], + "api_base": USER_API_BASE, + "vertex_project": "PROJECT_ID", + "vertex_location": "us-central1", + **overrides, + } + + +@pytest.mark.parametrize("header_kwarg", ["extra_headers", "headers"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_gemma_sends_the_callers_headers_with_the_vertex_token(gemma_chat_route, header_kwarg, sync_mode): + gemma_chat_route.mock(return_value=httpx.Response(200, json=GEMMA_RESPONSE)) + kwargs = _gemma_kwargs(**{header_kwarg: {"X-Custom-Tenant": "tenant-xyz"}}) + + response = litellm.completion(**kwargs) if sync_mode else await litellm.acompletion(**kwargs) + + assert response.choices[0].message.content == "Hello from Gemma!" + sent_headers = gemma_chat_route.calls[0].request.headers + assert sent_headers["x-custom-tenant"] == "tenant-xyz" + assert sent_headers["authorization"] == "Bearer fake-gemma-token" + assert sent_headers["content-type"] == "application/json" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_gemma_streaming_sends_the_callers_headers(gemma_chat_route, sync_mode): + gemma_chat_route.mock(return_value=httpx.Response(200, json=GEMMA_RESPONSE)) + kwargs = _gemma_kwargs(stream=True, extra_headers={"X-Custom-Tenant": "tenant-xyz"}) + + if sync_mode: + text = "".join(chunk.choices[0].delta.content or "" for chunk in litellm.completion(**kwargs)) + else: + stream = await litellm.acompletion(**kwargs) + text = "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + + assert text == "Hello from Gemma!" + sent_headers = gemma_chat_route.calls[0].request.headers + assert sent_headers["x-custom-tenant"] == "tenant-xyz" + assert sent_headers["authorization"] == "Bearer fake-gemma-token" + assert sent_headers["content-type"] == "application/json" + + +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_gemma_keeps_the_vertex_token_over_a_caller_authorization_header( + gemma_chat_route, +): + gemma_chat_route.mock(return_value=httpx.Response(200, json=GEMMA_RESPONSE)) + + await litellm.acompletion(**_gemma_kwargs(extra_headers={"Authorization": "Bearer from-caller"})) + + sent_headers = gemma_chat_route.calls[0].request.headers + assert sent_headers["authorization"] == "Bearer fake-gemma-token"