Support additional headers with ChatGoogleGenerativeAI requests

libs/genai/langchain_google_genai/chat_models.py

@@ -483,11 +483,15 @@ def is_lc_serializable(self) -> bool:
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validates params and passes them to google-generativeai package."""
+        additional_headers = values.get("additional_headers") or {}
+        default_metadata = tuple(additional_headers.items())
+
         if values.get("credentials"):
             genai.configure(
                 credentials=values.get("credentials"),
                 transport=values.get("transport"),
                 client_options=values.get("client_options"),
+                default_metadata=default_metadata,
             )
         else:
             google_api_key = get_from_dict_or_env(
@@ -500,6 +504,7 @@ def validate_environment(cls, values: Dict) -> Dict:
                 api_key=google_api_key,
                 transport=values.get("transport"),
                 client_options=values.get("client_options"),
+                default_metadata=default_metadata,
             )
         if (
             values.get("temperature") is not None

libs/genai/langchain_google_genai/llms.py

@@ -154,6 +154,12 @@ class _BaseGoogleGenerativeAI(BaseModel):
         None,
         description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
     )
+    additional_headers: Optional[Dict[str, str]] = Field(
+        None,
+        description=(
+            "A key-value dictionary representing additional headers for the model call"
+        ),
+    )
 
     safety_settings: Optional[Dict[HarmCategory, HarmBlockThreshold]] = None
     """The default safety settings to use for all generations. 

libs/genai/tests/unit_tests/test_chat_models.py

@@ -1,6 +1,7 @@
 """Test chat model integration."""
 
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
+from unittest.mock import Mock, patch
 
 import pytest
 from langchain_core.messages import (
@@ -73,3 +74,34 @@ def test_parse_history() -> None:
 def test_parse_function_history(content: Union[str, List[Union[str, Dict]]]) -> None:
     function_message = FunctionMessage(name="search_tool", content=content)
     _parse_chat_history([function_message], convert_system_message_to_human=True)
+
+
+@pytest.mark.parametrize(
+    "headers", (None, {}, {"X-User-Header": "Coco", "X-User-Header2": "Jamboo"})
+)
+def test_additional_headers_support(headers: Optional[Dict[str, str]]) -> None:
+    mock_configure = Mock()
+    params = {
+        "google_api_key": "[secret]",
+        "client_options": {"api_endpoint": "http://127.0.0.1:8000/ai"},
+        "transport": "rest",
+        "additional_headers": headers,
+    }
+
+    with patch("langchain_google_genai.chat_models.genai.configure", mock_configure):
+        chat = ChatGoogleGenerativeAI(model="gemini-pro", **params)
+
+    expected_default_metadata: tuple = ()
+    if not headers:
+        assert chat.additional_headers == headers
+    else:
+        assert chat.additional_headers
+        assert all(header in chat.additional_headers for header in headers.keys())
+        expected_default_metadata = tuple(headers.items())
+
+    mock_configure.assert_called_once_with(
+        api_key=params["google_api_key"],
+        transport=params["transport"],
+        client_options=params["client_options"],
+        default_metadata=expected_default_metadata,
+    )