1"""

2 AgentOS v1.1.9 — CodeAgent 模式。

3

4 基因来源: Smolagents CodeAgent (HuggingFace)

5

6 CodeAgent 允许 Agent 通过生成和执行 Python 代码来完成子任务，

7 而非仅调用预定义工具。代码可以调用已注册的 tools + 安全内置函数。

8

9 特性:

10- 多步执行：生成代码 → 执行 → 观察结果 → 继续

11- 安全沙箱：白名单模块、禁止危险操作、超时控制

12- Tools 集成：代码中直接调用 `tool_name(args)`

13- 内存持久：跨步骤的变量和结果通过 locals 传递

14"""

15

16 from __future__ import annotations

17

18 import ast

19 import asyncio

20 import inspect

21 import sys

22 import traceback

23 from dataclasses import dataclass, field

24 from typing import Any, Callable, Dict, List, Optional, Tuple

25

26 from agentos.models.router import ModelRouter

27

28

29# ── 安全常量 ───────────────────────────────────

30

31 DEFAULT_ALLOWED_MODULES = frozenset({

32 "math", "json", "re", "datetime", "collections",

33 "itertools", "functools", "typing", "dataclasses",

34 "decimal", "fractions", "statistics", "random",

35 "string", "textwrap", "unicodedata", "hashlib",

36 "base64", "binascii", "uuid", "copy", "pprint",

37 "enum", "pathlib", "logging", "warnings",

38 "csv", "html", "urllib.parse", "xml.etree.ElementTree",

39 "operator", "heapq", "bisect", "array",

40 "struct", "io", "os.path",

41})

42

43 FORBIDDEN_CALLS = frozenset({

44 "exec", "eval", "compile", "open", "__import__",

45 "getattr", "setattr", "delattr", "hasattr",

46 "globals", "locals", "vars",

47 "breakpoint", "input",

48 "os", "subprocess", "shutil", "sys",

49 "ctypes", "socket", "pickle", "marshal",

50 "multiprocessing", "threading", "signal",

51})

52

53 MAX_OUTPUT_LENGTH = 10000

54

55

56# ── 数据结构 ───────────────────────────────────

57

58@dataclass

59 class CodeStep:

60 """CodeAgent 单步执行记录。"""

61

62 step: int

63 code: str

64 result: Any = None

65 stdout: str = ""

66 error: Optional[str] = None

67 duration_ms: float = 0.0

68

69

70@dataclass

71 class CodeResult:

72 """CodeAgent 执行结果。"""

73

74 success: bool

75 final_answer: Any = None

76 steps: List[CodeStep] = field(default_factory=list)

77 total_duration_ms: float = 0.0

78 error: Optional[str] = None

79

80

81# ── 代码安全检查器 ─────────────────────────────

82

83 class CodeGuard(ast.NodeVisitor):

84 """Python 代码 AST 安全扫描器，拦截危险操作。"""

85

86 def __init__(self, allowed_modules: frozenset):

87 self.allowed_modules = allowed_modules

88 self.violations: List[str] = []

89

90 def visit_Import(self, node: ast.Import) -> None:

91 for alias in node.names:

92 if alias.name not in self.allowed_modules:

93 self.violations.append(f"import '{alias.name}' not allowed")

94

95 def visit_ImportFrom(self, node: ast.ImportFrom) -> None:

96 module = node.module or ""

97 base = module.split(".")[0]

98 if base not in self.allowed_modules:

99 self.violations.append(f"import from '{module}' not allowed")

100

101 def visit_Call(self, node: ast.Call) -> None:

102 if isinstance(node.func, ast.Name):

103 if node.func.id in FORBIDDEN_CALLS:

104 self.violations.append(f"call to '{node.func.id}()' is forbidden")

105 elif isinstance(node.func, ast.Attribute):

106 parts = []

107 curr = node.func

108 while isinstance(curr, ast.Attribute):

109 parts.append(curr.attr)

110 curr = curr.value

111 if isinstance(curr, ast.Name):

