From fc3f82b85a7c53e38984916e1943c10d2ef83a09 Mon Sep 17 00:00:00 2001 From: Damien Date: Mon, 15 Dec 2025 20:46:09 -0600 Subject: [PATCH] feat(gemini): support extra_headers in batch embeddings (#18004) * feat(vertex_ai): support extra_headers in batch embeddings * test(vertex_ai): add Gemini batch embeddings tests for custom api_base --- .../batch_embed_content_handler.py | 10 ++ litellm/main.py | 1 + .../vertex_ai/test_gemini_batch_embeddings.py | 145 ++++++++++++++++++ 3 files changed, 156 insertions(+) create mode 100644 tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 859bb0a6984..07f57a4a7f6 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -46,6 +46,7 @@ class GoogleBatchEmbeddings(VertexLLM): aembedding: Optional[bool] = False, timeout=300, client=None, + extra_headers: Optional[dict] = None, ) -> EmbeddingResponse: _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -90,6 +91,15 @@ class GoogleBatchEmbeddings(VertexLLM): headers = { "Content-Type": "application/json; charset=utf-8", } + if auth_header is not None: + if isinstance(auth_header, dict): + # For Gemini with custom api_base: auth_header is {"x-goog-api-key": "..."} + headers.update(auth_header) + else: + # For Vertex AI: auth_header is a Bearer token string + headers["Authorization"] = f"Bearer {auth_header}" + if extra_headers is not None: + headers.update(extra_headers) ## LOGGING logging_obj.pre_call( diff --git a/litellm/main.py b/litellm/main.py index b08ffd16e3d..216a3fe0ffd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4663,6 +4663,7 @@ def embedding( # noqa: PLR0915 api_key=gemini_api_key, api_base=api_base, client=client, + extra_headers=headers, ) elif custom_llm_provider == "vertex_ai": diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py new file mode 100644 index 00000000000..7047be4241b --- /dev/null +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -0,0 +1,145 @@ +""" +Test Gemini batch embeddings with custom api_base and extra_headers. + +This test ensures that: +1. Authentication headers are properly included when using custom api_base +2. The extra_headers parameter is correctly passed through +3. Both dict-based auth_header (Gemini) and Bearer token (Vertex AI) are handled +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../../../..")) + +import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + + +def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header(): + """ + Test that Gemini batch embeddings include auth_header when using custom api_base. + + This test verifies that when using Gemini embeddings with a custom api_base + (e.g., Cloudflare AI Gateway), the x-goog-api-key header is properly included + in the HTTP request. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return None, "test-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token", + side_effect=mock_auth_token + ), patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url" + ) as mock_get_token: + # Mock the _get_token_and_url to return auth_header dict and URL + mock_get_token.return_value = ( + {"x-goog-api-key": "test-gemini-api-key"}, + "https://gateway.ai.cloudflare.com/v1/test/noauth/google-ai-studio/v1beta" + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + { + "embeddings": { + "values": [0.1, 0.2, 0.3, 0.4, 0.5] + } + } + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="gemini/text-embedding-004", + input=["Hello, world!"], + api_key="test-gemini-api-key", + api_base="https://gateway.ai.cloudflare.com/v1/test/noauth/google-ai-studio/v1beta", + client=client + ) + + # Verify the POST was called + mock_post.assert_called_once() + + # Get the headers that were passed to the POST request + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + headers = kwargs.get("headers", {}) + + # Verify auth_header is included + assert "x-goog-api-key" in headers, f"x-goog-api-key not in headers: {headers}" + assert headers["x-goog-api-key"] == "test-gemini-api-key" + + # Verify Content-Type is still present + assert "Content-Type" in headers + assert headers["Content-Type"] == "application/json; charset=utf-8" + + +def test_gemini_batch_embeddings_with_extra_headers(): + """ + Test that extra_headers parameter is properly included in the request. + + This test verifies that custom headers passed via extra_headers are + properly merged into the request headers. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return None, "test-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token", + side_effect=mock_auth_token + ), patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url" + ) as mock_get_token: + # Mock the _get_token_and_url to return auth_header dict and URL + mock_get_token.return_value = ( + {"x-goog-api-key": "test-gemini-api-key"}, + "https://gateway.ai.cloudflare.com/v1/test/google-ai-studio/v1beta" + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + { + "embeddings": { + "values": [0.1, 0.2, 0.3] + } + } + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="gemini/text-embedding-004", + input=["Test"], + api_key="test-gemini-api-key", + api_base="https://gateway.ai.cloudflare.com/v1/test/google-ai-studio/v1beta", + headers={"Authorization": "Bearer test-token", "X-Custom": "custom-value"}, + client=client + ) + + # Verify the POST was called + mock_post.assert_called_once() + + # Get the headers that were passed to the POST request + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + headers = kwargs.get("headers", {}) + + # Verify all headers are included + assert "x-goog-api-key" in headers + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-token" + assert "X-Custom" in headers + assert headers["X-Custom"] == "custom-value" +