Skip to content

Commit 9e09299

Browse files
FronutSoulter
andauthored
feat: add python tool timeout param (#7953)
* feat: add python tool timeout param * Update python.py --------- Co-authored-by: Weilong Liao <37870767+Soulter@users.noreply.github.com>
1 parent 77fe2de commit 9e09299

1 file changed

Lines changed: 35 additions & 4 deletions

File tree

astrbot/core/tools/computer_tools/python.py

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,11 @@
3333
"description": "Whether to suppress the output of the code execution.",
3434
"default": False,
3535
},
36+
"timeout": {
37+
"type": "integer",
38+
"description": "Optional timeout in seconds for code execution.",
39+
"default": 30,
40+
},
3641
},
3742
"required": ["code"],
3843
}
@@ -77,16 +82,29 @@ class PythonTool(FunctionTool):
7782
parameters: dict = field(default_factory=lambda: param_schema)
7883

7984
async def call(
80-
self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False
85+
self,
86+
context: ContextWrapper[AstrAgentContext],
87+
code: str,
88+
silent: bool = False,
89+
timeout: int = 30,
8190
) -> ToolExecResult:
8291
if permission_error := check_admin_permission(context, "Python execution"):
8392
return permission_error
8493
sb = await get_booter(
8594
context.context.context,
8695
context.context.event.unified_msg_origin,
8796
)
97+
effective_timeout = (
98+
min(timeout, context.tool_call_timeout)
99+
if timeout > 0
100+
else context.tool_call_timeout
101+
)
88102
try:
89-
result = await sb.python.exec(code, silent=silent)
103+
result = await sb.python.exec(
104+
code,
105+
timeout=effective_timeout,
106+
silent=silent,
107+
)
90108
return await handle_result(result, context.context.event)
91109
except Exception as e:
92110
return f"Error executing code: {str(e)}"
@@ -104,13 +122,26 @@ class LocalPythonTool(FunctionTool):
104122
parameters: dict = field(default_factory=lambda: param_schema)
105123

106124
async def call(
107-
self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False
125+
self,
126+
context: ContextWrapper[AstrAgentContext],
127+
code: str,
128+
silent: bool = False,
129+
timeout: int = 30,
108130
) -> ToolExecResult:
109131
if permission_error := check_admin_permission(context, "Python execution"):
110132
return permission_error
111133
sb = get_local_booter()
134+
effective_timeout = (
135+
min(timeout, context.tool_call_timeout)
136+
if timeout > 0
137+
else context.tool_call_timeout
138+
)
112139
try:
113-
result = await sb.python.exec(code, silent=silent)
140+
result = await sb.python.exec(
141+
code,
142+
timeout=effective_timeout,
143+
silent=silent,
144+
)
114145
return await handle_result(result, context.context.event)
115146
except Exception as e:
116147
return f"Error executing code: {str(e)}"

0 commit comments

Comments
 (0)