Update ControllerAtomicFlow.py

ControllerAtomicFlow.py

@@ -22,14 +22,22 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
             commands: List[Command],
             plan_file_location: str,
             code_file_location: str,
+            plan_to_append: str,
+            function_signatures_to_append: str,
             **kwargs):
         super().__init__(**kwargs)
         if os.path.isdir(plan_file_location):
             plan_file_location = os.path.join(plan_file_location, "plan.txt")
         self.plan_file_location = plan_file_location
         self.code_file_location = code_file_location
+        self.plan_to_append = plan_to_append
+        self.function_signatures_to_append = function_signatures_to_append
         self.system_message_prompt_template = self.system_message_prompt_template.partial(
-            commands=self._build_commands_manual(commands)
+            commands=self._build_commands_manual(commands),
+            plan_file_location=self.plan_file_location,
+            code_file_location=self.code_file_location,
+            plan_to_append=self.plan_to_append,
+            function_signatures_to_append=self.function_signatures_to_append
         )
         self.hint_for_model = """
         Make sure your response is in the following format:
@@ -45,7 +53,6 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
                   }
               }
         """
-        self.original_system_template = self.system_message_prompt_template.template
 
     @staticmethod
     def _build_commands_manual(commands: List[Command]) -> str:
@@ -77,8 +84,12 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
         # ~~~ Set up file location for plan and code ~~~
         plan_file_location = flow_config["plan_file_location"]
         code_file_location = flow_config["code_file_location"]
+        plan_to_append = flow_config["plan_to_append"]
+        function_signatures_to_append = flow_config["function_signatures_to_append"]
         kwargs.update({"plan_file_location": plan_file_location})
         kwargs.update({"code_file_location": code_file_location})
+        kwargs.update({"plan_to_append": plan_to_append})
+        kwargs.update({"function_signatures_to_append": function_signatures_to_append})
 
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
@@ -109,14 +120,9 @@ class ControllerAtomicFlow(OpenAIChatAtomicFlow):
             input_data['goal'] += self.hint_for_model
         if 'human_feedback' in input_data:
             input_data['human_feedback'] += self.hint_for_model
-        plan_to_append = self._get_plan()
-        function_signatures_to_append = self._get_library_function_signatures()
-        self.system_message_prompt_template.template = \
-            self.original_system_template + "\n\n" + f"Here are the available functions at {self.code_file_location}\n" \
-            + function_signatures_to_append + "\n" \
-            + f"Here is the step-by-step plan to achieve the goal:\n" \
-            + 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}"
+        self.plan_to_append = self._get_plan()
+        self.function_signatures_to_append = self._get_library_function_signatures()
+        self.system_message_prompt_template = self.system_message_prompt_template.partial(plan_to_append = self.plan_to_append, function_signatures_to_append = self.function_signatures_to_append)
 
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         self._update_prompts_and_input(input_data)