python_interpreter_tool.py

import io
import sys
import platform
import faulthandler
from contextlib import contextmanager
import signal

from transformers import Tool


class PythonInterpreter(Tool):
    name = "python_interpreter_tool"
    description = "Executes Python code and returns results. Requires HF_ALLOW_CODE_EVAL = '1'"

    inputs = ["text"]
    outputs = ["text"]
    timeout = 10.0

    def __call__(self, task: str):
        import os
        import shutil

        if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1":
            return "Can't execute code, set code evaluation flag for PythonInterpreter"

        rmtree = shutil.rmtree
        rmdir = os.rmdir
        chdir = os.chdir

        reliability_guard()
        try:
            exec_globals = {}
            with capture_stdout() as output:
                with time_limit(self.timeout):
                    exec(task, exec_globals)
            
            captured_output = output.getvalue().strip()
        except Exception as e:
            captured_output = str(e)

        # Needed for cleaning up.
        shutil.rmtree = rmtree
        os.rmdir = rmdir
        os.chdir = chdir

        return captured_output


@contextmanager
def capture_stdout():
    output_buffer = io.StringIO()
    original_stdout = sys.stdout
    try:
        sys.stdout = output_buffer
        yield output_buffer
    finally:
        sys.stdout = original_stdout


@contextmanager
def time_limit(seconds):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")

    signal.setitimer(signal.ITIMER_REAL, seconds)
    signal.signal(signal.SIGALRM, signal_handler)
    try:
        yield
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)


def reliability_guard(maximum_memory_bytes=None):
    """
    This disables various destructive functions and prevents the generated code
    from interfering with the test (e.g. fork bomb, killing other processes,
    removing filesystem files, etc.)

    WARNING
    This function is NOT a security sandbox. Untrusted code, including, model-
    generated code, should not be blindly executed outside of one. See the
    Codex paper for more information about OpenAI's code sandbox, and proceed
    with caution.
    """

    if maximum_memory_bytes is not None:
        import resource

        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
        if not platform.uname().system == "Darwin":
            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))

    faulthandler.disable()

    import builtins

    builtins.exit = None
    builtins.quit = None

    import os

    os.environ["OMP_NUM_THREADS"] = "1"

    os.kill = None
    os.system = None
    os.putenv = None
    os.remove = None
    os.removedirs = None
    os.rmdir = None
    os.fchdir = None
    os.setuid = None
    os.fork = None
    os.forkpty = None
    os.killpg = None
    os.rename = None
    os.renames = None
    os.truncate = None
    os.replace = None
    os.unlink = None
    os.fchmod = None
    os.fchown = None
    os.chmod = None
    os.chown = None
    os.chroot = None
    os.fchdir = None
    os.lchflags = None
    os.lchmod = None
    os.lchown = None
    os.getcwd = None
    os.chdir = None

    import shutil

    shutil.rmtree = None
    shutil.move = None
    shutil.chown = None

    import subprocess

    subprocess.Popen = None  # type: ignore

    #__builtins__["help"] = None

    import sys

    sys.modules["ipdb"] = None
    sys.modules["joblib"] = None
    sys.modules["resource"] = None
    sys.modules["psutil"] = None
    sys.modules["tkinter"] = None

class TimeoutException(Exception):
    pass

"""
import os
tool = PythonInterpreter()
print(tool("import os; os.getcwd()"))
print(os.getcwd())
"""