AmmarFahmy
adding all files
105b369
import runpy
import functools
from pathlib import Path
from typing import Optional
from phi.tools import Toolkit
from phi.utils.log import logger
@functools.lru_cache(maxsize=None)
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}"