From 15227173b8c1241c9fbc761342a2344cd90c6593 Mon Sep 17 00:00:00 2001
From: Ashish Bhatia <ashishb@ashishb.net>
Date: Tue, 6 Aug 2024 02:36:40 -0700
Subject: [PATCH] feat: add default_headers to Gemini model (#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 <mpippi@gmail.com>

* Update pyproject.toml

---------

Co-authored-by: Massimiliano Pippi <mpippi@gmail.com>
---
 .../llama_index/llms/gemini/base.py                       | 8 ++++++++
 .../llms/llama-index-llms-gemini/pyproject.toml           | 2 +-
 2 files changed, 9 insertions(+), 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
index 119b252c0d37c..8779fcc4f5516 100644
--- 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
index fa783c1ecfdca..db29a199b57e8 100644
--- 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"