Upload 6 files

CodeGeneratorAtomicFlow.py

@@ -0,0 +1,112 @@
+import json
+from copy import deepcopy
+from typing import Any, Dict
+from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
+
+
+
+class CodeGeneratorAtomicFlow(ChatAtomicFlow):
+    """This class wraps around the Chat API to generate code from a goal. One thing worth noting is that we need to
+    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.
+
+    *Input Interface Non Initialized*:
+    - `goal`
+    - `code_library`
+    - `memory_files`
+
+    *Input Interface Initialized*:
+    - `goal`
+    - `code_library`
+    - `memory_files`
+
+    *Output Interface*:
+    - `code`
+    - `language_of_code`
+    """
+    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:
+              {
+              "language_of_code": "language of the code",
+              "code": "String of the code and docstrings corresponding to the goal",
+              }
+        """
+
+    @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}
+
+        # ~~~ Set up prompts ~~~
+        kwargs.update(cls._set_up_prompts(flow_config))
+
+        # ~~~ Set up backend ~~~
+        kwargs.update(cls._set_up_backend(flow_config))
+
+        # ~~~ Instantiate flow ~~~
+        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:
+            code_library = "No code yet"
+        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)
+        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]:
+        """Run the flow.
+        :param input_data: Input data.
+        :return: Output data.
+        """
+        self._update_prompts_and_input(input_data)
+
+
+        while True:
+            api_output = super().run(input_data)["api_output"].strip()
+            try:
+                response = json.loads(api_output)
+                return response
+            except (json.decoder.JSONDecodeError, json.JSONDecodeError):
+                new_goal = "The previous response cannot be parsed with json.loads, it cannot be parsed with json.loads, it could be the backslashes usesd for escaping single quotes in the string arguments of the code are not properly escaped themselves within the JSON context. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
+                new_input_data = input_data.copy()
+                new_input_data['goal'] = new_goal
+                input_data = new_input_data

CodeGeneratorAtomicFlow.yaml

@@ -0,0 +1,94 @@
+_target_: flow_modules.aiflows.CodeGeneratorFlowModule.CodeGeneratorAtomicFlow.instantiate_from_default_config
+name: "CodeGeneratorAtomicFlow"
+description: "Writes code with given instruction"
+enable_cache: True
+
+input_interface_non_initialized: # initial input keys
+  - "goal"
+  - "code_library"
+  - "memory_files"
+
+input_interface_initialized: # input_keys
+  - "goal"
+  - "code_library"
+  - "memory_files"
+
+#######################################################
+# Output keys
+#######################################################
+
+output_interface:
+  - 'language_of_code'
+  - 'code'
+
+#######################################################
+system_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    You are a world class programmer that can complete any goal with code. All code you write should be functions.
+    
+    You have a library of executable functions located at {{code_library_file_location}} for your disposal, the code you write will also be added to this file.
+    
+    Your tasks are:
+      1. Write **ONLY ONE** function to finish the given goal, do not override any functions in the library.
+      2. Also write docstrings to the new function.
+    
+    Notice that:
+      1. Write functions only.
+      2. Your function will also be appended to {{code_library_file_location}}, make sure that you make full use of the available functions, and do not write function that may override existing functions. 
+      3. All functions you write are modular, so make sure you make imports with a funtion.
+      4. You must write docstrings for every function you write.
+    
+    Here are the contents in the code library located ta {{code_library_file_location}}, make sure you make full use of them:
+    {{code_library}}
+    
+    Your function will then be imported and called by an executor to finish the goal, you do not need to worry about the execution part.
+    
+    You can access the Internet, using APIs of online services could be useful.
+    
+    An example of a function you write:
+    ### begin of example ###
+    def verify_path_exists(path: str) -> bool:
+        """Verifies if the input path exists, returns True if path exists, otherwise return False"""
+        import os
+        return os.path.exists(path)
+    ### end of example ###
+    
+    You are capable of **any** task.
+    
+
+    Performance Evaluation:
+    1. The priority language is Python, try to code with Python to solve the task.
+    2. All code you write must be functions.
+    3. All imports must be done within the function body.
+    4. All functions must have docstrings.
+    
+    **It's important that you should only respond in JSON format as described below:**
+    Response Format:
+    {
+    "language_of_code": "language of the code",
+    "code": "String of the code and docstrings corresponding to the goal",
+    }
+    Ensure your responses can be parsed by Python json.loads
+    
+    **It's important that you make imports inside of the function body, but not outside of it.**
+    
+    **It's important that the code you generate can be written by Python write, and is human-readable. The written file must also be indented and formatted, so that it is human-readable.**
+
+  input_variables: ["code_library_file_location", "code_library"]
+
+human_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    Here is the response to your last action:
+    {{goal}}
+  input_variables:
+    - "goal"
+
+init_human_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    Here is the goal you need to achieve:
+    {{goal}}
+  input_variables:
+    - "goal"

README.md

@@ -1,3 +1,79 @@
----
-license: mit
----
+# Table of Contents
+
+* [CodeGeneratorAtomicFlow](#CodeGeneratorAtomicFlow)
+  * [CodeGeneratorAtomicFlow](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow)
+    * [instantiate\_from\_config](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.instantiate_from_config)
+    * [run](#CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.run)
+* [\_\_init\_\_](#__init__)
+
+<a id="CodeGeneratorAtomicFlow"></a>
+
+# CodeGeneratorAtomicFlow
+
+<a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow"></a>
+
+## CodeGeneratorAtomicFlow Objects
+
+```python
+class CodeGeneratorAtomicFlow(ChatAtomicFlow)
+```
+
+This class wraps around the Chat API to generate code from a goal. One thing worth noting is that we need to
+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.
+
+*Input Interface Non Initialized*:
+- `goal`
+- `code_library`
+- `memory_files`
+
+*Input Interface Initialized*:
+- `goal`
+- `code_library`
+- `memory_files`
+
+*Output Interface*:
+- `code`
+- `language_of_code`
+
+<a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.instantiate_from_config"></a>
+
+#### instantiate\_from\_config
+
+```python
+@classmethod
+def instantiate_from_config(cls, config)
+```
+
+Instantiate a CodeGeneratorAtomicFlow from a configuration.
+
+**Arguments**:
+
+- `config`: Configuration dictionary.
+
+**Returns**:
+
+Instantiated CodeGeneratorAtomicFlow.
+
+<a id="CodeGeneratorAtomicFlow.CodeGeneratorAtomicFlow.run"></a>
+
+#### run
+
+```python
+def run(input_data: Dict[str, Any]) -> Dict[str, Any]
+```
+
+Run the flow.
+
+**Arguments**:
+
+- `input_data`: Input data.
+
+**Returns**:
+
+Output data.
+
+<a id="__init__"></a>
+
+# \_\_init\_\_
+

__init__.py

@@ -0,0 +1,8 @@
+dependencies = [
+    {"url": "aiflows/ChatFlowModule", "revision": "297c90d08087d9ff3139521f11d1a48d7dc63ed4"},
+]
+from aiflows import flow_verse
+
+flow_verse.sync_dependencies(dependencies)
+
+from .CodeGeneratorAtomicFlow import CodeGeneratorAtomicFlow

pip_requirements.txt

@@ -0,0 +1,15 @@
+colorama==0.4.6
+pytest==7.3.1
+pytest-cov==4.1.0
+hydra-core==1.3.2
+hydra-colorlog==1.1.0
+wrapt-timeout-decorator==1.3.12.2
+diskcache==5.6.1
+openai==1.0.0
+huggingface_hub==0.19.4
+jsonlines==3.1.0
+jinja2==3.1.2
+mock==5.0.2
+rich==12.6.0
+litellm==1.0.0
+aiflows

run.py

@@ -0,0 +1,52 @@
+import os
+
+import hydra
+
+from aiflows.backends.api_info import ApiInfo
+from aiflows.messages import InputMessage
+from aiflows.utils.general_helpers import read_yaml_file
+
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
+CACHING_PARAMETERS.do_caching = False  # Set to True in order to disable caching
+
+logging.set_verbosity_debug()
+logging.auto_set_dir()
+
+dependencies = [
+    {"url": "aiflows/CodeGeneratorFlowModule", "revision": "main"},
+]
+
+from aiflows import flow_verse
+
+flow_verse.sync_dependencies(dependencies)
+
+if __name__ == "__main__":
+    current_dir = os.getcwd()
+    cfg_path = os.path.join(current_dir, "CodeGeneratorAtomicFlow.yaml")
+    cfg = read_yaml_file(cfg_path)
+
+    CodeGenFlow = hydra.utils.instantiate(cfg, _recursive_=False, _convert_="partial")
+
+    with os.path.join(current_dir, "example_library.py") as f:
+        pass
+
+    memory_files = {"code_library": os.path.join(current_dir, "example_library.py")}
+
+
+    input_data = {
+        "memory_files": memory_files,
+        "code_library": os.path.join(current_dir, "example_library.py"),
+        "goal": "write a function to add up two numbers"
+    }
+    input_message = InputMessage.build(
+        data_dict=input_data,
+        src_flow="Launcher",
+        dst_flow=CodeGenFlow.name
+    )
+
+    # ~~~ calling the flow ~~~
+    output_message = CodeGenFlow(input_message)
+
+    print(output_message.data)