diff --git a/tests/test_litellm/test_vertex_ai_embedding_headers.py b/tests/test_litellm/test_vertex_ai_embedding_headers.py new file mode 100644 index 00000000000..a46e1e76bbb --- /dev/null +++ b/tests/test_litellm/test_vertex_ai_embedding_headers.py @@ -0,0 +1,16 @@ +import pytest +from litellm import embedding + +def test_vertex_ai_embedding_extra_headers(): + # Test that extra_headers are passed without crashing + try: + response = embedding( + model="vertex_ai/text-embedding-004", + input=["hello"], + extra_headers={"X-Custom-Header": "test-value"} + ) + except Exception as e: + # We expect a 401/404 if no real creds, + # but we are checking that it doesn't fail with a TypeError + if "extra_headers" in str(e): + pytest.fail("extra_headers not accepted by vertex_ai embedding")