add docs and comments

CodeGeneratorAtomicFlow.py

@@ -23,8 +23,25 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     *Output Interface*:
     - `code`
     - `language_of_code`
+
+
+    *Configuration Parameters*:
+    - Also refer to ChatAtomicFlow (https://huggingface.co/aiflows/ChatFlowModule/blob/main/ChatAtomicFlow.py)
+    - `input_interface_non_initialized`: The input interface when the conversation is not initialized.
+    - `input_interface_initialized`: The input interface when the conversation is initialized.
+    - `output_interface`: The output interface.
+    - `backend`: The backend to use for the Chat API.
+    - `system_message_prompt_template`: The template for the system message prompt.
+    - `human_message_prompt_template`: The template for the human message prompt.
+    - `init_human_message_prompt_template`: The initial human message prompt.
+
+
     """
     def __init__(self, **kwargs):
+        """Initialize the CodeGeneratorAtomicFlow.
+        :param kwargs: Keyword arguments.
+        :type kwargs: Any
+        """
         super().__init__(**kwargs)
         self.system_message_prompt_template = self.system_message_prompt_template.partial(
             code_library_file_location="no location yet",
@@ -43,7 +60,9 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     def instantiate_from_config(cls, config):
         """Instantiate a CodeGeneratorAtomicFlow from a configuration.
         :param config: Configuration dictionary.
+        :type config: Dict[str, Any]
         :return: Instantiated CodeGeneratorAtomicFlow.
+        :rtype: CodeGeneratorAtomicFlow
         """
         flow_config = deepcopy(config)
 
@@ -61,7 +80,11 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     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.
+        :type input_data: Dict[str, Any]
         :return: Code library file location.
+        :rtype: str
+        :raises AssertionError: If memory_files is not in input_data.
+        :raises AssertionError: If code_library is not in memory_files.
         """
         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"
@@ -71,7 +94,10 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     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.
+        :type input_data: Dict[str, Any]
         :return: Code library content.
+        :rtype: str
+        :raises AssertionError: If code_library is not in input_data.
         """
         assert "code_library" in input_data, "code_library not passed to CodeGeneratorAtomicFlow"
         code_library = input_data['code_library']
@@ -82,6 +108,7 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     def _update_prompts_and_input(self, input_data: Dict[str, Any]):
         """Update the prompts and input data.
         :param input_data: Input data.
+        :type input_data: Dict[str, Any]
         """
         if 'goal' in input_data:
             input_data['goal'] += self.hint_for_model
@@ -95,11 +122,12 @@ class CodeGeneratorAtomicFlow(ChatAtomicFlow):
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         """Run the flow.
         :param input_data: Input data.
+        :type input_data: Dict[str, Any]
         :return: Output data.
+        :rtype: Dict[str, Any]
         """
         self._update_prompts_and_input(input_data)
 
-
         while True:
             api_output = super().run(input_data)["api_output"].strip()
             try:

README.md

@@ -1,11 +1,17 @@
 # Table of Contents
 
+* [run](#run)
 * [CodeGeneratorAtomicFlow](#CodeGeneratorAtomicFlow)
   * [CodeGeneratorAtomicFlow](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow)
+    * [\_\_init\_\_](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.__init__)
     * [instantiate\_from\_config](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.instantiate_from_config)
     * [run](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.run)
 * [\_\_init\_\_](#__init__)
 
+<a id="run"></a>
+
+# run
+
 <a id="CodeGeneratorAtomicFlow"></a>
 
 # CodeGeneratorAtomicFlow
@@ -22,14 +28,12 @@ This class wraps around the Chat API to generate code from a goal. One thing wor
 make sure the code generator does not write repetitive code that is present in the library, so we need to inject
 the function signatures in the library to the system prompts.
 
-Please also visit https://huggingface.co/aiflows/ChatFlowModule for more parameter explaination.
-
-*Input Interface Non Initialized (i.e. input when the chat is first initialized)*:
+*Input Interface Non Initialized*:
 - `goal`
 - `code_library`
 - `memory_files`
 
-*Input Interface Initialized (i.e. input when the chat is already initialized, this happens when we are calling the chat again)*:
+*Input Interface Initialized*:
 - `goal`
 - `code_library`
 - `memory_files`
@@ -38,6 +42,31 @@ Please also visit https://huggingface.co/aiflows/ChatFlowModule for more paramet
 - `code`
 - `language_of_code`
 
+
+*Configuration Parameters*:
+- Also refer to ChatAtomicFlow (https://huggingface.co/aiflows/ChatFlowModule/blob/main/ChatAtomicFlow.py)
+- `input_interface_non_initialized`: The input interface when the conversation is not initialized.
+- `input_interface_initialized`: The input interface when the conversation is initialized.
+- `output_interface`: The output interface.
+- `backend`: The backend to use for the Chat API.
+- `system_message_prompt_template`: The template for the system message prompt.
+- `human_message_prompt_template`: The template for the human message prompt.
+- `init_human_message_prompt_template`: The initial human message prompt.
+
+<a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.__init__"></a>
+
+#### \_\_init\_\_
+
+```python
+def __init__(**kwargs)
+```
+
+Initialize the CodeGeneratorAtomicFlow.
+
+**Arguments**:
+
+- `kwargs` (`Any`): Keyword arguments.
+
 <a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.instantiate_from_config"></a>
 
 #### instantiate\_from\_config
@@ -51,11 +80,11 @@ Instantiate a CodeGeneratorAtomicFlow from a configuration.
 
 **Arguments**:
 
-- `config`: Configuration dictionary.
+- `config` (`Dict[str, Any]`): Configuration dictionary.
 
 **Returns**:
 
-Instantiated CodeGeneratorAtomicFlow.
+`CodeGeneratorAtomicFlow`: Instantiated CodeGeneratorAtomicFlow.
 
 <a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.run"></a>
 
@@ -69,11 +98,11 @@ Run the flow.
 
 **Arguments**:
 
-- `input_data`: Input data.
+- `input_data` (`Dict[str, Any]`): Input data.
 
 **Returns**:
 
-Output data.
+`Dict[str, Any]`: Output data.
 
 <a id="__init__"></a>