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

1""" 

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

3基因来源: Smolagents 

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

5""" 

6 

7from __future__ import annotations 

8 

9import subprocess 

10 

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

12 

13 

14class CodeAgentTool(BaseTool): 

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

16 

17 name = "execute_code" 

18 description = ( 

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

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

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

22 ) 

23 permission_level = PermissionLevel.SENSITIVE 

24 

25 @property 

26 def parameters(self) -> dict: 

27 return { 

28 "type": "object", 

29 "properties": { 

30 "code": { 

31 "type": "string", 

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

33 }, 

34 }, 

35 "required": ["code"], 

36 } 

37 

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

39 code = arguments["code"] 

40 

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

42 if sandbox: 

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

44 

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

46 try: 

47 result = subprocess.run( 

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

49 capture_output=True, 

50 text=True, 

51 timeout=60, 

52 ) 

53 return ToolResult( 

54 call_id="", 

55 output=result.stdout, 

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

57 exit_code=result.returncode, 

58 ) 

59 except subprocess.TimeoutExpired: 

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

61 except Exception as e: 

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

63 

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

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

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

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

68 

69 

70class ShellTool(BaseTool): 

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

72 

73 name = "shell" 

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

75 permission_level = PermissionLevel.SENSITIVE 

76 

77 @property 

78 def parameters(self) -> dict: 

79 return { 

80 "type": "object", 

81 "properties": { 

82 "command": { 

83 "type": "string", 

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

85 }, 

86 }, 

87 "required": ["command"], 

88 } 

89 

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

91 command = arguments["command"] 

92 

93 if sandbox: 

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

95 

96 try: 

97 result = subprocess.run( 

98 command, 

99 shell=True, 

100 capture_output=True, 

101 text=True, 

102 timeout=30, 

103 ) 

104 return ToolResult( 

105 call_id="", 

106 output=result.stdout or result.stderr, 

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

108 exit_code=result.returncode, 

109 ) 

110 except subprocess.TimeoutExpired: 

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

112 except Exception as e: 

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

114 

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

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

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

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