Update PlanGeneratorAtomicFlow.py

PlanGeneratorAtomicFlow.py

@@ -10,6 +10,13 @@ class PlanGeneratorAtomicFlow(OpenAIChatAtomicFlow):
             code_file_location: str,
             **kwargs):
         super().__init__(**kwargs)
+        self.hint_for_model = """
+        Make sure your response is in the following format:
+              Response Format:
+              {
+              "plan": "Python printable string of the plan corresponding to the goal",
+              }
+        """
         self.code_file_location = code_file_location
         self.original_system_template = self.system_message_prompt_template.template
 
@@ -44,23 +51,18 @@ class PlanGeneratorAtomicFlow(OpenAIChatAtomicFlow):
             return ret
         except FileNotFoundError:
             return 'There is no function available yet.'
-
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        hint_for_model = """
-        Make sure your response is in the following format:
-              Response Format:
-              {
-              "plan": "Python printable string of the plan corresponding to the goal",
-              }
-        """
+    
+    def _update_prompt_and_input(self, input_data: Dict[str, Any]):
         if 'goal' in input_data:
-            input_data['goal'] += hint_for_model
+            input_data['goal'] += self.hint_for_model
 
         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" \
+            + function_signatures_to_append + "\n"
 
+    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        self._update_prompt_and_input(input_data)
         api_output = super().run(input_data)["api_output"].strip()
         try:
             response = json.loads(api_output)