Update CodeGeneratorAtomicFlow.py

CodeGeneratorAtomicFlow.py

@@ -41,6 +41,10 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
 
     @classmethod
     def instantiate_from_config(cls, config):
+        """Instantiate a CodeGeneratorAtomicFlow from a configuration.
+        :param config: Configuration dictionary.
+        :return: Instantiated CodeGeneratorAtomicFlow.
+        """
         flow_config = deepcopy(config)
 
         kwargs = {"flow_config": flow_config}
@@ -55,12 +59,20 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
         return cls(**kwargs)
 
     def _get_code_library_file(self, input_data: Dict[str, Any]):
+        """Get the code library file location from the input data.
+        :param input_data: Input data.
+        :return: Code library file location.
+        """
         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]):
+        """Get the code library content from the input data.
+        :param input_data: Input data.
+        :return: Code library content.
+        """
         assert "code_library" in input_data, "code_library not passed to CodeGeneratorAtomicFlow"
         code_library = input_data['code_library']
         if len(code_library) == 0:
@@ -68,6 +80,9 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
         return code_library
 
     def _update_prompts_and_input(self, input_data: Dict[str, Any]):
+        """Update the prompts and input data.
+        :param input_data: Input data.
+        """
         if 'goal' in input_data:
             input_data['goal'] += self.hint_for_model
         code_library_file_location = self._get_code_library_file(input_data)
@@ -78,6 +93,10 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
         )
 
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Run the flow.
+        :param input_data: Input data.
+        :return: Output data.
+        """
         self._update_prompts_and_input(input_data)