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