ssh connection retry
frdel committed
Dec 3, 2024 at 14:56 UTC
0ebc142124fa3dcb370d95fd2a84bdba8f3145e8
1 file changed
+22
-9
python/tools/code_execution_tool.py
+22
-9
@@ -75,7 +75,7 @@ class CodeExecution(Tool):
75
await self.agent.hist_add_tool_result(self.name, response.message)
76
77
async def prepare_state(self, reset=False):
78
- self.state = self.agent.get_data("cot_state")
78
+ self.state = self.agent.get_data("_cot_state")
79
if not self.state or reset:
80
81
# initialize docker container if execution in docker is configured
@@ -106,7 +106,7 @@ class CodeExecution(Tool):
106
107
self.state = State(shell=shell, docker=docker)
108
await shell.connect()
109
- self.agent.set_data("cot_state", self.state)
109
+ self.agent.set_data("_cot_state", self.state)
110
111
async def execute_python_code(self, code: str, reset: bool = False):
112
escaped_code = shlex.quote(code)
@@ -124,15 +124,28 @@ class CodeExecution(Tool):
124
async def terminal_session(self, command: str, reset: bool = False):
125
126
await self.agent.handle_intervention() # wait for intervention and handle it, if paused
127
- if reset:
128
- await self.reset_terminal()
127
+ # try again on lost connection
128
+ for i in range(2):
129
+ try:
130
+
131
+ if reset:
132
+ await self.reset_terminal()
133
130
- self.state.shell.send_command(command)
134
+ self.state.shell.send_command(command)
135
132
- PrintStyle(background_color="white", font_color="#1B4F72", bold=True).print(
133
- f"{self.agent.agent_name} code execution output"
134
- )
135
- return await self.get_terminal_output()
136
+ PrintStyle(background_color="white", font_color="#1B4F72", bold=True).print(
137
+ f"{self.agent.agent_name} code execution output"
138
+ )
139
+ return await self.get_terminal_output()
140
+
141
+ except Exception as e:
142
+ if i==1:
143
+ # try again on lost connection
144
+ PrintStyle.error(str(e))
145
+ await self.prepare_state(reset=True)
146
+ continue
147
+ else:
148
+ raise e
149
150
async def get_terminal_output(
151
self,