MCPcopy Create free account
hub / github.com/THUDM/AgentTuning / Container

Class Container

AgentBench.old/src/tasks/os_interaction/task.py:23–103  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

21
22
23class Container:
24 def __init__(self, image):
25 self.image = image
26 self.client = docker.from_env()
27 self.container: docker.models.containers.Container = self.client.containers.run(
28 image, detach=True, tty=True, stdin_open=True, remove=True,
29 labels={"created_by": "os-pipeline"}
30 )
31 self.exec_id = self.client.api.exec_create(self.container.id, "bash --login", stdin=True, tty=True)["Id"]
32 self.sock = self.client.api.exec_start(self.exec_id, socket=True)._sock
33 self.sock.settimeout(5)
34 # clear buffer
35 self.sock.recv(1000)
36
37 def __del__(self):
38 try:
39 # print("Releasing", self.container.name, self.container.id)
40 self.container.stop()
41 except:
42 pass
43
44 def execute(self, command: str):
45 class DummyOutput:
46 output: bytes
47 exit_code: int
48
49 def __init__(self, code, o):
50 self.output = o
51 self.exit_code = code
52
53 # print("=== EXECUTING ===\n", command)
54 if not isinstance(command, str):
55 return DummyOutput(-1, b'')
56 self.sock.send(command.encode("utf-8") + b'\n')
57 # ignore input line
58 data = self.sock.recv(8)
59 _, n = struct.unpack('>BxxxL', data)
60 _ = self.sock.recv(n)
61 output = b''
62 while True:
63 try:
64 data = self.sock.recv(8)
65 # print(data)
66 if not data:
67 break
68 _, n = struct.unpack('>BxxxL', data)
69 line = self.sock.recv(n)
70 output += line
71 if re.search(b"\x1b.+@.+[#|$] ", line):
72 break
73 except TimeoutError:
74 break
75 except socket.timeout:
76 break
77 # replace the very end \x1b.+@.+[#|$] into nothing (required the suffix)
78 # output = re.sub(b"\x1b.+@.+[#|$] $", b"", output)
79 return DummyOutput(0, output)
80

Callers 1

predict_singleMethod · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected