| 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) |
|
|
| |
| content_to_write, file_location_to_write = self._generate_input_to_writer(input_data) |
|
|
| |
| result, plan_editor_output, temp_file_location = self._write_plan_content_to_file( |
| file_location_to_write, content_to_write) |
|
|
| |
| response = {} |
| response["plan_editor_output"] = plan_editor_output |
| response["temp_plan_file_location"] = temp_file_location |
| return response |
