| import subprocess | |
| import webbrowser | |
| from typing import Dict, Any | |
| from flows.base_flows.atomic import AtomicFlow | |
| import os | |
| class CodeFileEditAtomicFlow(AtomicFlow): | |
| def _write_code_to_temp_file(self, code_str, language_of_code, code_lib_location): | |
| directory = os.path.dirname(code_lib_location) | |
| if language_of_code.lower() == 'python': | |
| temp_file_location = os.path.join(directory, 'temp.py') | |
| else: | |
| raise NotImplemented | |
| try: | |
| 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" | |
| "# If you prefer not to give any thoughts, just leave it blank, the code will be added to the library as it is " | |
| "in this file.\n" | |
| "###########\n" | |
| "# Code:\n" + | |
| code_str + | |
| "\n############\n" | |
| "# Thoughts:" | |
| ) | |
| with open(temp_file_location, "w") as file: | |
| file.write(content) | |
| file_written_timestamp = os.path.getmtime(temp_file_location) | |
| return True, f"Code written to {temp_file_location}", temp_file_location, file_written_timestamp | |
| except Exception as e: | |
| return False, str(e), temp_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" | |
| def run( | |
| self, | |
| input_data: Dict[str, Any] | |
| ): | |
| self._check_input(input_data) | |
| code_str = input_data['code'] | |
| language_of_code = input_data["language_of_code"] | |
| code_lib_location = input_data["memory_files"]["code_library"] | |
| result, code_editor_output, temp_file_location, file_written_timestamp = self._write_code_to_temp_file(code_str, language_of_code, code_lib_location) | |
| 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 | |