Update app.py

app.py

@@ -5,12 +5,17 @@ import inspect
 import pandas as pd
 import smolagents
 from smolagents import DuckDuckGoSearchTool, VisitWebpageTool
-import time  # Added: For sleep functionality
-from functools import lru_cache  # Added: For caching search results
+import time 
+from functools import lru_cache  
+
+import google.generativeai as genai
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 
 from youtube_transcript_api import YouTubeTranscriptApi
 import re
 
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") 
+genai.configure(api_key=GOOGLE_API_KEY)
 
 class YouTubeVideoTool:
     def __init__(self):
@@ -155,10 +160,41 @@ class BasicAgent:
 
     
     
-    def _formulate_direct_answer(self, relevant_info, question):
-        if self.model:
+def _formulate_direct_answer(self, relevant_info, question):
+    if self.model and self.model.startswith('gemini'):
+        try:
+            # Configure the model
+            generation_config = {
+                "temperature": 0.7,
+                "top_p": 0.95,
+                "top_k": 40,
+                "max_output_tokens": 1024,
+            }
+            
+            safety_settings = {
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+            }
+            
+            # Initialize the model
+            model = genai.GenerativeModel(
+                model_name="gemini-pro",  # Adjust as needed based on your model string
+                generation_config=generation_config,
+                safety_settings=safety_settings
+            )
+            
+            # Prepare prompt and generate response
+            prompt = f"Question: {question}\n\nRelevant information: {relevant_info}\n\nProvide a concise answer based only on the given information."
+            response = model.generate_content(prompt)
+            return response.text
+            
+        except Exception as e:
+            print(f"Error using Gemini model: {e}")
             return f"Based on the search: {relevant_info}"
-        return relevant_info
+            
+    return relevant_info