From cdf0fae897779ce1e6e8edda4a372b8abe4ffc21 Mon Sep 17 00:00:00 2001 From: Aditya Kaushal Date: Tue, 30 Jun 2026 21:07:39 +0530 Subject: [PATCH 1/4] Update google-generativeai to google-genai version --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c531fd6..c184977 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,4 @@ rank_bm25==0.2.2 requests==2.31.0 sentence_transformers==2.7.0 setuptools==69.0.2 -google-generativeai=0.5.2 +google-genai>=1.0.0 From 18a900c3e50e42cf0d951376269fc038611714fa Mon Sep 17 00:00:00 2001 From: Aditya Kaushal Date: Tue, 30 Jun 2026 21:08:44 +0530 Subject: [PATCH 2/4] Rename google-generativeai to google-genai --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index ef0d3f5..6ed17f0 100644 --- a/setup.py +++ b/setup.py @@ -30,7 +30,7 @@ "sentence-transformers", "transformers", "requests", - "google-generativeai", + "google-genai", "rank_bm25", "pdftotext", ], From 8a60d87ba0a4e1196d41218bb3d45b965dff32bc Mon Sep 17 00:00:00 2001 From: Aditya Kaushal Date: Tue, 30 Jun 2026 21:09:08 +0530 Subject: [PATCH 3/4] Update google-generativeai to google-genai in requirements --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4b18816..a52fb3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,2 +1,2 @@ [build-system] -requires = ["setuptools", "wheel", "google-generativeai", "rank-bm25", "requests", "sentence-transformers", "google-generativeai"] +requires = ["setuptools", "wheel", "google-genai", "rank-bm25", "requests", "sentence-transformers"] From f179350a6208e74ef31bb8e114ac0127ee0dee09 Mon Sep 17 00:00:00 2001 From: Aditya Kaushal Date: Tue, 30 Jun 2026 21:09:52 +0530 Subject: [PATCH 4/4] Update GeminiLLM to integrate Google GenAI client Refactor GeminiLLM to use new Google GenAI client and model configuration. --- rag/llm.py | 40 ++++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/rag/llm.py b/rag/llm.py index 4343d89..9634afc 100644 --- a/rag/llm.py +++ b/rag/llm.py @@ -3,8 +3,6 @@ import traceback from abc import ABC -import google.api_core.exceptions -import google.generativeai as genai import requests @@ -24,43 +22,45 @@ class GeminiLLM(BaseLLM): def __init__(self, *args, **kwargs): """Initialize the Gemini language model. - api_key: str: The API key for the Gemini API. + - model_name: str: The Gemini model to use. """ super().__init__(*args, **kwargs) - GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") or kwargs.get("api_key") - if not GOOGLE_API_KEY: + from google import genai + from google.genai import types + + api_key = os.getenv("GOOGLE_API_KEY") or kwargs.get("api_key") + if not api_key: raise ValueError( "Please set GOOGLE_API_KEY environment variable or set api_key." ) - genai.configure(api_key=GOOGLE_API_KEY) - generation_config = genai.GenerationConfig( - temperature=0.0, - ) - self.model = genai.GenerativeModel( - "text-davinci-003", generation_config=generation_config - ) + self.model_name = kwargs.get("model_name", "gemini-2.0-flash") + self.client = genai.Client(api_key=api_key) + self._generate_config = types.GenerateContentConfig(temperature=0.0) def generate(self, prompt: str, **kwargs) -> str: """Generate text from the model. - prompt: str: The prompt to generate text from. - retried: int: The number of times the request has been retried. """ + from google.genai.errors import ServerError + retried = kwargs.get("retried", 0) if retried < 0: raise Exception("Retried too many times.") try: - response = self.model.generate_content(prompt) - if response.candidates and response.candidates[0].content.parts: - return response.candidates[0].content.parts[0].text - else: - return response.text - except google.api_core.exceptions.InternalServerError: + response = self.client.models.generate_content( + model=kwargs.get("model_name", self.model_name), + contents=prompt, + config=self._generate_config, + ) + return response.text or "" + except ServerError: print("Retrying 500...", retried) - return self.generate(prompt, kwargs=kwargs) + return self.generate(prompt, retried=retried) except Exception: print("An error occurred!") traceback.print_exc() - kwargs["retried"] = retried - 1 - return self.generate(prompt, kwargs=kwargs) + return self.generate(prompt, retried=retried - 1) def chat(self, prompt: str, **kwargs) -> str: return self.generate(prompt, **kwargs)