| import subprocess |
| import webbrowser |
| from typing import Dict, Any |
| from flows.base_flows.atomic import AtomicFlow |
| import os |
|
|
|
|
| class CodeFileEditAtomicFlow(AtomicFlow): |
| def _generate_content(self, code_lib_location, code_str) -> str: |
| 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" |
| "# Please do provide your feedback in the thoughts section so that JARVIS knows what you " |
| "are doing.\n" |
| "###########\n" |
| "# Code:\n" + |
| code_str + |
| "\n############\n" |
| "# Thoughts:" |
| ) |
| return content |
|
|
| def _generate_temp_file_location(self, code_lib_location): |
| 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): |
| try: |
| with open(file_location, "w") as file: |
| file.write(content) |
| file_written_timestamp = os.path.getmtime(file_location) |
|
|
| return True, f"Code written to {file_location}", file_location, file_written_timestamp |
|
|
| except Exception as e: |
| return False, str(e), file_location, 0 |
|
|
| def _check_input(self, input_data: Dict[str, Any]): |
| 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]): |
| 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] |
| ): |
| self._check_input(input_data) |
|
|
| |
| content_to_write, file_location_to_write = self._generate_input_to_writer(input_data) |
|
|
| |
| result, code_editor_output, temp_file_location, file_written_timestamp = self._write_code_content_to_file( |
| file_location_to_write, content_to_write) |
|
|
| |
| if result: |
| try: |
| subprocess.run(["code", temp_file_location], timeout=10) |
| except (subprocess.CalledProcessError, subprocess.TimeoutExpired): |
| webbrowser.open(temp_file_location) |
|
|
| |
| response = {} |
| response["code_editor_output"] = code_editor_output |
| response["temp_code_file_location"] = temp_file_location |
| response["temp_code_file_written_timestamp"] = file_written_timestamp |
| return response |
|
|
