"""Analyst script sandbox — run_analysis. Lets an agent write and run a real Python analysis script (pandas/numpy/yfinance/matplotlib) like a research analyst: fetch data, compute indicators/stats/backtests, save charts. Runs in a throwaway temp workdir with a **scrubbed environment** (no ALPACA_/FINNHUB_/OMLX_ secrets — model-generated code must never see keys), a wall-clock timeout, and network access (for data). Security posture (honest): running model-written networked code is a real surface. Mitigations: scrubbed env (no secrets), isolated temp workdir (no repo writes), wall-clock timeout, truncated output. Acceptable on a personal machine + local model; not a hardened multi-tenant sandbox. """ from __future__ import annotations import os import subprocess import sys import tempfile from pathlib import Path from langchain_core.tools import tool from config.settings import get_settings _STDOUT_LIMIT = 4000 _SECRET_PREFIXES = ("ALPACA", "FINNHUB", "OMLX", "EMAIL", "SMTP", "OPENAI", "ANTHROPIC") _ARTIFACT_EXTS = {".png", ".jpg", ".jpeg", ".svg", ".csv", ".txt", ".html", ".md", ".json"} _AVAILABLE = "pandas, numpy, yfinance, matplotlib (Agg), scipy + the standard library" # Common failure -> actionable hint, appended to the ERROR the agent reads back. Turns the # usual retry-loop ("same error again") into a one-shot self-correction. _ERROR_HINTS: list[tuple[str, str]] = [ ("ModuleNotFoundError", f"HINT: only {_AVAILABLE} are installed — rewrite using those."), ( "KeyError", "HINT: yf.download() returns MultiIndex columns by default. Flatten first: " "df.columns = df.columns.get_level_values(0) — or download with " "yf.download(ticker, period=..., multi_level_index=False).", ), ( "FileNotFoundError", "HINT: each run_analysis call starts in a FRESH empty directory — files from earlier " "calls are gone. Make the script self-contained (re-download data in this script).", ), ( "NameError", "HINT: nothing persists between run_analysis calls — define every variable and import " "in THIS script.", ), ( "JSONDecodeError", "HINT: a data request likely returned empty/error content — print the raw response " "first, and add a fallback ticker/period.", ), ] def _hint_for(stderr: str) -> str: for needle, hint in _ERROR_HINTS: if needle in stderr: return "\n" + hint if "No data found" in stderr or "no price data" in stderr.lower(): return "\nHINT: yfinance returned no data — check the ticker symbol and use a simple period like '6mo'." return "" def scrubbed_env(home: str) -> dict[str, str]: """A minimal environment with NO secrets — built from scratch, not inherited.""" env = { "PATH": os.environ.get("PATH", "/usr/bin:/bin"), "HOME": home, "TMPDIR": home, "LANG": os.environ.get("LANG", "en_US.UTF-8"), "MPLBACKEND": "Agg", # headless matplotlib "MPLCONFIGDIR": home, } # defensive: never leak anything secret-looking even if the whitelist grows return {k: v for k, v in env.items() if not k.upper().startswith(_SECRET_PREFIXES)} @tool def run_analysis(code: str) -> str: """Run a Python analysis script and return its stdout plus any files it wrote. ENVIRONMENT — read once, avoid the classic errors: - Installed: pandas, numpy, yfinance (market data), matplotlib (charts, headless Agg), scipy, plus the standard library. Nothing else — no requests/sklearn/talib. - Each call runs in a FRESH EMPTY temp directory. Files and variables do NOT persist between calls: every script must be fully self-contained (imports + data download + computation in one script). - You MUST print() every value you want to read back; save charts with plt.savefig("name.png") (saved files are listed in the reply). - ~90s wall-clock limit; no access to account secrets. - yfinance pitfalls: yf.download() returns MultiIndex columns — flatten with `df.columns = df.columns.get_level_values(0)` (or pass multi_level_index=False); it also auto-adjusts prices by default (auto_adjust=True). Use it like a real analyst: fetch data, compute indicators/stats/backtests, save charts. Example: import yfinance as yf, matplotlib.pyplot as plt df = yf.download("AAPL", period="6mo", multi_level_index=False) df["SMA50"] = df["Close"].rolling(50).mean() print(df[["Close","SMA50"]].tail()) df[["Close","SMA50"]].plot(); plt.savefig("aapl.png") """ s = get_settings() s.ensure_dirs() workdir = Path(tempfile.mkdtemp(dir=s.sandbox_dir)) (workdir / "script.py").write_text(code) # shared, persistent matplotlib cache so the font cache is built once (not per run) and # doesn't land in each run's artifact list mpl_cache = s.sandbox_dir / ".mplconfig" mpl_cache.mkdir(parents=True, exist_ok=True) env = scrubbed_env(str(workdir)) env["MPLCONFIGDIR"] = str(mpl_cache) try: proc = subprocess.run( [sys.executable, "script.py"], cwd=workdir, env=env, capture_output=True, text=True, timeout=s.sandbox_timeout_s, ) except subprocess.TimeoutExpired: return f"ERROR: analysis timed out ({s.sandbox_timeout_s}s limit)." except Exception as e: # pragma: no cover - defensive return f"ERROR: {type(e).__name__}: {e}" out = (proc.stdout or "").strip() err = (proc.stderr or "").strip() artifacts = [ str(p) for p in sorted(workdir.iterdir()) if p.is_file() and p.name != "script.py" and p.suffix.lower() in _ARTIFACT_EXTS ] if proc.returncode != 0: return f"ERROR (exit {proc.returncode}): {err[-1500:] or 'no stderr'}{_hint_for(err)}" result = out[:_STDOUT_LIMIT] or ( "(no stdout — you must print() every value you want to read back)" ) if len(out) > _STDOUT_LIMIT: result += "\n...[truncated]" if err: result += f"\n[stderr] {err[-400:]}" if artifacts: result += "\nArtifacts written: " + ", ".join(artifacts) return result