Update app.py

app.py

@@ -6,6 +6,7 @@ import pandas as pd
 from google import genai
 from google.genai import types
 import time
+from smolagents import CodeAgent, InferenceClientModel, load_tool, DuckDuckGoSearchTool
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -22,19 +23,15 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #         print(f"Agent returning fixed answer: {fixed_answer}")
 #         return fixed_answer
 
-class BasicAgent:
-    def __init__(self):
-        print("CustomAgent using Gemini 2.0 initialized.")
-
-        # Set up Gemini API
-        api_key = os.environ.get("GEMINI_API_KEY")
+class GeminiModel:
+    def __init__(self, model_name="gemini-2.0-flash-exp"):
+        api_key = os.getenv("GEMINI_API_KEY")
         if not api_key:
-            raise ValueError("GEMINI_API_KEY not found in environment variables.")
-        
-        os.environ["GOOGLE_API_KEY"] = api_key  # Required by google-genai
+            raise ValueError("GEMINI_API_KEY is missing.")
 
+        os.environ["GOOGLE_API_KEY"] = api_key
         self.client = genai.Client()
-        self.model_id = "gemini-2.0-flash-exp"  # Or "gemini-1.5-flash-002" if you want faster
+        self.model_id = model_name
 
         self.generation_config = types.GenerateContentConfig(
             temperature=0.4,
@@ -46,26 +43,45 @@ class BasicAgent:
             frequency_penalty=0.0,
         )
 
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-
+    def run(self, prompt: str) -> str:
+        """Send prompt to Gemini."""
         try:
-            response = self.client.models.generate_content(
+            response = self.client.generate_content(
                 model=self.model_id,
-                contents=f"""You are a smart, factual assistant. Answer clearly and concisely:\n\n{question}\n\nProvide only the final answer without extra commentary.""",
-                config=self.generation_config,
+                contents=[{"role": "user", "parts": [{"text": prompt}]}],
+                generation_config=self.generation_config
             )
-            answer = response.text.strip()
+            return response.candidates[0].content.parts[0].text.strip()
+        except Exception as e:
+            return f"Error during Gemini call: {str(e)}"
 
-            # ✨ Add a short sleep to avoid hitting rate limits
-            time.sleep(7)  # Wait 7 seconds after each question
-            print(f"Returning answer (first 100 chars): {answer[:100]}")
-            return answer
+# Define BasicAgent properly
+class BasicAgent:
+    def __init__(self):
+        print("Initializing CodeAgent with Gemini + tools.")
 
-        except Exception as e:
-            print(f"Error during Gemini call: {str(e)}")
-            return f"Error: {str(e)}"
+        # Load tools
+        self.search_tool = DuckDuckGoSearchTool()
+
+        # Build the agent
+        self.agent = CodeAgent(
+            tools=[self.search_tool],
+            model=GeminiModel(),  # Our simple Gemini wrapper
+            planning_interval=3  # Activate planning
+        )
 
+    def __call__(self, question: str) -> str:
+        """Call the CodeAgent."""
+        print(f"Running agent for task: {question[:50]}...")
+        try:
+            result = self.agent.run(question)
+
+            # Sleep to respect rate limits
+            time.sleep(7)
+            return result
+        except Exception as e:
+            return f"Error running agent: {str(e)}"
+            
 # class BasicAgent(ReActAgent):
 #     def __init__(self):
 #         print("BasicAgent using local LLM initialized.")