from langchain.tools import BaseTool
from typing import Dict


class PythonExecutionTool(BaseTool):
    name = "PythonExecutor"
    description = "Executes Python code and returns the code and its output. "

    def _run(self, code: str) -> Dict[str, str]:
        return "Code: \n" + code
        try:
            local_vars = {}
            exec(code, {}, local_vars)
            output = local_vars["result"]

            return "Code: \n" + code + " \n Code output: " + str(output)
        except Exception:
            return "Code: \n" + code
