import subprocess
import sys
from smolagents import Tool

class PyCodeInterpreterTool(Tool):
    name = "py_code_interpreter"
    description = """
        Executes file with python code. 
        Strongly use only for answering to the questions from GAIA benchmark. 
        Use default python_interpreter for other cases.
    """
    inputs = {
        "file_path": {
            "type": "string",
            "description": "The file with python code to execute.",
        }
    }
    output_type = "string"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, file_path: str) -> str:
        cmd = f"python {file_path}".split()

        try:
            output = subprocess.check_output(cmd, stderr=subprocess.STDOUT, timeout=60.0)
            output_text = output.decode(sys.stdout.encoding).strip()
        except subprocess.TimeoutExpired:
            output_text = "Error: code execution timed out."
        except subprocess.CalledProcessError as e:
            output_text = e.output.decode(sys.stdout.encoding).strip()
        except subprocess.SubprocessError as ex:
            output_text = f"Error: {str(ex)}\n"

        return output_text