Update app.py

app.py

@@ -3,9 +3,8 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from langchain_google_genai import ChatGoogleGenerativeAI
-from tools import InternetSearchTool, WebscraperTool, ExcelTool, WikipediaTool, MediaTool
-from smolagents import CodeAgent, InferenceClientModel
+from tools import InternetSearchTool, WebscraperTool, ExcelTool, WikipediaTool, MediaTool, LiteLLMModel
+from smolagents import CodeAgent, InferenceClientModel, PythonInterpreterTool
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -17,7 +16,9 @@ class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
         # Initialize the Hugging Face model
-        model = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest")
+        llm = LiteLLMModel(
+            model_id="gemini-2.0-flash", api_key=os.getenv("GEMINI_API_KEY")
+        )
         
         # Initialize tools
         search_tool = InternetSearchTool()
@@ -25,18 +26,19 @@ class BasicAgent:
         excel_tool = ExcelTool()
         media_tool = MediaTool()
         wikipedia_tool = WikipediaTool()
+        python_tool = PythonInterpreterTool()
 
         # Create agent
         self.agent= CodeAgent(
-            tools=[search_tool, webscraper_tool, excel_tool, media_tool, wikipedia_tool],
-            model=InferenceClientModel(token=os.getenv("HF_TOKEN"), provider="together"),
+            tools=[search_tool, webscraper_tool, excel_tool, media_tool, wikipedia_tool, python_tool],
+            model=llm,
             add_base_tools=True,  # Add any additional base tools
             planning_interval=3,   # Enable planning every 3 steps
             max_steps=10
         )
 
     def summarize(self, question: str) -> str:
-        prompt = """You are an AI assistant for summarizing tasks. You will receive a long request message, your task is to make it shorter by removing irrelevant details. Do no attempt to find an answer to the request or any part of the request.
+        prompt = """You are an AI assistant for summarizing tasks. You will receive a long request message, your task is to make it shorter by removing irrelevant details. Do no attempt to find an answer to the request or any part of the request. Make sure to always include any requirement towards the answer's format in your summary.
         
         EXAMPLE:
         'Hi, we've been learning about reptiles in biology and I'd like to know more about crocodiles. Can you tell me how many legs an average crocodile has? Please answer only with a number! Thanks in advance!'