1"""

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

3 基因来源: Smolagents

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

5"""

6

7 from __future__ import annotations

8

9 import subprocess

10 import tempfile

11 import os

12

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

14

15

16 class 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

72 class 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)