Update CodeGeneratorAtomicFlow.py

CodeGeneratorAtomicFlow.py

@@ -6,9 +6,14 @@ from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
 
 
 class CodeGeneratorAtomicFlow(ChatAtomicFlow):
-    """Generates code and docstrings with given goal (from controller)"""
+    """Generates one function with docstrings to finish the given goal (from the controller).
+    """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        self.system_message_prompt_template = self.system_message_prompt_template.partial(
+            code_library_file_location="no location yet",
+            code_library="no code yet"
+        )
         self.hint_for_model = """
         Make sure your response is in the following format:
               Response Format:
@@ -33,12 +38,31 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
 
-    def _update_input(self, input_data: Dict[str, Any]):
+    def _get_code_library_file(self, input_data: Dict[str, Any]):
+        assert "memory_files" in input_data, "memory_files not passed to CodeGeneratorAtomicFlow"
+        assert "code_library" in input_data['memory_files'], "code_library not in memory_files"
+        code_library_file_location = input_data['memory_files']['code_library']
+        return code_library_file_location
+
+    def _get_code_library_content(self, input_data: Dict[str, Any]):
+        assert "code_library" in input_data, "code_library not passed to CodeGeneratorAtomicFlow"
+        code_library = input_data['code_library']
+        if len(code_library) == 0:
+            code_library = "No code yet"
+        return code_library
+
+    def _update_prompts_and_input(self, input_data: Dict[str, Any]):
         if 'goal' in input_data:
             input_data['goal'] += self.hint_for_model
+        code_library_file_location = self._get_code_library_file(input_data)
+        code_library = self._get_code_library_content(input_data)
+        self.system_message_prompt_template = self.system_message_prompt_template.partial(
+            code_library_file_location=code_library_file_location,
+            code_library=code_library
+        )
 
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        self._update_input(input_data)
+        self._update_prompts_and_input(input_data)
         api_output = super().run(input_data)["api_output"].strip()
         try:
             response = json.loads(api_output)