Coverage for agentos/core/code_agent.py: 26%

1""" 

2AgentOS v1.1.9 — CodeAgent 模式。 

3 

4基因来源: Smolagents CodeAgent (HuggingFace) 

5 

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

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

8 

9特性: 

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

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

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

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

14""" 

15 

16from __future__ import annotations 

17 

18import ast 

19import asyncio 

20import inspect 

21import sys 

22import traceback 

23from collections.abc import Callable 

24from dataclasses import dataclass, field 

25from typing import Any 

26 

27from agentos.models.router import ModelRouter 

28 

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

30 

31DEFAULT_ALLOWED_MODULES = frozenset( 

32 { 

33 "math", 

34 "json", 

35 "re", 

36 "datetime", 

37 "collections", 

38 "itertools", 

39 "functools", 

40 "typing", 

41 "dataclasses", 

42 "decimal", 

43 "fractions", 

44 "statistics", 

45 "random", 

46 "string", 

47 "textwrap", 

48 "unicodedata", 

49 "hashlib", 

50 "base64", 

51 "binascii", 

52 "uuid", 

53 "copy", 

54 "pprint", 

55 "enum", 

56 "pathlib", 

57 "logging", 

58 "warnings", 

59 "csv", 

60 "html", 

61 "urllib.parse", 

62 "xml.etree.ElementTree", 

63 "operator", 

64 "heapq", 

65 "bisect", 

66 "array", 

67 "struct", 

68 "io", 

69 "os.path", 

70 } 

71) 

72 

73FORBIDDEN_CALLS = frozenset( 

74 { 

75 "exec", 

76 "eval", 

77 "compile", 

78 "open", 

79 "__import__", 

80 "getattr", 

81 "setattr", 

82 "delattr", 

83 "hasattr", 

84 "globals", 

85 "locals", 

86 "vars", 

87 "breakpoint", 

88 "input", 

89 "os", 

90 "subprocess", 

91 "shutil", 

92 "sys", 

93 "ctypes", 

94 "socket", 

95 "pickle", 

96 "marshal", 

97 "multiprocessing", 

98 "threading", 

99 "signal", 

100 } 

101) 

102 

103MAX_OUTPUT_LENGTH = 10000 

104 

105 

106# ── 数据结构 ─────────────────────────────────── 

107 

108 

109@dataclass 

110class CodeStep: 

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

112 

113 step: int 

114 code: str 

115 result: Any = None 

116 stdout: str = "" 

117 error: str | None = None 

118 duration_ms: float = 0.0 

119 

120 

121@dataclass 

122class CodeResult: 

123 """CodeAgent 执行结果。""" 

124 

125 success: bool 

126 final_answer: Any = None 

127 steps: list[CodeStep] = field(default_factory=list) 

128 total_duration_ms: float = 0.0 

129 error: str | None = None 

130 

131 

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

133 

134 

135class CodeGuard(ast.NodeVisitor): 

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

137 

138 def __init__(self, allowed_modules: frozenset): 

139 self.allowed_modules = allowed_modules 

140 self.violations: list[str] = [] 

141 

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

143 for alias in node.names: 

144 if alias.name not in self.allowed_modules: 

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

146 

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

148 module = node.module or "" 

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

150 if base not in self.allowed_modules: 

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

152 

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

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

155 if node.func.id in FORBIDDEN_CALLS: 

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

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

158 parts = [] 

159 curr = node.func 

160 while isinstance(curr, ast.Attribute): 

161 parts.append(curr.attr) 

162 curr = curr.value 

163 if isinstance(curr, ast.Name): 

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

165 for forbidden in FORBIDDEN_CALLS: 

166 if full.startswith(forbidden): 

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

168 break 

169 self.generic_visit(node) 

170 

171 

172def scan_code(code: str, allowed_modules: frozenset) -> list[str]: 

173 try: 

174 tree = ast.parse(code) 

175 except SyntaxError: 

176 return [] 

177 guard = CodeGuard(allowed_modules) 

178 guard.visit(tree) 

179 return guard.violations 

180 

181 

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

183 

184 

185def safe_exec( 

186 code: str, 

187 tools: dict[str, Callable], 

188 state: dict[str, Any], 

189 timeout: float, 

190) -> tuple[Any, str, str | None]: 

191 from io import StringIO 

192 

193 stdout_capture = StringIO() 

194 old_stdout = sys.stdout 

195 sys.stdout = stdout_capture 

196 result = None 

197 error = None 

198 

199 try: 

200 exec_globals = {"__builtins__": __builtins__} 

