Update app.py

app.py

@@ -6,7 +6,8 @@ import pandas as pd
 from google import genai
 from google.genai import types
 import time
-from smolagents import CodeAgent, InferenceClientModel, load_tool, DuckDuckGoSearchTool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, PythonInterpreterTool, FinalAnswerTool, LiteLLMModel, tool
+
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -153,59 +154,79 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
 
-# Setup Gemini Client
-api_key = os.getenv("GEMINI_API_KEY")
-if not api_key:
-    raise ValueError("GEMINI_API_KEY is missing.")
-os.environ["GOOGLE_API_KEY"] = api_key
-client = genai.Client()
-model_id = "gemini-2.0-flash-exp"
-
-generation_config = types.GenerateContentConfig(
-    temperature=0.4,
-    top_p=0.9,
-    top_k=40,
-    candidate_count=1,
-    seed=42,
-    presence_penalty=0.0,
-    frequency_penalty=0.0,
-)
-
+# # Setup Gemini Client
+# api_key = os.getenv("GEMINI_API_KEY")
+# if not api_key:
+#     raise ValueError("GEMINI_API_KEY is missing.")
+# os.environ["GOOGLE_API_KEY"] = api_key
+# client = genai.Client()
+# model_id = "gemini-2.0-flash-exp"
+
+# generation_config = types.GenerateContentConfig(
+#     temperature=0.4,
+#     top_p=0.9,
+#     top_k=40,
+#     candidate_count=1,
+#     seed=42,
+#     presence_penalty=0.0,
+#     frequency_penalty=0.0,
+# )
+
+@tool
+def reverse_string(input_string: str) -> str:
+    return input_string[::-1]
+    
 # Define the real agent
 class BasicAgent:
     def __init__(self):
-        print("Initializing BasicAgent with direct Gemini calls + DuckDuckGo tool.")
+        print("Improved BasicAgent initialized with Gemini and enhanced tools.")
 
-        self.search_tool = DuckDuckGoSearchTool()
+        # Load Gemini through LiteLLM (ensure GEMINI_API_TOKEN is set)
+        self.model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash-lite",
+            api_key=os.getenv("GEMINI_API_TOKEN"),
+        )
 
+        # Setup CodeAgent with tools
         self.agent = CodeAgent(
-            tools=[self.search_tool],
-            model=self,  # Pass self as model
-            planning_interval=3
+            tools=[
+                DuckDuckGoSearchTool(),
+                WikipediaSearchTool(),
+                PythonInterpreterTool(),
+                FinalAnswerTool(),
+                reverse_string,
+            ],
+            model=self.model,
+            max_steps=10,
+            add_base_tools=True,
+            additional_authorized_imports=["pandas", "*"]
         )
 
-    def __call__(self, prompt: str, **kwargs) -> str:
-        """Directly call Gemini API for each prompt, returning a string."""
-        try:
-            response = client.models.generate_content(
-                model=model_id,
-                contents=[{"role": "user", "parts": [{"text": prompt}]}],
-                config=generation_config
-            )
-            answer = response.text
-            time.sleep(7)
-            return answer 
-        except Exception as e:
-            return f"Error during Gemini call: {str(e)}"
-
-    def run_agent(self, question: str) -> str:
-        """Run the agent on a question."""
+    def __call__(self, question: str) -> str:
+        """Main callable for processing questions with rate limiting."""
         print(f"Running agent for task: {question[:50]}...")
+
+        # Prompt with strict output formatting
+        system_instruction = (
+            "I will ask you a question. Report your thoughts step by step. "
+            "Finish your answer only with the final answer. In the final answer don't write explanations. "
+            "The answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. "
+            "Avoid units, abbreviations, or articles unless specified. "
+            "Pay attention to each sentence in the question and verify the answer against every part. "
+            "Try searching more sources if initial results are insufficient.\nQUESTION: "
+        )
+
+        prompt = system_instruction + question
         try:
-            result = self.agent.run(question)
-            return result
+            result = self.agent.run(prompt)
         except Exception as e:
-            return f"Error running agent: {str(e)}"
+            result = f"Error: {str(e)}"
+
+        # Rate limiting to avoid 429 errors or API limits
+        print("Waiting 3 seconds to respect rate limits...")
+        time.sleep(3)
+
+        return result
             
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """