From cd257823593d0abbf65fce07af3a5e117ed646ba Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 7 Oct 2025 01:29:19 +0530 Subject: [PATCH] fix passing headers for gemini (#15231) --- litellm/main.py | 4 +- .../test_gemini_header_forwarding.py | 172 ++++++++++++++++++ 2 files changed, 174 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py diff --git a/litellm/main.py b/litellm/main.py index cfb0bef0797..46a024a7631 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2872,7 +2872,7 @@ def completion( # type: ignore # noqa: PLR0915 custom_llm_provider=custom_llm_provider, # type: ignore client=client, api_base=api_base, - extra_headers=extra_headers, + extra_headers=headers, ) elif custom_llm_provider == "vertex_ai": @@ -2941,7 +2941,7 @@ def completion( # type: ignore # noqa: PLR0915 custom_llm_provider=custom_llm_provider, # type: ignore client=client, api_base=api_base, - extra_headers=extra_headers, + extra_headers=headers, ) elif "openai" in model: # Vertex Model Garden - OpenAI compatible models diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py b/tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py new file mode 100644 index 00000000000..16eaaac265a --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py @@ -0,0 +1,172 @@ +""" +Test to verify that custom headers are correctly forwarded to Gemini/Vertex AI API calls. + +This test verifies the fix for the issue where headers configured via +forward_client_headers_to_llm_api were not being passed to Gemini/Vertex AI providers. +""" + +import pytest +from unittest.mock import Mock, patch, MagicMock +import litellm +from litellm import completion + + +class TestGeminiHeaderForwarding: + """Test cases for verifying header forwarding to Gemini/Vertex AI.""" + + @pytest.mark.parametrize( + "custom_llm_provider,model", + [ + ("gemini", "gemini/gemini-1.5-pro"), + ("vertex_ai_beta", "gemini-1.5-pro"), + ("vertex_ai", "gemini-1.5-pro"), + ], + ) + def test_headers_forwarded_to_gemini(self, custom_llm_provider, model): + """ + Test that headers from kwargs are correctly merged and passed to Gemini completion. + + This test verifies that when headers are passed via kwargs (as the proxy does when + forward_client_headers_to_llm_api is configured), they are correctly merged with + extra_headers and passed to the Vertex AI completion handler. + """ + messages = [{"role": "user", "content": "Hello"}] + + # Headers that would be set by the proxy when forwarding client headers + custom_headers = { + "X-Custom-Header": "CustomValue", + "X-BYOK-Token": "secret-token", + } + + # Mock the vertex completion handler + with patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM.completion" + ) as mock_vertex_completion: + # Configure the mock to return a proper response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message.content = "Hello back!" + mock_vertex_completion.return_value = mock_response + + try: + # Call completion with custom headers via kwargs + # This simulates what the proxy does when forward_client_headers_to_llm_api is set + completion( + model=model, + messages=messages, + headers=custom_headers, # This is how proxy passes forwarded headers + custom_llm_provider=custom_llm_provider, + api_key="dummy-key", + ) + + # Verify that the completion handler was called + assert mock_vertex_completion.called, "Vertex completion handler should be called" + + # Get the actual call arguments + call_kwargs = mock_vertex_completion.call_args.kwargs + + # Verify that extra_headers parameter contains our custom headers + assert "extra_headers" in call_kwargs, "extra_headers should be passed to completion" + + passed_headers = call_kwargs["extra_headers"] + assert passed_headers is not None, "extra_headers should not be None" + + # Verify our custom headers are present in the passed headers + for header_key, header_value in custom_headers.items(): + assert ( + header_key in passed_headers + or header_key.lower() in passed_headers + ), f"Header {header_key} should be in extra_headers" + + print(f"✓ Test passed for {custom_llm_provider}/{model}") + print(f" Headers correctly forwarded: {passed_headers}") + + except Exception as e: + pytest.fail( + f"Failed to forward headers to {custom_llm_provider}/{model}: {str(e)}" + ) + + def test_extra_headers_and_headers_merge(self): + """ + Test that both extra_headers and headers parameters are correctly merged. + + This ensures that headers from kwargs (forwarded by proxy) and extra_headers + (passed explicitly) are both included in the final headers sent to the provider. + """ + messages = [{"role": "user", "content": "Hello"}] + + # Headers from proxy (via kwargs["headers"]) + proxy_headers = {"X-Forwarded-Header": "ProxyValue"} + + # Explicit extra_headers + explicit_headers = {"X-Explicit-Header": "ExplicitValue"} + + with patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM.completion" + ) as mock_vertex_completion: + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message.content = "Response" + mock_vertex_completion.return_value = mock_response + + try: + completion( + model="gemini/gemini-1.5-pro", + messages=messages, + headers=proxy_headers, # From proxy forwarding + extra_headers=explicit_headers, # Explicitly passed + custom_llm_provider="gemini", + api_key="dummy-key", + ) + + call_kwargs = mock_vertex_completion.call_args.kwargs + passed_headers = call_kwargs.get("extra_headers", {}) + + # Both sets of headers should be present + assert ( + "X-Forwarded-Header" in passed_headers + or "x-forwarded-header" in passed_headers + ), "Proxy forwarded header should be present" + + assert ( + "X-Explicit-Header" in passed_headers + or "x-explicit-header" in passed_headers + ), "Explicitly passed header should be present" + + print("✓ Both header sources correctly merged and forwarded") + print(f" Final headers: {passed_headers}") + + except Exception as e: + pytest.fail(f"Failed to merge and forward headers: {str(e)}") + + +if __name__ == "__main__": + # Run the tests + test_instance = TestGeminiHeaderForwarding() + + print("\n" + "="*80) + print("Testing Gemini/Vertex AI Header Forwarding") + print("="*80 + "\n") + + # Test each provider + for provider, model in [ + ("gemini", "gemini/gemini-1.5-pro"), + ("vertex_ai_beta", "gemini-1.5-pro"), + ("vertex_ai", "gemini-1.5-pro"), + ]: + print(f"\nTesting {provider}/{model}...") + try: + test_instance.test_headers_forwarded_to_gemini(provider, model) + except Exception as e: + print(f"✗ Test failed: {e}") + + print("\n\nTesting header merging...") + try: + test_instance.test_extra_headers_and_headers_merge() + except Exception as e: + print(f"✗ Test failed: {e}") + + print("\n" + "="*80) + print("All tests completed!") + print("="*80 + "\n") +