201 exec_globals.update(tools) 

202 exec_globals.update( 

203 { 

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

205 "__result__": None, 

206 "state": state, 

207 } 

208 ) 

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

210 exec(compiled, exec_globals) 

211 result = exec_globals.get("__result__") 

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

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

214 continue 

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

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

217 except Exception as e: 

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

219 finally: 

220 sys.stdout = old_stdout 

221 

222 stdout = stdout_capture.getvalue() 

223 if len(stdout) > MAX_OUTPUT_LENGTH: 

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

225 return result, stdout, error 

226 

227 

228# ── 代码生成 Prompt ───────────────────────────── 

229 

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

231 

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

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

234 

235Available tools (callable as functions): 

236{tools_description} 

237 

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

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

240 

241Example: 

242```python 

243# Use tools 

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

245# Compute 

246result = len(data) 

247# Return 

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

249``` 

250 

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

252 

253 

254# ── CodeAgent ─────────────────────────────────── 

255 

256 

257class CodeAgent: 

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

259 

260 def __init__( 

261 self, 

262 tools: list[Callable] | None = None, 

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

264 max_steps: int = 10, 

265 timeout_per_step: float = 30.0, 

266 allowed_modules: frozenset = DEFAULT_ALLOWED_MODULES, 

267 ): 

268 self.model = model 

269 self.max_steps = max_steps 

270 self.timeout_per_step = timeout_per_step 

271 self.allowed_modules = allowed_modules 

272 self._tools: dict[str, Callable] = {} 

273 if tools: 

274 for tool in tools: 

275 self._tools[tool.__name__] = tool 

276 

277 @property 

278 def tools(self) -> dict[str, Callable]: 

279 return self._tools 

280 

281 def _tools_description(self) -> str: 

282 lines = [] 

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

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

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

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

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

288 

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

290 if state is None: 

291 state = {"_vars": {}} 

292 tools_desc = self._tools_description() 

293 steps: list[CodeStep] = [] 

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

295 

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

297 if step_num == 1: 

298 user_prompt = task 

299 else: 

300 last = steps[-1] 

301 if last.error: 

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

303 else: 

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

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

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

307 user_prompt = ( 

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

309 ) 

310 

311 router = ModelRouter() 

312 try: 

313 response = await router.chat( 

314 model=self.model, 

315 messages=[ 

316 { 

317 "role": "system", 

318 "content": CODE_AGENT_SYSTEM_PROMPT.format( 

319 tools_description=tools_desc 

320 ), 

321 }, 

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

323 ], 

324 temperature=0.0, 

325 max_tokens=2048, 

326 ) 

327 except Exception as e: 

328 return CodeResult( 

329 success=False, 

330 steps=steps, 

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

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

333 ) 

334 

335 code = self._extract_code(response.content) 

336 if not code: 

337 if steps: 

338 return CodeResult( 

339 success=True, 

340 final_answer=steps[-1].result, 

341 steps=steps, 

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

343 ) 

344 continue 

345 

346 violations = scan_code(code, self.allowed_modules) 

347 if violations: 

348 steps.append( 

349 CodeStep( 

350 step=step_num, 

351 code=code, 

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

353 ) 

354 ) 

355 continue 

356 

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

358 try: 

359 loop = asyncio.get_event_loop() 

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

361 loop.run_in_executor( 

362 None, safe_exec, code, self._tools, state, self.timeout_per_step 

363 ), 

364 timeout=self.timeout_per_step + 5, 

365 ) 

366 except TimeoutError: 

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

368 

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

370 cs = CodeStep( 

371 step=step_num, 

372 code=code, 

373 result=result, 

374 stdout=stdout, 

375 error=error, 

376 duration_ms=step_duration, 

377 ) 

378 steps.append(cs) 

379 

380 if error: 

381 continue 

382 

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

384 return CodeResult( 

385 success=True, 

386 final_answer=result, 

387 steps=steps, 

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

389 ) 

390 

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

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

393 return CodeResult( 

394 success=True, 

395 final_answer=result, 

396 steps=steps, 

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

398 ) 

399 

400 return CodeResult( 

401 success=False, 

402 final_answer=steps[-1].result if steps else None, 

403 steps=steps, 

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

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

406 ) 

407 

408 @staticmethod 

409 def _extract_code(content: str) -> str | None: 

410 if "```python" in content: 

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

412 if len(parts) > 1: 

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

414 if "```" in content: 

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

416 if len(parts) > 1: 

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

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

419 return content.strip() 

420 return None