CodeFileEditAtomicFlow.py

from typing import Dict, Any
from aiflows.base_flows.atomic import AtomicFlow
import os


class CodeFileEditAtomicFlow(AtomicFlow):
    """This class is used to write code to a temp code file, with commented instructions to give information
    to the user.
    
    *Input Interface*:
    - `code`: str
    - `language_of_code`: str
    - `memory_files`: Dict[str, str]
    
    *Output Interface*:
    - `code_editor_output`: str
    - `temp_code_file_location`: str

    *Configuration Parameters*:
    - `input_interface`: The input interface of the atomic flow.
    - `output_interface`: The output interface of the atomic flow.

    """
    def _generate_content(self, code_lib_location, code_str) -> str:
        """This function generates the content to be written to the temp file.
        :param code_lib_location (str): The location of the code library file.
        :param code_str (str): The code to be written to the temp file.
        :return: (str) The content to be written to the temp file.
        """
        content = (
                "# The below code will be appended to " +
                code_lib_location + "\n"
                                    "# Edit the code directly or provide your thoughts down below if you have any suggestions.\n"
                                    "# When you are done editing code and providing feedback, save file and close the current VSCode session to continue. \n"
                                    "###########\n"
                                    "# Code:\n" +
                code_str +
                "\n############\n"
                "# Thoughts:"
        )
        return content

    def _generate_temp_file_location(self, code_lib_location):
        """This function generates the location of the temp file.
        :param code_lib_location (str): The location of the code library file.
        :return: (str) The location of the temp file.
        """
        directory = os.path.dirname(code_lib_location)
        ret = os.path.join(directory, 'temp.py')
        return ret

    def _write_code_content_to_file(self, file_location, content: str):
        """This function writes the content to the temp file.
        :param file_location (str): The location of the temp file.
        :param content (str): The content to be written to the temp file.
        :return: (bool, str, str) The first value is a boolean indicating if the write was successful.
        The second value is a string containing the output of the write operation.
        The third value is a string containing the location of the temp file.
        """
        try:
            with open(file_location, "w") as file:
                file.write(content)

            return True, f"Code written to {file_location}", file_location

        except Exception as e:
            return False, str(e), file_location

    def _check_input(self, input_data: Dict[str, Any]):
        """This function checks if the input data is valid.
        :param input_data (Dict[str, Any]): The input data.
        :return: None
        """
        assert "code" in input_data, "code is not passed to CodeFileEditAtomicFlow"
        assert "language_of_code" in input_data, "language_of_code is not passed to CodeFileEditAtomicFlow"
        assert "memory_files" in input_data, "memory_files is not passed to CodeFileEditAtomicFlow"
        assert "code_library" in input_data["memory_files"], "code_library not in memory files"
        code_lib_loc = input_data["memory_files"]["code_library"]
        assert os.path.exists(code_lib_loc), f"{code_lib_loc} does not exist"
        assert os.path.isfile(code_lib_loc), f"{code_lib_loc} is not a file"
        language_of_code = input_data["language_of_code"]
        assert language_of_code.lower() == 'python', "sorry!! only writing python code is supported."

    def _generate_input_to_writer(self, input_data: Dict[str, Any]):
        """This function generates the input to the writer function.
        :param input_data (Dict[str, Any]): The input data.
        :return: (str, str) The first value is a string containing the content to be written to the temp file.
        The second value is a string containing the location of the temp file.
        """
        code_str = input_data['code']
        code_lib_location = input_data["memory_files"]["code_library"]
        content_to_write = self._generate_content(code_lib_location, code_str)
        file_location_to_write = self._generate_temp_file_location(code_lib_location)
        return content_to_write, file_location_to_write

    def run(
            self,
            input_data: Dict[str, Any]
    ):
        """This function runs the atomic flow.

        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: The output data containing code editor output and temporary file location.
        :rtype: Dict[str, Any]
        """
        self._check_input(input_data)

        # ~~~ Getting input data to the file editor ~~~
        content_to_write, file_location_to_write = self._generate_input_to_writer(input_data)

        # ~~~ Calling the writer function ~~~
        result, code_editor_output, temp_file_location = self._write_code_content_to_file(
            file_location_to_write, content_to_write)

        # ~~~ Generating return variables ~~~
        response = {}
        response["code_editor_output"] = code_editor_output
        response["temp_code_file_location"] = temp_file_location
        return response