Spaces:
Sleeping
Sleeping
File size: 4,528 Bytes
be745f3 | 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 | """
SAFE PYTHON REPL TOOL FOR LANGCHAIN AGENTS
=========================================
This tool allows an AI agent to execute Python code safely.
Why this exists:
LLM agents often need to:
• do math
• analyze data
• run small scripts
But a raw Python REPL is EXTREMELY dangerous.
So we add guardrails:
✔ restricted builtins
✔ import whitelist
✔ timeout protection
✔ output capture
✔ error handling
"""
# =========================
# Imports
# =========================
from langchain_core.tools import tool
import io
import contextlib
import threading
# =========================
# 1️⃣ Timeout Protection
# =========================
# WHY threading.Timer instead of signal.SIGALRM?
# signal.SIGALRM only exists on Unix/Linux — it crashes on Windows.
# threading.Timer is cross-platform and works identically.
# It fires a callback after N seconds on a background thread,
# which sets a flag that the main thread checks after exec().
class TimeoutException(Exception):
pass
# =========================
# 2️⃣ Restricted Builtins (Sandbox)
# =========================
# WHY:
# Remove dangerous Python functions such as:
# open(), exec(), eval(), __import__(), etc.
SAFE_BUILTINS = {
"print": print,
"len": len,
"range": range,
"sum": sum,
"min": min,
"max": max,
"abs": abs,
"round": round,
"sorted": sorted,
}
# =========================
# 3️⃣ Allowed Libraries
# =========================
# WHY:
# Agent should only use safe scientific libraries.
# Blocks OS/system access.
ALLOWED_IMPORTS = {
"math",
"statistics",
"random",
"numpy",
"pandas",
}
# =========================
# 4️⃣ Import Validator
# =========================
def validate_imports(code: str):
"""
Block dangerous imports like os, sys, subprocess.
"""
for line in code.split("\n"):
line = line.strip()
if line.startswith("import") or line.startswith("from"):
lib = line.split()[1].split(".")[0]
if lib not in ALLOWED_IMPORTS:
raise ValueError(f"Import '{lib}' is not allowed.")
# =========================
# 5️⃣ Safe Execution Engine
# =========================
def execute_python(code: str) -> str:
"""
Executes Python code safely and captures output.
"""
# Validate imports before execution
validate_imports(code)
# Capture print() output
output_buffer = io.StringIO()
# WHY a list instead of a plain bool?
# exec() runs in the same thread, so we need a mutable container
# that the timer callback can write to and the main thread can read.
# A plain bool variable would be a new local binding — not shared.
timed_out = [False]
def _trigger_timeout():
timed_out[0] = True
# WHY threading.Timer?
# It fires _trigger_timeout() after 3 seconds on a background thread.
# Cross-platform — works on Windows, Linux, and Mac.
timer = threading.Timer(3.0, _trigger_timeout)
try:
timer.start()
# Execute code inside restricted environment
with contextlib.redirect_stdout(output_buffer):
exec(code, {"__builtins__": SAFE_BUILTINS}, {})
# Check if timer fired during execution
if timed_out[0]:
return "Execution timed out (3s limit)."
output = output_buffer.getvalue()
return output if output else "Code executed successfully."
except Exception as e:
return f"Execution error: {str(e)}"
finally:
# WHY finally?
# Always cancel the timer — even if exec() raises an exception.
# Without this, the timer thread would keep running in the background.
timer.cancel()
# =========================
# 6️⃣ LangChain Tool Wrapper
# =========================
@tool
def python_repl(code: str) -> str:
"""
Safely execute short Python code.
Use this tool for:
• math calculations
• numpy / pandas data analysis
• quick scripts
Restrictions:
• No file/system/network access
• Only safe libraries allowed
• Execution limited to 3 seconds
"""
# Guardrail: prevent huge code execution
if len(code) > 1000:
return "Code too long. Please shorten."
return execute_python(code)
# =========================
# Example CLI Test
# =========================
if __name__ == "__main__":
print(python_repl.run("print(2+2)"))
print(python_repl.run("import numpy as np\nprint(np.mean([1,2,3]))")) |