101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
"""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"}
|
|
|
|
|
|
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.
|
|
|
|
Available: pandas, numpy, yfinance (data), matplotlib (charts, Agg), scipy. Use this like a
|
|
real analyst: fetch data with yfinance, compute indicators/stats/backtests, and save charts
|
|
with `plt.savefig("chart.png")`. You MUST print() anything you want to read back. Runs in an
|
|
isolated temp dir with a ~90s limit and NO access to account secrets.
|
|
|
|
Example:
|
|
import yfinance as yf, matplotlib.pyplot as plt
|
|
df = yf.download("AAPL", period="6mo")
|
|
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'}"
|
|
result = out[:_STDOUT_LIMIT] or "(no stdout — did you print()?)"
|
|
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
|