Update app.py

app.py

@@ -121,23 +121,32 @@ class BasicAgent:
             return None
 
     # --- Robust Wikipedia Tool ---
-    def wiki_search(self, query):
+     def wiki_search(self, query):
         try:
-            # Clean common AI hallucinations from the query
+            # Clean the query
             query = query.strip(' "').replace('TOOL:', '').replace('INPUT:', '')
-            url = f"https://en.wikipedia.org{query}&format=json"
-            r = requests.get(url, timeout=10).json()
+            # Step 1: Search for the page
+            search_url = "https://en.wikipedia.org"
+            params = {"action": "query", "list": "search", "srsearch": query, "format": "json"}
+            r = requests.get(search_url, params=params, timeout=10).json()
+            
             if not r.get("query", {}).get("search"):
-                return "No results found. Try a broader search term."
+                return "No results. Try simpler keywords."
+
+            # Step 2: Get the top result's snippet and title
+            best_result = r["query"]["search"][0]
+            title = best_result["title"]
+            snippet = BeautifulSoup(best_result["snippet"], "html.parser").get_text()
             
-            title = r["query"]["search"][0]["title"]
-            # Fetch the actual content summary
+            # Step 3: Get the page summary content
             summary_url = f"https://en.wikipedia.org{title.replace(' ', '_')}"
-            page = requests.get(summary_url, timeout=10).json()
-            return page.get("extract", "No summary available.")
-        except:
-            return "Wikipedia access error."
-
+            sum_r = requests.get(summary_url, timeout=10).json()
+            extract = sum_r.get("extract", snippet)
+            
+            return f"Source: {title}\nContent: {extract}"
+        except Exception as e:
+            return f"Wiki error: {str(e)}"
+            
     def youtube_captions(self, url):
         try:
             ydl_opts = {"skip_download": True, "writesubtitles": True, "writeautomaticsub": True}
@@ -149,24 +158,31 @@ class BasicAgent:
             return f"YouTube error: {e}"
 
     def execute_tool(self, tool, input_data, file_url):
-        # Clean input_data (remove extra quotes or labels)
         input_data = input_data.strip(' "')
-
-        if tool == "read_image" or tool == "read_excel":
-            # GAIA Fix: If agent says "none", use the file_url provided by the system
-            target = file_url if (not input_data or input_data.lower() == "none") else input_data
-            if not target: return "Error: No file provided."
-            
+        
+        # Tool: Image or Excel
+        if tool in ["read_image", "read_excel"]:
+            target = file_url if (not input_data or "http" not in input_data) else input_data
+            if not target: return "Error: No file available for this task."
             try:
                 r = requests.get(target, timeout=15)
                 with open("temp_file", "wb") as f: f.write(r.content)
-                
                 if tool == "read_image":
                     return pytesseract.image_to_string(Image.open("temp_file"))
-                else:
-                    return pd.read_excel("temp_file").to_string()[:5000]
+                return pd.read_excel("temp_file").to_string()[:5000]
             except Exception as e:
-                return f"File tool error: {str(e)}"
+                return f"File error: {str(e)}"
+
+        if tool == "wiki_search":
+            return self.wiki_search(input_data)
+        
+        if tool == "calculator":
+            try:
+                expr = re.sub(r'[^\d\+\-\*\/\(\)\.]', '', input_data)
+                return str(eval(expr, {"__builtins__": {}}))
+            except: return "Math error."
+
+        return f"Tool {tool} not recognized."
 
         if tool == "wiki_search": return self.wiki_search(input_data)
         
@@ -200,45 +216,49 @@ class BasicAgent:
 
     def agent_loop(self, question, file_url):
         memory = ""
-        # Explicitly tell the agent about the file
-        context = f"A file for this task is located at: {file_url}" if file_url else "No file attached."
+        context = f"File URL: {file_url}" if file_url else "No file provided."
         
         for step in range(5):
-            prompt = f"""You are a GAIA solver.
-Available tools: read_excel, read_image, wiki_search, calculator.
+            prompt = f"""You are a GAIA solver. Use tools only when necessary.
+Available tools: wiki_search, read_image, read_excel, calculator.
 {context}
 
 Question: {question}
 History: {memory}
 
-If you have the answer, respond with FINAL: [answer].
-If you need to use a tool, respond with:
-TOOL: [tool_name]
-INPUT: [input_data]"""
+Output Format:
+TOOL: tool_name
+INPUT: tool_input
+OR
+FINAL: your_answer"""
 
             response = self.client.chat.completions.create(
                 model="gpt-4o-mini",
-                temperature=0, # Crucial for accuracy
-                messages=[{"role": "system", "content": "You are a precise agent. For 'opposite' questions, think carefully. For files, use the provided URL."},
+                temperature=0,
+                messages=[{"role": "system", "content": "Be concise. If you see a file URL in context, use it for 'read' tools."},
                           {"role": "user", "content": prompt}]
             )
             
-            resp_text = response.choices[0].message.content.strip()
-            print(f"Step {step}: {resp_text}")
-
-            if "FINAL:" in resp_text:
-                return resp_text.split("FINAL:")[-1].strip()
-
-            if "TOOL:" in resp_text and "INPUT:" in resp_text:
-                try:
-                    tool_name = re.search(r"TOOL:\s*(.*)", resp_text).group(1).split('\n')[0].strip()
-                    tool_input = re.search(r"INPUT:\s*(.*)", resp_text).group(1).strip()
-                    result = self.execute_tool(tool_name, tool_input, file_url)
-                    memory += f"\n- {tool_name} output: {result[:1000]}"
-                except:
-                    memory += "\n- Failed to parse tool call."
-            else:
-                memory += f"\n- {resp_text}"
+            resp = response.choices[0].message.content.strip()
+            print(f"Step {step}: {resp}")
+
+            if "FINAL:" in resp:
+                return resp.split("FINAL:")[-1].strip()
+
+            # Robust Tool Parsing
+            try:
+                t_match = re.search(r"TOOL:\s*(.*)", resp, re.I)
+                i_match = re.search(r"INPUT:\s*(.*)", resp, re.I)
+                
+                if t_match and i_match:
+                    t_name = t_match.group(1).strip().lower()
+                    t_input = i_match.group(1).strip()
+                    result = self.execute_tool(t_name, t_input, file_url)
+                    memory += f"\nStep {step} {t_name} output: {result[:800]}"
+                else:
+                    memory += f"\nStep {step} info: {resp}"
+            except:
+                memory += f"\nStep {step}: Parsing error."
         
         return "No answer found."