112 full = f"{curr.id}.{'.'.join(reversed(parts))}"

113 for forbidden in FORBIDDEN_CALLS:

114 if full.startswith(forbidden):

115 self.violations.append(f"call to '{full}()' is forbidden")

116 break

117 self.generic_visit(node)

118

119

120 def scan_code(code: str, allowed_modules: frozenset) -> List[str]:

121 try:

122 tree = ast.parse(code)

123 except SyntaxError:

124 return []

125 guard = CodeGuard(allowed_modules)

126 guard.visit(tree)

127 return guard.violations

128

129

130# ── 受控执行环境 ───────────────────────────────

131

132 def safe_exec(

133 code: str,

134 tools: Dict[str, Callable],

135 state: Dict[str, Any],

136 timeout: float,

137) -> Tuple[Any, str, Optional[str]]:

138 from io import StringIO

139

140 stdout_capture = StringIO()

141 old_stdout = sys.stdout

142 sys.stdout = stdout_capture

143 result = None

144 error = None

145

146 try:

147 exec_globals = {"__builtins__": __builtins__}

148 exec_globals.update(tools)

149 exec_globals.update({

150 "print": lambda *a, **kw: print(*a, **kw),

151 "__result__": None,

152 "state": state,

153 })

154 compiled = compile(code, "<code_agent>", "exec")

155 exec(compiled, exec_globals)

156 result = exec_globals.get("__result__")

157 for key in list(exec_globals.keys()):

158 if key.startswith("_") or key in tools or key in ("state", "print"):

159 continue

160 if key not in ("__builtins__",):

161 state.setdefault("_vars", {})[key] = exec_globals[key]

162 except Exception as e:

163 error = f"{type(e).__name__}: {e}\n{traceback.format_exc(limit=3)}"

164 finally:

165 sys.stdout = old_stdout

166

167 stdout = stdout_capture.getvalue()

168 if len(stdout) > MAX_OUTPUT_LENGTH:

169 stdout = stdout[:MAX_OUTPUT_LENGTH] + "\n... [truncated]"

170 return result, stdout, error

171

172

173# ── 代码生成 Prompt ─────────────────────────────

174

175 CODE_AGENT_SYSTEM_PROMPT = """You are a CodeAgent that solves tasks by writing and executing Python code.

176

177 YOU MUST respond ONLY with Python code inside ```python ... ``` blocks.

178 NO explanations, NO markdown outside the code block. Just the code.

179

180 Available tools (callable as functions):

181{tools_description}

182

183 To output the final answer, assign it to the variable `__result__`.

184 You can store persistent data in the `state` dict.

185

186 Example:

187```python

188# Use tools

189 data = web_search("Python 3.12 release date")

190# Compute

191 result = len(data)

192# Return

193__result__ = f"Found {{result}} results"

194```

195

196 Now solve the following task. ONLY output the code block."""

197

198

199# ── CodeAgent ───────────────────────────────────

200

201 class CodeAgent:

202 """代码执行型 Agent。"""

203

204 def __init__(

205 self,

206 tools: List[Callable] | None = None,

207 model: str = "gpt-4o",

208 max_steps: int = 10,

209 timeout_per_step: float = 30.0,

210 allowed_modules: frozenset = DEFAULT_ALLOWED_MODULES,

211 ):

212 self.model = model

213 self.max_steps = max_steps

214 self.timeout_per_step = timeout_per_step

215 self.allowed_modules = allowed_modules

216 self._tools: Dict[str, Callable] = {}

217 if tools:

218 for tool in tools:

219 self._tools[tool.__name__] = tool

220

221 @property

222 def tools(self) -> Dict[str, Callable]:

223 return self._tools

224

225 def _tools_description(self) -> str:

226 lines = []

227 for name, fn in self._tools.items():

228 sig = str(inspect.signature(fn))

229 doc = (inspect.getdoc(fn) or "No description").split("\n")[0]

230 lines.append(f" {name}{sig}: {doc}")

