Coverage for agentos/tools/code_agent.py: 43%

1""" 

2CodeAgent 工具 — Agent直接写代码执行，不输出JSON。 

3基因来源: Smolagents 

4核心洞察: 代码的表达力远超JSON（循环/条件/异常/变量作用域）。 

5""" 

6 

7from __future__ import annotations 

8 

9import subprocess 

10import tempfile 

11import os 

12 

13from agentos.tools.base import BaseTool, PermissionLevel, ToolResult 

14 

15 

16class CodeAgentTool(BaseTool): 

17 """代码执行工具 — Agent不输出JSON，直接写Python代码。""" 

18 

19 name = "execute_code" 

20 description = ( 

21 "执行Python代码并返回结果。支持任意Python标准库。" 

22 "适用场景：数据处理、文件列表、字符串操作、复杂逻辑。" 

23 "返回值包含stdout/stderr/exit_code。" 

24 ) 

25 permission_level = PermissionLevel.SENSITIVE 

26 

27 @property 

28 def parameters(self) -> dict: 

29 return { 

30 "type": "object", 

31 "properties": { 

32 "code": { 

33 "type": "string", 

34 "description": "要执行的Python代码", 

35 }, 

36 }, 

37 "required": ["code"], 

38 } 

39 

40 async def execute(self, arguments: dict, sandbox=None) -> ToolResult: 

41 code = arguments["code"] 

42 

43 # 如果传入了sandbox则在沙箱中执行 

44 if sandbox: 

45 return await sandbox.execute_code(code, "python") 

46 

47 # 默认在当前进程执行（Python环境） 

48 try: 

49 result = subprocess.run( 

50 ["python3", "-c", code], 

51 capture_output=True, 

52 text=True, 

53 timeout=60, 

54 ) 

55 return ToolResult( 

56 call_id="", 

57 output=result.stdout, 

58 error=result.stderr if result.returncode != 0 else None, 

59 exit_code=result.returncode, 

60 ) 

61 except subprocess.TimeoutExpired: 

62 return ToolResult(call_id="", error="Code execution timed out (60s)") 

63 except Exception as e: 

64 return ToolResult(call_id="", error=str(e)) 

65 

66 def is_write_operation(self, arguments: dict) -> bool: 

67 code = arguments.get("code", "") 

68 write_keywords = ("open(", "write(", "mkdir(", "remove(", "shutil.rmtree") 

69 return any(kw in code for kw in write_keywords) 

70 

71 

72class ShellTool(BaseTool): 

73 """Shell命令执行工具。""" 

74 

75 name = "shell" 

76 description = "执行Shell命令并返回结果。用于文件操作、系统查询等。" 

77 permission_level = PermissionLevel.SENSITIVE 

78 

79 @property 

80 def parameters(self) -> dict: 

81 return { 

82 "type": "object", 

83 "properties": { 

84 "command": { 

85 "type": "string", 

86 "description": "要执行的Shell命令", 

87 }, 

88 }, 

89 "required": ["command"], 

90 } 

91 

92 async def execute(self, arguments: dict, sandbox=None) -> ToolResult: 

93 command = arguments["command"] 

94 

95 if sandbox: 

96 return await sandbox.execute_code(command, "shell") 

97 

98 try: 

99 result = subprocess.run( 

100 command, 

101 shell=True, 

102 capture_output=True, 

103 text=True, 

104 timeout=30, 

105 ) 

106 return ToolResult( 

107 call_id="", 

108 output=result.stdout or result.stderr, 

109 error=None if result.returncode == 0 else result.stderr, 

110 exit_code=result.returncode, 

111 ) 

112 except subprocess.TimeoutExpired: 

113 return ToolResult(call_id="", error="Command timed out (30s)") 

114 except Exception as e: 

115 return ToolResult(call_id="", error=str(e)) 

116 

117 def is_write_operation(self, arguments: dict) -> bool: 

118 cmd = arguments.get("command", "") 

119 write_keywords = ("rm ", "rmdir", "mv ", "cp ", "touch ", "mkdir ", ">") 

120 return any(kw in cmd for kw in write_keywords)