From f94bde875b2ff060bb3c33f9b018d1e89b248c60 Mon Sep 17 00:00:00 2001
From: Steve White <stwhite@mac.com>
Date: Fri, 28 Feb 2025 15:04:07 -0600
Subject: [PATCH] Fix Gemini model integration with correct model name and
 endpoint

---
 config/config.yaml         | 4 ++--
 query/llm_interface.py     | 5 ++++-
 report/report_synthesis.py | 5 ++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/config/config.yaml b/config/config.yaml
index 923a0cd..da32dc1 100644
--- a/config/config.yaml
+++ b/config/config.yaml
@@ -84,11 +84,11 @@ models:
 
   gemini-2.0-flash-lite:
     provider: "google"
-    model_name: "gemini-2.0-flash-lite-001"
+    model_name: "gemini-2.0-flash-lite"
     temperature: 0.5
     max_tokens: 2048
     top_p: 1.0
-    endpoint: "https://generativelanguage.googleapis.com/v1"
+    endpoint: "https://generativelanguage.googleapis.com/v1beta"
 
 # Default model to use if not specified for a module
 default_model: "llama-3.1-8b-instant"  # Using Groq's Llama 3.1 8B model for testing
diff --git a/query/llm_interface.py b/query/llm_interface.py
index ed18090..e997d6e 100644
--- a/query/llm_interface.py
+++ b/query/llm_interface.py
@@ -106,7 +106,10 @@ class LLMInterface:
             # Special handling for Google Gemini models
             params['model'] = f"gemini/{self.model_config.get('model_name', self.model_name)}"
             # Google Gemini uses a different API base
-            params['api_base'] = self.model_config.get('endpoint', 'https://generativelanguage.googleapis.com/v1')
+            params['api_base'] = self.model_config.get('endpoint', 'https://generativelanguage.googleapis.com/v1beta')
+            
+            # Add additional parameters for Gemini
+            params['custom_llm_provider'] = 'gemini'
         else:
             # Standard provider (OpenAI, Anthropic, etc.)
             params['model'] = self.model_name
diff --git a/report/report_synthesis.py b/report/report_synthesis.py
index ca0eacd..5b32aa1 100644
--- a/report/report_synthesis.py
+++ b/report/report_synthesis.py
@@ -102,7 +102,10 @@ class ReportSynthesizer:
             # Special handling for Google Gemini models
             params['model'] = f"gemini/{self.model_config.get('model_name', self.model_name)}"
             # Google Gemini uses a different API base
-            params['api_base'] = self.model_config.get('endpoint', 'https://generativelanguage.googleapis.com/v1')
+            params['api_base'] = self.model_config.get('endpoint', 'https://generativelanguage.googleapis.com/v1beta')
+            
+            # Add additional parameters for Gemini
+            params['custom_llm_provider'] = 'gemini'
         else:
             # Standard provider (OpenAI, Anthropic, etc.)
             params['model'] = self.model_name