Spaces:
Running
Running
File size: 8,375 Bytes
0084130 | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 | """
Code executor for spatial transcriptomics agent.
Executes generated Python code safely with session management.
"""
from __future__ import annotations
import io
import traceback
import contextlib
import re
from dataclasses import dataclass
from typing import Dict, Any, Optional, List
import logging
logger = logging.getLogger(__name__)
@dataclass
class ExecutionResult:
"""Result of code execution."""
success: bool
stdout: str = ""
stderr: str = ""
error: Optional[str] = None
variables: Dict[str, Any] = None
plots: Optional[List[str]] = None # List of base64 encoded plot images
def get_display_output(self) -> str:
"""Get formatted output for display."""
if not self.success:
return f"Error: {self.error}\n\n{self.stderr}"
output = self.stdout
if self.stderr:
output += f"\n\nWarnings:\n{self.stderr}"
return output or "Code executed successfully (no output)"
class CodeExecutor:
"""
Executes Python code with session management.
Features:
- Safe code execution in controlled namespace
- Session state management (variables persist)
- Capture stdout/stderr
- Error handling and reporting
Parameters
----------
timeout : int
Execution timeout in seconds (default: 60)
safe_mode : bool
Enable safety checks (default: True)
Examples
--------
>>> executor = CodeExecutor()
>>> result = executor.execute("x = 5\\nprint(x * 2)", namespace={"session": session})
>>> print(result.stdout)
10
"""
def __init__(self, timeout: int = 60, safe_mode: bool = True):
self.timeout = timeout
self.safe_mode = safe_mode
self.namespace: Dict[str, Any] = {}
logger.info(f"Initialized code executor (timeout={timeout}s, safe_mode={safe_mode})")
def execute(
self,
code: str,
namespace: Optional[Dict[str, Any]] = None,
capture_output: bool = True
) -> ExecutionResult:
"""
Execute Python code.
Parameters
----------
code : str
Python code to execute
namespace : Optional[Dict]
Namespace dictionary (variables available to code)
capture_output : bool
Whether to capture stdout/stderr
Returns
-------
ExecutionResult
Execution result with output and any errors
"""
# Clean code (remove markdown fences if present)
code = self._clean_code(code)
if not code or not code.strip():
return ExecutionResult(
success=False,
error="Empty code provided"
)
# Safety checks
if self.safe_mode:
safety_check = self._check_code_safety(code)
if not safety_check[0]:
return ExecutionResult(
success=False,
error=f"Safety check failed: {safety_check[1]}"
)
# Prepare namespace
exec_namespace = self.namespace.copy()
if namespace:
exec_namespace.update(namespace)
# Capture output
stdout = io.StringIO()
stderr = io.StringIO()
try:
# Compile code
compiled_code = compile(code, "<spatial_agent>", "exec")
# Execute with output capture
if capture_output:
with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
exec(compiled_code, exec_namespace)
else:
exec(compiled_code, exec_namespace)
# Update persistent namespace (exclude builtins and special vars)
for key, value in exec_namespace.items():
if not key.startswith('_') and key not in ['__builtins__']:
self.namespace[key] = value
# Capture matplotlib figures
plots = self._capture_plots()
result = ExecutionResult(
success=True,
stdout=stdout.getvalue(),
stderr=stderr.getvalue(),
variables=self.namespace.copy(),
plots=plots if plots else None
)
logger.info(f"Code executed successfully ({len(plots) if plots else 0} plots captured)")
return result
except SyntaxError as e:
error_msg = f"Syntax Error: {e}"
logger.error(f"Syntax error in code execution: {e}")
return ExecutionResult(
success=False,
stdout=stdout.getvalue(),
stderr=stderr.getvalue(),
error=error_msg
)
except Exception as e:
error_msg = f"{type(e).__name__}: {e}"
error_trace = traceback.format_exc()
logger.error(f"Error executing code: {e}\n{error_trace}")
return ExecutionResult(
success=False,
stdout=stdout.getvalue(),
stderr=stderr.getvalue() + "\n" + error_trace,
error=error_msg
)
def _clean_code(self, code: str) -> str:
"""Remove markdown code fences from code."""
code = code.strip()
# Remove markdown code blocks
if code.startswith("```"):
lines = code.split("\n")
# Remove first line if it's ```python or ```
if lines[0].strip().startswith("```"):
lines = lines[1:]
# Remove last line if it's ```
if lines and lines[-1].strip() == "```":
lines = lines[:-1]
code = "\n".join(lines)
return code.strip()
def _check_code_safety(self, code: str) -> tuple[bool, Optional[str]]:
"""
Check code for unsafe operations.
Returns (is_safe, error_message)
"""
# Dangerous patterns
dangerous_patterns = [
(r'\bos\.system\b', "os.system() not allowed"),
(r'\beval\b', "eval() not allowed"),
(r'\bexec\b', "nested exec() not allowed"),
(r'\b__import__\b', "dynamic imports not allowed"),
(r'\bopen\([^)]*[\'"]w', "file writing not allowed in safe mode"),
(r'\bsubprocess\b', "subprocess not allowed"),
]
for pattern, message in dangerous_patterns:
if re.search(pattern, code):
return False, message
return True, None
def _capture_plots(self) -> list:
"""
Capture all open matplotlib figures as base64 encoded images.
Returns
-------
list
List of base64 encoded PNG images
"""
plots = []
try:
import matplotlib.pyplot as plt
import base64
# Get all figure numbers
fig_nums = plt.get_fignums()
if not fig_nums:
return plots
for fig_num in fig_nums:
fig = plt.figure(fig_num)
# Save figure to bytes buffer
buf = io.BytesIO()
fig.savefig(buf, format='png', bbox_inches='tight', dpi=100)
buf.seek(0)
# Encode to base64
img_base64 = base64.b64encode(buf.read()).decode('utf-8')
plots.append(img_base64)
# Close figure to free memory
plt.close(fig)
logger.info(f"Captured {len(plots)} matplotlib figures")
except ImportError:
logger.debug("matplotlib not available, skipping plot capture")
except Exception as e:
logger.warning(f"Failed to capture plots: {e}")
return plots
def reset_namespace(self) -> None:
"""Reset the execution namespace."""
self.namespace.clear()
logger.info("Reset execution namespace")
def get_variables(self) -> Dict[str, Any]:
"""Get current namespace variables."""
return self.namespace.copy()
def set_variable(self, name: str, value: Any) -> None:
"""Set a variable in the namespace."""
self.namespace[name] = value
def __repr__(self) -> str:
return f"CodeExecutor(vars={len(self.namespace)}, safe_mode={self.safe_mode})"
__all__ = ["ExecutionResult", "CodeExecutor"]
|