Update PlanGeneratorAtomicFlow.py

PlanGeneratorAtomicFlow.py

@@ -40,12 +40,13 @@ class PlanGeneratorAtomicFlow(ChatAtomicFlow):
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         self._update_prompts_and_input(input_data)
         api_output = super().run(input_data)["api_output"].strip()
-        try:
-            response = json.loads(api_output)
-            return response
-        except json.decoder.JSONDecodeError:
-            new_goal = f"Here is your previous response {api_output}, it cannot be parsed with json.loads, please fix this issue. Provide your response only in JSON format."
-            new_input_data = input_data.copy()
-            new_input_data['goal'] = new_goal
-            new_api_output = super().run(new_input_data)["api_output"].strip()
-            return json.loads(new_api_output)
+        while True:
+            api_output = super().run(input_data)["api_output"].strip()
+            try:
+                response = json.loads(api_output)
+                return response
+            except (json.decoder.JSONDecodeError, json.JSONDecodeError):
+                new_goal = "The previous respond cannot be parsed with json.loads. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
+                new_input_data = input_data.copy()
+                new_input_data['goal'] = new_goal
+                input_data = new_input_data