fix(vertex_ai): pass extra_headers to vertex embedding handler

Forward `extra_headers` to `vertex_embedding.embedding()` in the
embedding call path, consistent with how bedrock, gemini, and vertex
batch embeddings already handle it.

Fixes #21020
This commit is contained in:
Chesars 2026-03-13 23:54:42 -03:00
parent bca8730041
commit 63a15c5ffd
2 changed files with 26 additions and 0 deletions

View file

@ -5206,6 +5206,7 @@ def embedding( # noqa: PLR0915
api_key=api_key,
api_base=api_base,
client=client,
extra_headers=headers,
)
elif custom_llm_provider == "oobabooga":
response = oobabooga.embedding(

View file

@ -1459,6 +1459,31 @@ def test_aaavertex_embeddings_distances(
text_embedding = text_response.data[0].embedding
def test_vertex_embedding_passes_extra_headers():
"""
Test that extra_headers are forwarded to the vertex embedding handler.
Relevant issue: https://github.com/BerriAI/litellm/issues/21020
"""
from unittest.mock import MagicMock, patch
mock_response = MagicMock()
with patch.object(
litellm.main.vertex_embedding, "embedding", return_value=mock_response
) as mock_embedding:
litellm.embedding(
model="vertex_ai/textembedding-gecko",
input=["hello"],
extra_headers={"X-Custom-Header": "test-value"},
)
mock_embedding.assert_called_once()
call_kwargs = mock_embedding.call_args.kwargs
assert "extra_headers" in call_kwargs
assert call_kwargs["extra_headers"]["X-Custom-Header"] == "test-value"
def test_vertex_parallel_tool_calls_true():
"""
Test that parallel_tool_calls = True sets the correct tool_config.