Update CodeWriterCtrlFlow.py

CodeWriterCtrlFlow.py

@@ -80,13 +80,16 @@ class CodeWriterCtrlFlow(ChatAtomicFlow):
             self._state_update_add_chat_message(content=updated_system_message_content,
                                                 role=self.flow_config["system_name"])
 
-        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):
+                updated_system_message_content = self._get_message(self.system_message_prompt_template, input_data)
+                self._state_update_add_chat_message(content=updated_system_message_content,
+                                                    role=self.flow_config["system_name"])
+                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['feedback'] = new_goal
+                input_data = new_input_data