fix passing headers for gemini (#15231)

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

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")
+