Update ControllerAtomicFlow.py

ControllerAtomicFlow.py

@@ -96,21 +96,19 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
             return ret
         except FileNotFoundError:
             return 'There is no function available yet.'
-
+    
     def _get_plan(self):
         try:
             with open(self.plan_file_location, 'r') as file:
                 return file.read()
         except FileNotFoundError:
             return "There is no plan yet"
-
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        
+    def _update_prompts_and_input(self, input_data: Dict[str, Any]):
         if 'goal' in input_data:
             input_data['goal'] += self.hint_for_model
         if 'human_feedback' in input_data:
             input_data['human_feedback'] += self.hint_for_model
-
-        # self.system_message_prompt_template.template
         plan_to_append = self._get_plan()
         function_signatures_to_append = self._get_library_function_signatures()
         self.system_message_prompt_template.template = \
@@ -120,8 +118,9 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
             + plan_to_append + "\n\n" + f"Make sure the plan your write is at {self.plan_file_location}\n" \
             + f"Make sure the code you call the code writer to write is at {self.code_file_location}"
 
+    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