gaia_agent/tools/python_tool.py

"""A restricted Python REPL for exact arithmetic, string, and date work."""

from __future__ import annotations

import contextlib
import io
import math

from langchain_core.tools import tool

# A small, deliberately limited set of safe builtins/modules.
_SAFE_BUILTINS = {
    "abs": abs, "all": all, "any": any, "bin": bin, "bool": bool, "chr": chr,
    "dict": dict, "divmod": divmod, "enumerate": enumerate, "filter": filter,
    "float": float, "hex": hex, "int": int, "len": len, "list": list, "map": map,
    "max": max, "min": min, "oct": oct, "ord": ord, "pow": pow, "range": range,
    "reversed": reversed, "round": round, "set": set, "sorted": sorted, "str": str,
    "sum": sum, "tuple": tuple, "zip": zip, "print": print,
}


@tool
def python_repl(code: str) -> str:
    """Execute a short Python snippet and return its stdout (and `result` if set).

    Use for exact arithmetic, string manipulation, sorting, and date math. ALWAYS use
    this for numeric work instead of computing in your head. ``math``, ``statistics``,
    ``datetime``, ``re``, ``itertools``, ``collections`` and ``pd`` (pandas) are
    available -- e.g. read a spreadsheet with ``pd.read_excel(path)``. Assign to a
    variable named ``result`` to return a value.

    Args:
        code: Python source to execute.
    """
    import collections
    import datetime
    import itertools
    import re
    import statistics

    env = {
        "__builtins__": _SAFE_BUILTINS,
        "math": math,
        "statistics": statistics,
        "datetime": datetime,
        "re": re,
        "itertools": itertools,
        "collections": collections,
    }
    # Preload pandas as `pd` so spreadsheet math is exact, e.g.
    #   df = pd.read_excel(path); result = df[["Burgers","Fries"]].sum().sum()
    try:
        import pandas as pd

        env["pd"] = pd
    except Exception:  # noqa: BLE001
        pass
    buf = io.StringIO()
    try:
        with contextlib.redirect_stdout(buf):
            exec(code, env)  # noqa: S102 - sandboxed builtins, internal use only
    except Exception as exc:  # noqa: BLE001
        return f"Error: {exc}\nOutput so far:\n{buf.getvalue()}"

    out = buf.getvalue()
    if "result" in env:
        out += ("\n" if out else "") + f"result = {env['result']!r}"
    return out or "(no output)"


@tool
def run_python_file(path: str) -> str:
    """Execute a downloaded .py file in a subprocess and return its stdout.

    Use for questions asking for the output of an attached Python script.

    Args:
        path: Local path to the .py file (e.g. from download_task_file).
    """
    import subprocess
    import sys

    try:
        proc = subprocess.run(
            [sys.executable, path],
            capture_output=True,
            text=True,
            timeout=30,
        )
    except subprocess.TimeoutExpired:
        return "Execution timed out after 30s."
    except Exception as exc:  # noqa: BLE001
        return f"Could not run file: {exc}"
    out = (proc.stdout or "").strip()
    err = (proc.stderr or "").strip()
    result = f"STDOUT:\n{out}"
    if err:
        result += f"\nSTDERR:\n{err[:2000]}"
    return result or "(no output)"