Update PlanGeneratorAtomicFlow.py

PlanGeneratorAtomicFlow.py

@@ -5,8 +5,13 @@ from flow_modules.aiflows.OpenAIChatFlowModule import OpenAIChatAtomicFlow
 
 from dataclasses import dataclass
 class PlanGeneratorAtomicFlow(OpenAIChatAtomicFlow):
-    def __init__(self, **kwargs):
+    def __init__(
+            self,
+            code_file_location: str,
+            **kwargs):
         super().__init__(**kwargs)
+        self.code_file_location = code_file_location
+        self.original_system_template = self.system_message_prompt_template.template
 
     @classmethod
     def instantiate_from_config(cls, config):
@@ -18,9 +23,28 @@ class PlanGeneratorAtomicFlow(OpenAIChatAtomicFlow):
         kwargs.update(cls._set_up_prompts(flow_config))
         kwargs.update(cls._set_up_backend(flow_config))
 
+        # ~~~ Set up code file location
+        code_file_location = flow_config["code_file_location"]
+        kwargs.update({"code_file_location": code_file_location})
+
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
 
+    def _get_library_function_signatures(self):
+        import importlib.util
+        try:
+            spec = importlib.util.spec_from_file_location("code_library", self.code_file_location)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            ret = ''
+            import inspect
+            for name, obj in inspect.getmembers(module):
+                if inspect.isfunction(obj):
+                    ret += f"{name}: {inspect.signature(obj)}\n"
+            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:
@@ -31,6 +55,12 @@ class PlanGeneratorAtomicFlow(OpenAIChatAtomicFlow):
         """
         if 'goal' in input_data:
             input_data['goal'] += 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" \
+
         api_output = super().run(input_data)["api_output"].strip()
         try:
             response = json.loads(api_output)