Update `setup` method to handle additional files with extensions, adapt prompt template for new variable

src/gaia_solving_agent/agent.py

@@ -1,4 +1,6 @@
 import re
+from pathlib import Path
+from typing import Any
 
 from llama_index.core.agent.workflow import FunctionAgent, AgentWorkflow
 from llama_index.core.prompts import RichPromptTemplate
@@ -31,6 +33,8 @@ def get_llm(model_name=model_name):
 
 class QueryEvent(Event):
     query: str
+    additional_file: Any | None
+    additional_file_path: str = ""
     plan: str
 
 class AnswerEvent(Event):
@@ -42,12 +46,12 @@ class GaiaWorkflow(Workflow):
     @step
     async def setup(self, ev: StartEvent) -> QueryEvent:
         llm = get_llm()
-        prompt_template = RichPromptTemplate(
-            PLANING_PROMPT,
-            template_var_mappings={"query": "user_request"},
-        )
-        plan = llm.complete(prompt_template.format(query=ev.query))
-        return QueryEvent(query=ev.query, plan=plan.text)
+        prompt_template = RichPromptTemplate(PLANING_PROMPT)
+        plan = llm.complete(prompt_template.format(
+            user_request=ev.query,
+            additional_file_extension=Path(ev.additional_file_name).suffix,
+        ))
+        return QueryEvent(query=ev.query, additional_file=ev.additional_file, plan=plan.text)
 
     @step()
     async def multi_agent_process(self, ev: QueryEvent) -> AnswerEvent:

src/gaia_solving_agent/prompts.py

@@ -30,6 +30,9 @@ Respond with exactly:
 - xxxxxxx
 - xxxxxxx
 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 %}
 - xxxxxxx
 - xxxxxxx
 '''