Spaces:
Running on Zero
Running on Zero
File size: 6,427 Bytes
9936912 48ee375 9936912 48ee375 9936912 48ee375 9936912 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """Jupyter-style deterministic Python execution tool for Control-LLM."""
from __future__ import annotations
import contextlib
import hashlib
import io
import sys
import time
from pathlib import Path
from typing import Any
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import scipy
from scipy import linalg, signal
try:
import control as ct
except ImportError:
ct = None
from controlai_agent.registry import registry
PLOTS_DIR = Path("outputs/plots")
PLOTS_DIR.mkdir(parents=True, exist_ok=True)
@registry.register(
name="execute_python_code",
description=(
"Execute Python code for control engineering, numerical simulations, differential equations, "
"optimization, and signal plotting (similar to a Jupyter Notebook cell). Can use numpy (np), "
"scipy (scipy), scipy.signal (signal), scipy.linalg (linalg), control (ct, control), and "
"matplotlib.pyplot (plt). Captures stdout and any generated Matplotlib figures. "
"IMPORTANT -- the `control` package (ct) takes POSITIONAL arguments only, never num=/den= "
"or sys1=/sys2= keywords (they raise 'Needs 1, 2, or 3 arguments'): "
"ct.tf(num, den) not ct.tf(num=num, den=den); "
"ct.feedback(sys1, sys2=1, sign=-1) for a closed loop; "
"ct.series(sys1, sys2) and ct.parallel(sys1, sys2); "
"ct.step_response(sys, T=t_array) returns (T, yout); "
"ct.poles(sys) and ct.zeros(sys) for pole/zero locations; "
"ct.bode(sys) / ct.bode_plot(sys) is PLOT-ONLY and does not return (mag, phase, omega) arrays "
"-- for numeric Bode data use resp = ct.frequency_response(sys, omega); "
"resp.magnitude, resp.phase (radians), resp.omega. "
"Every other registered tool (continuous_lqr, discrete_lqr, place_state_feedback, "
"stability_margins, exact_zoh, etc.) is ALSO directly callable here by its exact name with "
"its normal arguments -- e.g. `result = place_state_feedback(A=A, B=B, desired_poles=poles)`. "
"Each returns a dict just like the standalone tool call does, so pull out the field you need, "
"e.g. `K = np.array(result['K'])`, before using it in further computation."
),
parameters_schema={
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Valid Python code to execute. Can use numpy, scipy, control, matplotlib.pyplot, etc.",
},
},
"required": ["code"],
},
)
def execute_python_code(code: str) -> dict[str, Any]:
"""Execute Python code in an isolated namespace and capture stdout + plots."""
# Setup execution environment
stdout_capture = io.StringIO()
plt.close("all") # Reset existing figures
# Configure plot defaults for dark/clean aesthetic
plt.style.use("dark_background")
plt.rcParams["figure.facecolor"] = "#151b23"
plt.rcParams["axes.facecolor"] = "#0f1217"
plt.rcParams["axes.edgecolor"] = "#30363d"
plt.rcParams["axes.labelcolor"] = "#8b949e"
plt.rcParams["xtick.color"] = "#8b949e"
plt.rcParams["ytick.color"] = "#8b949e"
plt.rcParams["grid.color"] = "#30363d"
plt.rcParams["grid.linestyle"] = ":"
plt.rcParams["font.sans-serif"] = ["DejaVu Sans", "Helvetica", "Arial"]
exec_globals: dict[str, Any] = {
"np": np,
"numpy": np,
"scipy": scipy,
"linalg": linalg,
"signal": signal,
"plt": plt,
"matplotlib": matplotlib,
"ct": ct,
"control": ct,
}
# Every other registered deterministic tool (continuous_lqr, place_state_feedback,
# stability_margins, ...) is also callable directly by name here, with its
# normal keyword arguments and dict return value -- the model otherwise
# reasonably expects a tool it knows to be usable in code it writes, not
# only through the separate tool-call protocol, and hits a NameError.
exec_globals.update(registry.get_callables(exclude={"execute_python_code"}))
saved_plots = []
# Intercept any direct plt.savefig / fig.savefig calls so they are redirected into outputs/plots
orig_plt_savefig = plt.savefig
orig_fig_savefig = matplotlib.figure.Figure.savefig
def safe_plt_savefig(fname, *args, **kwargs):
fname_name = Path(fname).name if fname else f"plot_{int(time.time()*1000)}.png"
target_path = PLOTS_DIR / fname_name
saved_plots.append(str(target_path))
return orig_plt_savefig(str(target_path), *args, **kwargs)
def safe_fig_savefig(self, fname, *args, **kwargs):
fname_name = Path(fname).name if fname else f"plot_{int(time.time()*1000)}.png"
target_path = PLOTS_DIR / fname_name
saved_plots.append(str(target_path))
return orig_fig_savefig(self, str(target_path), *args, **kwargs)
plt.savefig = safe_plt_savefig
matplotlib.figure.Figure.savefig = safe_fig_savefig
try:
with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stdout_capture):
exec(code, exec_globals)
stdout_val = stdout_capture.getvalue().strip()
# Check if any matplotlib figures remain unsaved
fig_nums = plt.get_fignums()
if fig_nums:
for num in fig_nums:
fig = plt.figure(num)
hash_val = hashlib.md5(f"{code}_{num}_{time.time()}".encode()).hexdigest()[:8]
out_path = PLOTS_DIR / f"plot_exec_{hash_val}.png"
orig_fig_savefig(fig, str(out_path), dpi=130, bbox_inches="tight", facecolor=fig.get_facecolor())
saved_plots.append(str(out_path))
plt.close("all")
result = {
"status": "success",
"stdout": stdout_val if stdout_val else "Code executed successfully.",
"executed_code": code,
}
if saved_plots:
result["plot_path"] = saved_plots[0]
result["all_plots"] = list(dict.fromkeys(saved_plots))
return result
except Exception as exc:
plt.close("all")
return {
"status": "error",
"error": str(exc),
"stdout": stdout_capture.getvalue().strip(),
"executed_code": code,
}
finally:
plt.savefig = orig_plt_savefig
matplotlib.figure.Figure.savefig = orig_fig_savefig
|