231 return "\n".join(lines) if lines else " (no tools available)"

232

233 async def run(self, task: str, state: Dict[str, Any] | None = None) -> CodeResult:

234 if state is None:

235 state = {"_vars": {}}

236 tools_desc = self._tools_description()

237 steps: List[CodeStep] = []

238 total_start = asyncio.get_event_loop().time()

239

240 for step_num in range(1, self.max_steps + 1):

241 if step_num == 1:

242 user_prompt = task

243 else:

244 last = steps[-1]

245 if last.error:

246 feedback = f"Error: {last.error}"

247 else:

248 rp = str(last.result)[:500] if last.result is not None else "None"

249 op = last.stdout[:500] if last.stdout else ""

250 feedback = f"Output: {op}\nResult: {rp}"

251 user_prompt = f"Step {step_num}: Continue.\nPrevious result:\n{feedback}\n\nTask: {task}"

252

253 router = ModelRouter()

254 try:

255 response = await router.chat(

256 model=self.model,

257 messages=[

258 {"role": "system", "content": CODE_AGENT_SYSTEM_PROMPT.format(tools_description=tools_desc)},

259 {"role": "user", "content": user_prompt},

260 ],

261 temperature=0.0,

262 max_tokens=2048,

263 )

264 except Exception as e:

265 return CodeResult(

266 success=False, steps=steps,

267 total_duration_ms=(asyncio.get_event_loop().time() - total_start) * 1000,

268 error=f"LLM error: {e}",

269 )

270

271 code = self._extract_code(response.content)

272 if not code:

273 if steps:

274 return CodeResult(

275 success=True, final_answer=steps[-1].result, steps=steps,

276 total_duration_ms=(asyncio.get_event_loop().time() - total_start) * 1000,

277 )

278 continue

279

280 violations = scan_code(code, self.allowed_modules)

281 if violations:

282 steps.append(CodeStep(step=step_num, code=code,

283 error=f"Security violation: {'; '.join(violations)}"))

284 continue

285

286 step_start = asyncio.get_event_loop().time()

287 try:

288 loop = asyncio.get_event_loop()

289 result, stdout, error = await asyncio.wait_for(

290 loop.run_in_executor(None, safe_exec, code, self._tools, state, self.timeout_per_step),

291 timeout=self.timeout_per_step + 5,

292 )

293 except asyncio.TimeoutError:

294 result, stdout, error = None, "", "TimeoutError: exceeded limit"

295

296 step_duration = (asyncio.get_event_loop().time() - step_start) * 1000

297 cs = CodeStep(step=step_num, code=code, result=result, stdout=stdout, error=error, duration_ms=step_duration)

298 steps.append(cs)

299

300 if error:

301 continue

302

303 if "__result__" in code or (result is not None and "__result__" in code):

304 return CodeResult(

305 success=True, final_answer=result, steps=steps,

306 total_duration_ms=(asyncio.get_event_loop().time() - total_start) * 1000,

307 )

308

309 # heuristic: non-trivial result without error = likely done

310 if result is not None and step_num >= 1:

311 return CodeResult(

312 success=True, final_answer=result, steps=steps,

313 total_duration_ms=(asyncio.get_event_loop().time() - total_start) * 1000,

314 )

315

316 return CodeResult(

317 success=False, final_answer=steps[-1].result if steps else None, steps=steps,

318 total_duration_ms=(asyncio.get_event_loop().time() - total_start) * 1000,

319 error=f"Max steps ({self.max_steps}) reached",

320 )

321

322 @staticmethod

323 def _extract_code(content: str) -> Optional[str]:

324 if "```python" in content:

325 parts = content.split("```python", 1)

326 if len(parts) > 1:

327 return parts[1].split("```", 1)[0].strip()

328 if "```" in content:

329 parts = content.split("```", 1)

330 if len(parts) > 1:

331 return parts[1].split("```", 1)[0].strip()

332 if "print(" in content or "def " in content or "result" in content:

333 return content.strip()

334 return None