Add support for handling `additional_file` and `additional_file_path` in workflow agents, update `setup` method, and adapt prompt template accordingly

src/gaia_solving_agent/agent.py

@@ -34,7 +34,7 @@ def get_llm(model_name=model_name):
 class QueryEvent(Event):
     query: str
     additional_file: Any | None
-    additional_file_path: str = ""
+    additional_file_path: str | Path | None = None
     plan: str
 
 class AnswerEvent(Event):
@@ -47,11 +47,17 @@ class GaiaWorkflow(Workflow):
     async def setup(self, ev: StartEvent) -> QueryEvent:
         llm = get_llm()
         prompt_template = RichPromptTemplate(PLANING_PROMPT)
+        file_extension = Path(ev.additional_file_path).suffix if ev.additional_file_path else ""
         plan = llm.complete(prompt_template.format(
-            user_request=ev.query,
-            additional_file_extension=Path(ev.additional_file_name).suffix,
+            user_request=ev.user_msg,
+            additional_file_extension=file_extension,
         ))
-        return QueryEvent(query=ev.query, additional_file=ev.additional_file, plan=plan.text)
+        return QueryEvent(
+            query=ev.user_msg,
+            additional_file=ev.additional_file,
+            additional_file_path=ev.additional_file_path,
+            plan=plan.text,
+        )
 
     @step()
     async def multi_agent_process(self, ev: QueryEvent) -> AnswerEvent:
@@ -59,7 +65,12 @@ class GaiaWorkflow(Workflow):
         from llama_index.core.memory import ChatMemoryBuffer
         memory = ChatMemoryBuffer.from_defaults(token_limit=100000)
 
-        agent_output = await gaia_solving_agent.run(user_msg=ev.plan, memory=memory)
+        agent_output = await gaia_solving_agent.run(
+            user_msg=ev.plan,
+            memory=memory,
+            additional_file=ev.additional_file,
+            additional_file_path=ev.additional_file_path,
+        )
         return AnswerEvent(plan=ev.plan, answer=str(agent_output))
 
     @step

src/gaia_solving_agent/hf_submission_api.py

@@ -82,9 +82,17 @@ async def run_agent(agent, questions_data):
             continue
         try:
             if isinstance(agent, BaseWorkflowAgent):
-                submitted_answer = await agent.run(user_msg=question_text)
+                submitted_answer = await agent.run(
+                    user_msg=question_text,
+                    additional_file=additional_file,
+                    additional_file_path=additional_file_path,
+                )
             elif isinstance(agent, Workflow):
-                submitted_answer = await agent.run(query=question_text)
+                submitted_answer = await agent.run(
+                    user_msg=question_text,
+                    additional_file=additional_file,
+                    additional_file_path=additional_file_path,
+                )
             else:
                 raise NotImplementedError(f"Invalid agent type: {type(agent)}")
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})

src/gaia_solving_agent/prompts.py

@@ -32,7 +32,7 @@ Respond with exactly:
 3. Sub-tasks:
 {% if additional_file_extension|length -%}
 - There is an additional file to analyse. It's file extension is {{ additional_file_extension }}. You must analyse it to find out
-{%- else %}
+{%- endif %}
 - xxxxxxx
 - xxxxxxx
 '''