Update CodeGeneratorAtomicFlow.py

CodeGeneratorAtomicFlow.py

@@ -63,14 +63,15 @@ class CodeGeneratorAtomicFlow(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 = "The previous response cannot be parsed with json.loads, it cannot be parsed with json.loads, it could be the backslashes usesd for escaping single quotes in the string arguments of the code are not properly escaped themselves within the JSON context. 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
-            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 response cannot be parsed with json.loads, it cannot be parsed with json.loads, it could be the backslashes usesd for escaping single quotes in the string arguments of the code are not properly escaped themselves within the JSON context. 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