Make `run_agent` asynchronous, enhance agent handling, and improve answer extraction logic

src/gaia_solving_agent/hf_submission_api.py

@@ -1,15 +1,19 @@
+import re
+
 import pandas as pd
 import requests
+from llama_index.core.agent.workflow import BaseWorkflowAgent
+from workflows import Workflow
 
-from agent import BasicAgent
+from gaia_solving_agent.agent import GaiaWorkflow
 
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
-def instantiate_agent(space_id: str):
+def instantiate_agent(space_id: str, **agent_kwargs):
     try:
-        agent = BasicAgent()
+        agent = GaiaWorkflow(**agent_kwargs)
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -45,7 +49,7 @@ def fetching_questions(api_url: str = DEFAULT_API_URL):
     return questions_data
 
 
-def run_agent(agent, questions_data):
+async def run_agent(agent, questions_data):
     results_log = []
     answers_payload = []
     print(f"Running agent on {len(questions_data)} questions...")
@@ -56,7 +60,12 @@ def run_agent(agent, questions_data):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            if isinstance(agent, BaseWorkflowAgent):
+                submitted_answer = await agent.run(user_msg=question_text)
+            elif isinstance(agent, Workflow):
+                submitted_answer = await agent.run(query=question_text)
+            else:
+                raise NotImplementedError(f"Invalid agent type: {type(agent)}")
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
@@ -70,6 +79,18 @@ def run_agent(agent, questions_data):
 
 
 def prepare_submission_data(username:str, agent_code: str, answers_payload: list[dict]):
+    def extract_answer(answer: str) -> str | int:
+        pattern = r"FINAL ANSWER:\s*(.*)"
+        search = re.search(pattern, answer)
+        extracted_answer = search.group(1) if search else ""
+        try:
+            return eval(extracted_answer)
+        except:
+            return extracted_answer
+
+    for answer in answers_payload:
+        answer["submitted_answer"] = extract_answer(str(answer["submitted_answer"]))
+
     submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
     status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
     print(status_update)