Spaces:
Sleeping
Sleeping
| import runpy | |
| import functools | |
| from pathlib import Path | |
| from typing import Optional | |
| from phi.tools import Toolkit | |
| from phi.utils.log import logger | |
| def warn() -> None: | |
| logger.warning("PythonTools can run arbitrary code, please provide human supervision.") | |
| class PythonTools(Toolkit): | |
| def __init__( | |
| self, | |
| base_dir: Optional[Path] = None, | |
| save_and_run: bool = True, | |
| pip_install: bool = False, | |
| run_code: bool = False, | |
| list_files: bool = False, | |
| run_files: bool = False, | |
| read_files: bool = False, | |
| safe_globals: Optional[dict] = None, | |
| safe_locals: Optional[dict] = None, | |
| ): | |
| super().__init__(name="python_tools") | |
| self.base_dir: Path = base_dir or Path.cwd() | |
| # Restricted global and local scope | |
| self.safe_globals: dict = safe_globals or globals() | |
| self.safe_locals: dict = safe_locals or locals() | |
| if run_code: | |
| self.register(self.run_python_code, sanitize_arguments=False) | |
| if save_and_run: | |
| self.register(self.save_to_file_and_run, sanitize_arguments=False) | |
| if pip_install: | |
| self.register(self.pip_install_package) | |
| if run_files: | |
| self.register(self.run_python_file_return_variable) | |
| if read_files: | |
| self.register(self.read_file) | |
| if list_files: | |
| self.register(self.list_files) | |
| def save_to_file_and_run( | |
| self, file_name: str, code: str, variable_to_return: Optional[str] = None, overwrite: bool = True | |
| ) -> str: | |
| """This function saves Python code to a file called `file_name` and then runs it. | |
| If successful, returns the value of `variable_to_return` if provided otherwise returns a success message. | |
| If failed, returns an error message. | |
| Make sure the file_name ends with `.py` | |
| :param file_name: The name of the file the code will be saved to. | |
| :param code: The code to save and run. | |
| :param variable_to_return: The variable to return. | |
| :param overwrite: Overwrite the file if it already exists. | |
| :return: if run is successful, the value of `variable_to_return` if provided else file name. | |
| """ | |
| try: | |
| warn() | |
| file_path = self.base_dir.joinpath(file_name) | |
| logger.debug(f"Saving code to {file_path}") | |
| if not file_path.parent.exists(): | |
| file_path.parent.mkdir(parents=True, exist_ok=True) | |
| if file_path.exists() and not overwrite: | |
| return f"File {file_name} already exists" | |
| file_path.write_text(code) | |
| logger.info(f"Saved: {file_path}") | |
| logger.info(f"Running {file_path}") | |
| globals_after_run = runpy.run_path(str(file_path), init_globals=self.safe_globals, run_name="__main__") | |
| if variable_to_return: | |
| variable_value = globals_after_run.get(variable_to_return) | |
| if variable_value is None: | |
| return f"Variable {variable_to_return} not found" | |
| logger.debug(f"Variable {variable_to_return} value: {variable_value}") | |
| return str(variable_value) | |
| else: | |
| return f"successfully ran {str(file_path)}" | |
| except Exception as e: | |
| logger.error(f"Error saving and running code: {e}") | |
| return f"Error saving and running code: {e}" | |
| def run_python_file_return_variable(self, file_name: str, variable_to_return: Optional[str] = None) -> str: | |
| """This function runs code in a Python file. | |
| If successful, returns the value of `variable_to_return` if provided otherwise returns a success message. | |
| If failed, returns an error message. | |
| :param file_name: The name of the file to run. | |
| :param variable_to_return: The variable to return. | |
| :return: if run is successful, the value of `variable_to_return` if provided else file name. | |
| """ | |
| try: | |
| warn() | |
| file_path = self.base_dir.joinpath(file_name) | |
| logger.info(f"Running {file_path}") | |
| globals_after_run = runpy.run_path(str(file_path), init_globals=self.safe_globals, run_name="__main__") | |
| if variable_to_return: | |
| variable_value = globals_after_run.get(variable_to_return) | |
| if variable_value is None: | |
| return f"Variable {variable_to_return} not found" | |
| logger.debug(f"Variable {variable_to_return} value: {variable_value}") | |
| return str(variable_value) | |
| else: | |
| return f"successfully ran {str(file_path)}" | |
| except Exception as e: | |
| logger.error(f"Error running file: {e}") | |
| return f"Error running file: {e}" | |
| def read_file(self, file_name: str) -> str: | |
| """Reads the contents of the file `file_name` and returns the contents if successful. | |
| :param file_name: The name of the file to read. | |
| :return: The contents of the file if successful, otherwise returns an error message. | |
| """ | |
| try: | |
| logger.info(f"Reading file: {file_name}") | |
| file_path = self.base_dir.joinpath(file_name) | |
| contents = file_path.read_text() | |
| return str(contents) | |
| except Exception as e: | |
| logger.error(f"Error reading file: {e}") | |
| return f"Error reading file: {e}" | |
| def list_files(self) -> str: | |
| """Returns a list of files in the base directory | |
| :return: Comma separated list of files in the base directory. | |
| """ | |
| try: | |
| logger.info(f"Reading files in : {self.base_dir}") | |
| files = [str(file_path.name) for file_path in self.base_dir.iterdir()] | |
| return ", ".join(files) | |
| except Exception as e: | |
| logger.error(f"Error reading files: {e}") | |
| return f"Error reading files: {e}" | |
| def run_python_code(self, code: str, variable_to_return: Optional[str] = None) -> str: | |
| """This function to runs Python code in the current environment. | |
| If successful, returns the value of `variable_to_return` if provided otherwise returns a success message. | |
| If failed, returns an error message. | |
| Returns the value of `variable_to_return` if successful, otherwise returns an error message. | |
| :param code: The code to run. | |
| :param variable_to_return: The variable to return. | |
| :return: value of `variable_to_return` if successful, otherwise returns an error message. | |
| """ | |
| try: | |
| warn() | |
| logger.debug(f"Running code:\n\n{code}\n\n") | |
| exec(code, self.safe_globals, self.safe_locals) | |
| if variable_to_return: | |
| variable_value = self.safe_locals.get(variable_to_return) | |
| if variable_value is None: | |
| return f"Variable {variable_to_return} not found" | |
| logger.debug(f"Variable {variable_to_return} value: {variable_value}") | |
| return str(variable_value) | |
| else: | |
| return "successfully ran python code" | |
| except Exception as e: | |
| logger.error(f"Error running python code: {e}") | |
| return f"Error running python code: {e}" | |
| def pip_install_package(self, package_name: str) -> str: | |
| """This function installs a package using pip in the current environment. | |
| If successful, returns a success message. | |
| If failed, returns an error message. | |
| :param package_name: The name of the package to install. | |
| :return: success message if successful, otherwise returns an error message. | |
| """ | |
| try: | |
| warn() | |
| logger.debug(f"Installing package {package_name}") | |
| import sys | |
| import subprocess | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", package_name]) | |
| return f"successfully installed package {package_name}" | |
| except Exception as e: | |
| logger.error(f"Error installing package {package_name}: {e}") | |
| return f"Error installing package {package_name}: {e}" | |