Update app.py

app.py

@@ -19,7 +19,7 @@ import logging
 import shutil
 import pandas as pd
 from docx import Document as DocxDocument
-
+import google.generativeai as genai
 
 # Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -40,6 +40,7 @@ MODELS = [
     "@cf/meta/llama-3.1-8b-instruct",
     "mistralai/Mistral-Nemo-Instruct-2407",
     "mistralai/Mathstral-7B-v0.1",
+    "gemini-1.5-flash",
     "duckduckgo/gpt-4o-mini",
     "duckduckgo/claude-3-haiku",
     "duckduckgo/llama-3.1-70b",
@@ -352,6 +353,53 @@ def summarize_web_results(query: str, search_results: List[Dict[str, str]], conv
     except Exception as e:
         return f"An error occurred during summarization: {str(e)}"
 
+
+def get_response_from_gemini(query, context, file_type, num_calls=1, temperature=0.2):
+    # Configure the Gemini API
+    genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+
+    # Define the model
+    model = genai.GenerativeModel(
+        model_name="gemini-1.5-flash",
+        generation_config={
+            "temperature": temperature,
+            "top_p": 1,
+            "top_k": 1,
+            "max_output_tokens": 2048,
+        },
+    )
+
+    # Create the system instruction based on file type
+    if file_type == "excel":
+        system_instruction = """You are a highly specialized data analyst with expertise in Excel spreadsheets. 
+        Your task is to analyze the provided Excel data and answer the user's query accurately and concisely. 
+        Focus on identifying key metrics, trends, and significant details relevant to the query. 
+        Do not make assumptions or include information not explicitly supported by the dataset."""
+    elif file_type == "pdf":
+        system_instruction = """You are a highly specialized document analyst with expertise in extracting information from PDF documents. 
+        Your task is to analyze the provided PDF content and answer the user's query accurately and comprehensively. 
+        Focus on key points, important details, and relevant information from the document. 
+        Ensure your response is strictly based on the provided context."""
+    else:
+        raise ValueError("Invalid file type. Use 'excel' or 'pdf'.")
+
+    # Prepare the chat session
+    chat_session = model.start_chat(history=[])
+
+    full_response = ""
+    for _ in range(num_calls):
+        try:
+            # Send the message with context and query
+            response = chat_session.send_message(
+                f"{system_instruction}\n\nContext:\n{context}\n\nUser query: {query}"
+            )
+            full_response += response.text + "\n"
+        except Exception as e:
+            print(f"Error in generating response from Gemini: {str(e)}")
+            return f"An error occurred with the Gemini model: {str(e)}. Please try again."
+
+    return full_response.strip()
+
 def get_response_from_excel(query, model, context, num_calls=3, temperature=0.2):
     logging.info(f"Getting response from Excel using model: {model}")
     
@@ -460,12 +508,19 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
 
             # Process Excel documents
             if excel_docs:
-                for response in get_response_from_excel(message, model, excel_context, num_calls, temperature):
+                if model == "gemini-1.5-flash":
+                    response = get_response_from_gemini(message, excel_context, "excel", num_calls, temperature)
                     yield response
+                else:
+                    for response in get_response_from_excel(message, model, excel_context, num_calls, temperature):
+                        yield response
 
             # Process other documents (PDF, Word)
             if other_docs:
-                if model == "@cf/meta/llama-3.1-8b-instruct":
+                if model == "gemini-1.5-flash":
+                    response = get_response_from_gemini(message, other_context, "pdf", num_calls, temperature)
+                    yield response
+                elif model == "@cf/meta/llama-3.1-8b-instruct":
                     for response in get_response_from_cloudflare(prompt="", context=other_context, query=message, num_calls=num_calls, temperature=temperature, search_type="document"):
                         yield response
                 else: