Add retry mechanism and `StopEvent` handling in `make_plan` for better formatting enforcement

src/gaia_solving_agent/agent.py

@@ -57,6 +57,7 @@ def get_llm(model_name=cheap_model_name):
 class PlanEvent(Event):
     to_do: Literal["Initialize", "Format", "Replan"] = "Initialize"
     plan: str | None = None
+    n_retries: int = 0
 
 class QueryEvent(Event):
     pass
@@ -75,7 +76,7 @@ class GaiaWorkflow(Workflow):
         return PlanEvent()
 
     @step
-    async def make_plan(self, ctx: Context, ev: PlanEvent) -> PlanEvent | QueryEvent:
+    async def make_plan(self, ctx: Context, ev: PlanEvent) -> PlanEvent | QueryEvent | StopEvent:
         additional_file_path = await ctx.store.get("additional_file_path")
         user_msg = await ctx.store.get("user_msg")
 
@@ -92,16 +93,22 @@ class GaiaWorkflow(Workflow):
             ...
             # TODO : Placeholder for future update
         elif ev.to_do == "Format":
+            if ev.n_retries > 3:
+                return StopEvent(result="Cannot provide a plan. Format may be wrong.", reasoning=ev.plan)
             prompt = f"""
 The original plan is not in the correct format.
 
+______________
 There is the query and constraints you must respect :
 {prompt}
 
+______________
 There is the original plan you must reformat :
 {ev.plan}
 
-Stick strictly to the formatting constraints ! 
+______________
+Ask yourself what you did wrong and fix it.
+Stick strictly to the formatting constraints !
 """
 
         plan = llm.complete(prompt)
@@ -114,7 +121,7 @@ Stick strictly to the formatting constraints !
             extracted is None
             for extracted in [question, known_facts, sub_tasks]
         ):
-            return PlanEvent(to_do="Format", plan=plan.text)
+            return PlanEvent(to_do="Format", plan=plan.text, n_retries=ev.n_retries + 1)
         else:
             await ctx.store.set("question", question if question is not None else "")
             await ctx.store.set("known_facts", known_facts if known_facts is not None else "")