Update app.py

app.py

@@ -22,7 +22,6 @@ wiki_search = WikipediaSearchTool(
     content_type="text",
     extract_format="WIKI",
 )
-# web_search = DuckDuckGoSearchTool()
 web_search = GoogleSearchTool(provider="serper")
 visit_webpage = VisitWebpageTool()
 final_answer = FinalAnswerTool()
@@ -96,7 +95,7 @@ class EnhancedAgent:
         self.provider = "auto"
         self.timeout = 360
         self.tools = [wiki_search, web_search, visit_webpage, read_file_as_text, reverse_string, final_answer]
-        self.auth_imports = ["pandas", "numpy", "datetime", "json", "re", "math", "os", "requests", "csv", "urllib"]
+        self.auth_imports = ["pandas", "numpy", "datetime", "json", "re", "math", "os", "io", "requests", "csv", "urllib"]
         self.max_steps = 5
         self.agent = CodeAgent(
             model=InferenceClientModel(
@@ -179,17 +178,28 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         if not task_id or question_text is None:
             print(f"Skipping item with missing task_id or question: {item}")
             continue
+        
         try:
-            answer = agent(question_text, file_name)
-            if answer.startswith("Error in generating final LLM output: 504 Server Error: Gateway Time-out for"):
-                print(f"Error running agent on task {task_id}: {submitted_answer}")
-                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {answer}"})
+            answer = None
+            tries = 0
+            while (answer is None) and (tries < 5): 
+                maybe_answer = agent(question_text, file_name)
+                if maybe_answer.startswith("Error in generating final LLM output: 504 Server Error: Gateway Time-out for"):
+                    answer = None
+                    tries += 1
+                    time.sleep(10)
+                else:
+                    answer = maybe_answer
+            if answer is None:
+                print(f"Error running agent on task {task_id}: {maybe_answer}")
+                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {maybe_answer}"})
             else:
+                print(f"Successfully ran agent on task {task_id}.")
                 answers_payload.append({"task_id": task_id, "submitted_answer": answer})
                 results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": answer})
         except Exception as e:
-             print(f"Error running agent on task {task_id}: {e}")
-             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
+            print(f"Error running agent on task {task_id}: {e}")
+            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
 
     if not answers_payload:
         print("Agent did not produce any answers to submit.")