code exex updates
output update, termination
frdel committed
Sep 4, 2024 at 23:12 UTC
8c46a0dbe22e034e6a7a35c0c586153a8cfdd7fc
1 file changed
+12
-7
python/tools/code_execution_tool.py
+12
-7
@@ -33,7 +33,9 @@ class CodeExecution(Tool):
33
elif runtime == "terminal":
34
response = await self.execute_terminal_command(self.args["code"])
35
elif runtime == "output":
36
- response = await self.get_terminal_output()
36
+ response = await self.get_terminal_output(wait_with_output=5, wait_without_output=20)
37
+ elif runtime == "reset":
38
+ response = await self.reset_terminal()
39
else:
40
response = self.agent.read_prompt("fw.code_runtime_wrong.md", runtime=runtime)
41
@@ -54,9 +56,9 @@ class CodeExecution(Tool):
56
msg_response = self.agent.read_prompt("fw.tool_response.md", tool_name=self.name, tool_response=response.message)
57
await self.agent.append_message(msg_response, human=True)
58
57
- async def prepare_state(self):
59
+ async def prepare_state(self, reset=False):
60
self.state = self.agent.get_data("cot_state")
59
- if not self.state:
61
+ if not self.state or reset:
62
63
#initialize docker container if execution in docker is configured
64
if self.agent.config.code_exec_docker_enabled:
@@ -95,10 +97,11 @@ class CodeExecution(Tool):
97
PrintStyle(background_color="white",font_color="#1B4F72",bold=True).print(f"{self.agent.agent_name} code execution output:")
98
return await self.get_terminal_output()
99
98
- async def get_terminal_output(self):
100
+ async def get_terminal_output(self, wait_with_output=3, wait_without_output=10):
101
idle=0
102
+ SLEEP_TIME = 0.1
103
while True:
101
- await asyncio.sleep(0.1) # Wait for some output to be generated
104
+ await asyncio.sleep(SLEEP_TIME) # Wait for some output to be generated
105
full_output, partial_output = await self.state.shell.read_output()
106
107
await self.agent.handle_intervention() # wait for intervention and handle it, if paused
@@ -109,5 +112,7 @@ class CodeExecution(Tool):
112
idle=0
113
else:
114
idle+=1
112
- if ( full_output and idle > 30 ) or ( not full_output and idle > 100 ): return full_output
113
-
\ No newline at end of file
115
+ if ( full_output and idle > wait_with_output / SLEEP_TIME ) or ( not full_output and idle > wait_without_output / SLEEP_TIME ): return full_output
116
+
117
+ async def reset_terminal(self):
118
+ await self.prepare_state(reset=True)
\ No newline at end of file