src/tools/code_interpreter.py

"""
CodeInterpreterTool — local sandboxed Python execution.
Replaces the crewai_tools.CodeInterpreterTool removed in crewai v1.x.
Runs Python code in a subprocess with a strict timeout and whitelist.
"""
from __future__ import annotations

import json
import subprocess
import sys
import textwrap
from pathlib import Path

from crewai.tools import BaseTool
from loguru import logger


# ─── Allowlist of safe imports ────────────────────────────────────────────────
_ALLOWED_IMPORTS = frozenset(
    [
        "pandas", "numpy", "math", "statistics", "json", "csv", "io",
        "datetime", "re", "collections", "itertools", "functools",
        "matplotlib", "matplotlib.pyplot",
    ]
)

_BLOCKED_KEYWORDS = frozenset(
    [
        "__import__", "exec", "eval", "compile", "open", "os", "sys",
        "subprocess", "socket", "urllib", "requests", "shutil", "glob",
        "importlib", "ctypes", "pickle", "marshal",
    ]
)


def _is_safe_code(code: str) -> tuple[bool, str]:
    """Basic static check — not a full sandbox, just a sanity layer."""
    for kw in _BLOCKED_KEYWORDS:
        if kw in code:
            return False, f"Blocked keyword in code: '{kw}'"
    return True, ""


class CodeInterpreterTool(BaseTool):
    """
    Sandboxed Python code execution tool for financial calculations.
    Implements the interface of the deprecated crewai_tools.CodeInterpreterTool.
    Accepts Python code as a string. Returns stdout + last expression result.
    """

    name: str = "code_interpreter"
    description: str = (
        "Execute Python code for financial calculations. "
        "Supports pandas, numpy, math, statistics, datetime. "
        "Input: a string of Python code. "
        "Output: stdout output and/or the result of the last expression. "
        "Use for: revenue projections, unit economics, CAC/LTV, burn rate, market sizing."
    )

    timeout_seconds: int = 30

    def _run(self, code: str) -> str:
        """Execute Python code in a subprocess and return the output."""
        if not code or not code.strip():
            return "Error: no code provided."

        safe, reason = _is_safe_code(code)
        if not safe:
            return f"Error: unsafe code detected — {reason}"

        # Wrap in a try/except and auto-print last expression
        wrapped = textwrap.dedent(f"""
import pandas as pd
import numpy as np
import math
import statistics
import json
import csv
import io
import datetime
import re
import collections
import itertools
import functools

try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except Exception:
    pass

# ─── User code ────────────────────────
{code}
""").strip()

        try:
            result = subprocess.run(
                [sys.executable, "-c", wrapped],
                capture_output=True,
                text=True,
                timeout=self.timeout_seconds,
            )
            output = result.stdout.strip()
            err = result.stderr.strip()

            if result.returncode != 0:
                # Surface the traceback but strip local paths
                safe_err = "\n".join(
                    ln for ln in err.splitlines()
                    if "/home/" not in ln and "/tmp/" not in ln
                )
                return f"Code execution error:\n{safe_err}" if safe_err else f"Exit code {result.returncode}"

            if output:
                return output
            return "(code ran successfully, no output)"

        except subprocess.TimeoutExpired:
            return f"Error: code execution timed out after {self.timeout_seconds}s."
        except Exception as exc:
            logger.error(f"CodeInterpreterTool: {exc}")
            return f"Error: {exc}"