feat: add default_headers to Gemini model (run-llama#15141)

* feat: add default_headers to Gemini model Similar to [OpenAI](https://github.com/run-llama/llama_index/blob/9900560183a2e212515e6273e63159316e056e96/llama-index-integrations/llms/llama-index-llms-openai/llama_index/llms/openai/base.py#L190-L192), allow `default_headers` for Gemini model as well * Update llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py Co-authored-by: Massimiliano Pippi <[email protected]> * Update pyproject.toml --------- Co-authored-by: Massimiliano Pippi <[email protected]>
Tonic-AI · Aug 6, 2024 · 1522717 · 1522717
1 parent da2f3fa
commit 1522717
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py b/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py
@@ -96,6 +96,7 @@ def __init__(
         api_base: Optional[str] = None,
         transport: Optional[str] = None,
         model_name: Optional[str] = None,
+        default_headers: Optional[Dict[str, str]] = None,
         **generate_kwargs: Any,
     ):
         """Creates a new Gemini model interface."""
@@ -123,6 +124,13 @@ def __init__(
             config_params["client_options"] = {"api_endpoint": api_base}
         if transport:
             config_params["transport"] = transport
+        if default_headers:
+            default_metadata: Sequence[Dict[str, str]] = []
+            for key, value in default_headers.items():
+                default_metadata.append((key, value))
+            # `default_metadata` contains (key, value) pairs that will be sent with every request.
+            # When using `transport="rest"`, these will be sent as HTTP headers.
+            config_params["default_metadata"] = default_metadata
         # transport: A string, one of: [`rest`, `grpc`, `grpc_asyncio`].
         genai.configure(**config_params)
 

diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-gemini/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-gemini"
 readme = "README.md"
-version = "0.1.12"
+version = "0.2.0"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"