mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
bca8730041
commit
63a15c5ffd
2 changed files with 26 additions and 0 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue