PlanFileEditAtomicFlow.py

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


class PlanFileEditAtomicFlow(AtomicFlow):
    """This class is used to write plan to a temp code file, with commented instructions to give information
    to the user.

    *Input Interface*:
    - `plan`: str
    - `plan_file_location`: str

    *Output Interface*:
    - `plan_editor_output`: str
    - `temp_plan_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, plan_file_location, plan_str) -> str:
        """
        This function generates the content to be written to the plan file.
        :param plan_file_location: The location of the plan file.
        :type plan_file_location: str
        :param plan_str: The plan to be written to the plan file.
        :type plan_str: str
        :return: The content to be written to the plan file.
        :rtype: str
        """
        content = (
                "The below plan will be written to " +
                plan_file_location + "\n"
                "Edit the plan directly or provide your thoughts down below if you have any suggestions.\n"
                "When you are done editing plan and providing feedback, save file and close the current VSCode session to continue.\n"
                "###########\n"
                "Plan:\n" +
                plan_str +
                "\n############\n"
                "Thoughts:"
        )
        return content

    def _generate_temp_file_location(self, plan_file_location):
        """
        This function generates the location of the temp plan file.
        :param plan_file_location: The location of the plan file.
        :type plan_file_location: str
        :return: The location of the temp plan file.
        :rtype: str
        """
        directory = os.path.dirname(plan_file_location)
        ret = os.path.join(directory, 'temp_plan.txt')
        return ret

    def _write_plan_content_to_file(self, file_location, content: str):
        """
        This function writes the plan content to the plan file.
        :param file_location: The location of the plan file.
        :type file_location: str
        :param content: The content to be written to the plan file.
        :type content: str
        :return: The result of writing the plan content to the plan file.
        :rtype: Tuple[bool, str, str]
        """
        try:
            with open(file_location, "w") as file:
                file.write(content)

            return True, f"Plan 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: The input data.
        :type input_data: Dict[str, Any]
        :return: None
        :rtype: None
        :raises AssertionError: If the input data is invalid.
        """
        assert any(item in input_data for item in ["plan", "new_plan"]), "plan or new_plan is not passed to PlanFileEditAtomicFlow"
        assert "plan_file_location" in input_data, "plan_file_location not passed to PlanFileEditAtomicFlow"
        plan_file_loc = input_data["plan_file_location"]
        assert os.path.exists(plan_file_loc), f"{plan_file_loc} does not exist"
        assert os.path.isfile(plan_file_loc), f"{plan_file_loc} is not a file"

    def _generate_input_to_writer(self, input_data: Dict[str, Any]):
        """
        sometimes the plan generator will return an array of indexed plans, like
        [
            "1. Extend the code library with a function named 'import_libraries'. This function should import necessary libraries for the task...",
            "2. Extend the code library with a function named 'fetch_stock_prices'. This function should take two inputs: 'company_code' and 'duration'...",
            "3. Investigate the issue with importing the 'fetch_stock_prices' function from the library..."
        ]
        In this case we need to parse this format accordingly.

        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: The content to be written to the plan file and the location of the temp plan file.
        :rtype: Tuple[str, str]
        """
        
        plan = input_data['plan'] if "plan" in input_data else input_data['new_plan']
        if isinstance(plan, str):
            plan_str = plan
        elif isinstance(plan, list):
            plan_str = "\n".join(plan)
        else:
             raise TypeError("plan is neither a string nor a list")
        plan_file_location = input_data["plan_file_location"]
        content_to_write = self._generate_content(plan_file_location, plan_str)
        file_location_to_write = self._generate_temp_file_location(plan_file_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.
        :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, plan_editor_output, temp_file_location = self._write_plan_content_to_file(
            file_location_to_write, content_to_write)

        # ~~~ Generating return variables ~~~
        response = {}
        response["plan_editor_output"] = plan_editor_output
        response["temp_plan_file_location"] = temp_file_location
        return response