Spaces:
Running
Running
File size: 1,843 Bytes
2ebf3cc | 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 | import asyncio
import sys
import io
import os
import json
import traceback
from schemas.agent import ToolOutput
MAX_EXECUTION_TIME = 15
MAX_OUTPUT_SIZE = 10000
async def execute_python(code: str) -> ToolOutput:
if not code:
return ToolOutput(tool_name="python", output="", error="No code provided", latency_ms=0)
local_vars = {}
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()
old_stdout = sys.stdout
old_stderr = sys.stderr
try:
sys.stdout = stdout_capture
sys.stderr = stderr_capture
compiled = compile(code.strip(), "<agent_exec>", "exec", flags=0)
loop = asyncio.get_running_loop()
def run_code():
try:
exec(compiled, {"__builtins__": __builtins__, "os": os, "json": json}, local_vars)
except Exception:
traceback.print_exc()
try:
await asyncio.wait_for(loop.run_in_executor(None, run_code), timeout=MAX_EXECUTION_TIME)
except asyncio.TimeoutError:
sys.stdout = old_stdout
sys.stderr = old_stderr
return ToolOutput(tool_name="python", output="", error=f"Execution timed out ({MAX_EXECUTION_TIME}s)", latency_ms=0)
output = stdout_capture.getvalue()[:MAX_OUTPUT_SIZE]
error = stderr_capture.getvalue()[:MAX_OUTPUT_SIZE]
if error and not output:
output = error
return ToolOutput(tool_name="python", output=output or "(no output)", latency_ms=0)
except SyntaxError as e:
return ToolOutput(tool_name="python", output="", error=f"SyntaxError: {e}", latency_ms=0)
except Exception as e:
return ToolOutput(tool_name="python", output="", error=f"{type(e).__name__}: {e}", latency_ms=0)
finally:
sys.stdout = old_stdout
sys.stderr = old_stderr
|