| 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 |
| """ |
| def _generate_content(self, plan_file_location, plan_str) -> 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): |
| 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): |
| 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]): |
| 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 arrary 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 prase this format accorrdingly. |
| """ |
| |
| 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] |
| ): |
| 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 |
