Skip to content

Commit d910487

Browse files
committed
feat: add python tool timeout param
1 parent dee4f14 commit d910487

1 file changed

Lines changed: 41 additions & 4 deletions

File tree

astrbot/core/tools/computer_tools/python.py

Lines changed: 41 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. Omit or set to 0 to use tool_call_timeout.",
39+
"default": 0,
40+
},
3641
},
3742
"required": ["code"],
3843
}
@@ -77,16 +82,32 @@ 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 | None = None,
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 = context.tool_call_timeout
98+
if timeout is not None:
99+
try:
100+
timeout_value = int(timeout)
101+
except (TypeError, ValueError):
102+
timeout_value = None
103+
if timeout_value and timeout_value > 0:
104+
effective_timeout = min(timeout_value, context.tool_call_timeout)
88105
try:
89-
result = await sb.python.exec(code, silent=silent)
106+
result = await sb.python.exec(
107+
code,
108+
timeout=effective_timeout,
109+
silent=silent,
110+
)
90111
return await handle_result(result, context.context.event)
91112
except Exception as e:
92113
return f"Error executing code: {str(e)}"
@@ -104,13 +125,29 @@ class LocalPythonTool(FunctionTool):
104125
parameters: dict = field(default_factory=lambda: param_schema)
105126

106127
async def call(
107-
self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False
128+
self,
129+
context: ContextWrapper[AstrAgentContext],
130+
code: str,
131+
silent: bool = False,
132+
timeout: int | None = None,
108133
) -> ToolExecResult:
109134
if permission_error := check_admin_permission(context, "Python execution"):
110135
return permission_error
111136
sb = get_local_booter()
137+
effective_timeout = context.tool_call_timeout
138+
if timeout is not None:
139+
try:
140+
timeout_value = int(timeout)
141+
except (TypeError, ValueError):
142+
timeout_value = None
143+
if timeout_value and timeout_value > 0:
144+
effective_timeout = min(timeout_value, context.tool_call_timeout)
112145
try:
113-
result = await sb.python.exec(code, silent=silent)
146+
result = await sb.python.exec(
147+
code,
148+
timeout=effective_timeout,
149+
silent=silent,
150+
)
114151
return await handle_result(result, context.context.event)
115152
except Exception as e:
116153
return f"Error executing code: {str(e)}"

0 commit comments

Comments
 (0)