ai-computer-client 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ai_computer/__init__.py +13 -3
- ai_computer/client.py +171 -548
- ai_computer/models.py +45 -0
- ai_computer/submodules/__init__.py +5 -0
- ai_computer/submodules/base.py +81 -0
- ai_computer/submodules/code.py +295 -0
- ai_computer/submodules/filesystem.py +373 -0
- ai_computer/submodules/shell.py +52 -0
- {ai_computer_client-0.3.2.dist-info → ai_computer_client-0.3.3.dist-info}/METADATA +1 -1
- ai_computer_client-0.3.3.dist-info/RECORD +12 -0
- ai_computer_client-0.3.2.dist-info/RECORD +0 -6
- {ai_computer_client-0.3.2.dist-info → ai_computer_client-0.3.3.dist-info}/WHEEL +0 -0
- {ai_computer_client-0.3.2.dist-info → ai_computer_client-0.3.3.dist-info}/licenses/LICENSE +0 -0
ai_computer/client.py
CHANGED
@@ -7,6 +7,9 @@ import os
|
|
7
7
|
import mimetypes
|
8
8
|
from pathlib import Path
|
9
9
|
|
10
|
+
from .models import SandboxResponse, StreamEvent, FileOperationResponse
|
11
|
+
from .submodules import FileSystemModule, ShellModule, CodeModule
|
12
|
+
|
10
13
|
@dataclass
|
11
14
|
class SandboxResponse:
|
12
15
|
"""Response from sandbox operations.
|
@@ -56,6 +59,11 @@ class SandboxClient:
|
|
56
59
|
This client provides methods to execute Python code in an isolated sandbox environment.
|
57
60
|
It handles authentication, sandbox creation/deletion, and code execution.
|
58
61
|
|
62
|
+
The client is organized into submodules for different types of operations:
|
63
|
+
- fs: File system operations (upload, download, read, write)
|
64
|
+
- shell: Shell command execution
|
65
|
+
- code: Python code execution
|
66
|
+
|
59
67
|
Args:
|
60
68
|
base_url: The base URL of the sandbox service
|
61
69
|
token: Optional pre-existing authentication token
|
@@ -70,6 +78,26 @@ class SandboxClient:
|
|
70
78
|
self.token = token
|
71
79
|
self.sandbox_id = None
|
72
80
|
|
81
|
+
# Initialize submodules
|
82
|
+
self._fs = FileSystemModule(self)
|
83
|
+
self._shell = ShellModule(self)
|
84
|
+
self._code = CodeModule(self)
|
85
|
+
|
86
|
+
@property
|
87
|
+
def fs(self) -> FileSystemModule:
|
88
|
+
"""File system operations submodule."""
|
89
|
+
return self._fs
|
90
|
+
|
91
|
+
@property
|
92
|
+
def shell(self) -> ShellModule:
|
93
|
+
"""Shell operations submodule."""
|
94
|
+
return self._shell
|
95
|
+
|
96
|
+
@property
|
97
|
+
def code(self) -> CodeModule:
|
98
|
+
"""Code execution operations submodule."""
|
99
|
+
return self._code
|
100
|
+
|
73
101
|
async def setup(self) -> SandboxResponse:
|
74
102
|
"""Initialize the client and create a sandbox.
|
75
103
|
|
@@ -98,179 +126,124 @@ class SandboxClient:
|
|
98
126
|
if response.status == 200:
|
99
127
|
data = await response.json()
|
100
128
|
self.sandbox_id = data["sandbox_id"]
|
101
|
-
# Wait for sandbox to be ready
|
102
|
-
ready = await self.wait_for_ready()
|
103
|
-
if not ready.success:
|
104
|
-
return ready
|
105
|
-
return SandboxResponse(success=True, data=data)
|
106
129
|
else:
|
107
130
|
text = await response.text()
|
108
131
|
return SandboxResponse(success=False, error=text)
|
132
|
+
|
133
|
+
# Wait for sandbox to be ready
|
134
|
+
return await self.wait_for_ready()
|
109
135
|
|
110
|
-
async def wait_for_ready(self,
|
111
|
-
"""Wait for the sandbox to be
|
136
|
+
async def wait_for_ready(self, max_attempts: int = 10, delay: float = 1.0) -> SandboxResponse:
|
137
|
+
"""Wait for the sandbox to be ready.
|
112
138
|
|
113
139
|
Args:
|
114
|
-
|
115
|
-
delay: Delay between
|
140
|
+
max_attempts: Maximum number of attempts to check if sandbox is ready
|
141
|
+
delay: Delay between attempts in seconds
|
116
142
|
|
117
143
|
Returns:
|
118
|
-
SandboxResponse indicating if sandbox is ready
|
144
|
+
SandboxResponse indicating if the sandbox is ready
|
119
145
|
"""
|
120
146
|
if not self.token or not self.sandbox_id:
|
121
|
-
return SandboxResponse(success=False, error="Client not properly initialized")
|
147
|
+
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
122
148
|
|
123
149
|
headers = {"Authorization": f"Bearer {self.token}"}
|
124
150
|
|
125
|
-
for
|
126
|
-
|
127
|
-
async with
|
128
|
-
|
129
|
-
|
130
|
-
|
131
|
-
|
151
|
+
for attempt in range(max_attempts):
|
152
|
+
try:
|
153
|
+
async with aiohttp.ClientSession() as session:
|
154
|
+
async with session.get(
|
155
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/status",
|
156
|
+
headers=headers
|
157
|
+
) as response:
|
158
|
+
if response.status != 200:
|
159
|
+
# If we get an error, wait and try again
|
160
|
+
await asyncio.sleep(delay)
|
161
|
+
continue
|
162
|
+
|
132
163
|
data = await response.json()
|
133
|
-
|
164
|
+
status = data.get("status")
|
165
|
+
|
166
|
+
if status == "ready":
|
134
167
|
return SandboxResponse(success=True, data=data)
|
135
|
-
|
136
|
-
|
137
|
-
|
168
|
+
elif status == "error":
|
169
|
+
return SandboxResponse(
|
170
|
+
success=False,
|
171
|
+
error=data.get("error", "Unknown error initializing sandbox")
|
172
|
+
)
|
173
|
+
|
174
|
+
# If not ready yet, wait and try again
|
175
|
+
await asyncio.sleep(delay)
|
176
|
+
|
177
|
+
except Exception as e:
|
178
|
+
# If we get an exception, wait and try again
|
179
|
+
await asyncio.sleep(delay)
|
180
|
+
|
181
|
+
return SandboxResponse(success=False, error="Timed out waiting for sandbox to be ready")
|
138
182
|
|
139
|
-
async def
|
140
|
-
|
141
|
-
code: Union[str, bytes],
|
142
|
-
timeout: int = 30
|
143
|
-
) -> SandboxResponse:
|
144
|
-
"""Execute Python code in the sandbox and return the combined output.
|
145
|
-
|
146
|
-
This method collects all output from the streaming response and returns it as a single result.
|
147
|
-
It captures both stdout and stderr, and handles any errors during execution.
|
183
|
+
async def cleanup(self) -> SandboxResponse:
|
184
|
+
"""Delete the sandbox.
|
148
185
|
|
149
|
-
Args:
|
150
|
-
code: Python code to execute
|
151
|
-
timeout: Maximum execution time in seconds
|
152
|
-
|
153
186
|
Returns:
|
154
|
-
SandboxResponse
|
187
|
+
SandboxResponse indicating success/failure
|
155
188
|
"""
|
156
189
|
if not self.token or not self.sandbox_id:
|
157
190
|
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
158
191
|
|
159
|
-
|
160
|
-
ready = await self.wait_for_ready()
|
161
|
-
if not ready.success:
|
162
|
-
return ready
|
163
|
-
|
164
|
-
headers = {
|
165
|
-
"Authorization": f"Bearer {self.token}",
|
166
|
-
"Content-Type": "application/json"
|
167
|
-
}
|
168
|
-
|
169
|
-
data = {
|
170
|
-
"code": code,
|
171
|
-
"timeout": timeout
|
172
|
-
}
|
192
|
+
headers = {"Authorization": f"Bearer {self.token}"}
|
173
193
|
|
174
194
|
try:
|
175
195
|
async with aiohttp.ClientSession() as session:
|
176
|
-
async with session.
|
177
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}
|
178
|
-
headers=headers
|
179
|
-
json=data
|
196
|
+
async with session.delete(
|
197
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}",
|
198
|
+
headers=headers
|
180
199
|
) as response:
|
181
200
|
if response.status != 200:
|
182
|
-
|
183
|
-
return SandboxResponse(success=False, error=
|
201
|
+
text = await response.text()
|
202
|
+
return SandboxResponse(success=False, error=text)
|
184
203
|
|
185
|
-
#
|
186
|
-
|
187
|
-
return SandboxResponse(success=True
|
204
|
+
# Reset sandbox ID
|
205
|
+
self.sandbox_id = None
|
206
|
+
return SandboxResponse(success=True)
|
188
207
|
|
189
208
|
except Exception as e:
|
190
209
|
return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
|
191
210
|
|
192
|
-
|
193
|
-
|
194
|
-
|
195
|
-
|
196
|
-
) -> AsyncGenerator[StreamEvent, None]:
|
197
|
-
"""Execute Python code in the sandbox and stream the output.
|
211
|
+
# Backward compatibility methods
|
212
|
+
|
213
|
+
async def execute_code(self, code: str, timeout: int = 30) -> SandboxResponse:
|
214
|
+
"""Execute Python code in the sandbox.
|
198
215
|
|
199
|
-
This
|
200
|
-
the type of event and the associated data.
|
216
|
+
This is a backward compatibility method that delegates to the code submodule.
|
201
217
|
|
202
218
|
Args:
|
203
|
-
code: Python code to execute
|
219
|
+
code: The Python code to execute
|
204
220
|
timeout: Maximum execution time in seconds
|
205
221
|
|
206
|
-
|
207
|
-
|
222
|
+
Returns:
|
223
|
+
SandboxResponse containing execution results
|
208
224
|
"""
|
209
|
-
|
210
|
-
|
211
|
-
|
212
|
-
|
213
|
-
# Ensure sandbox is ready
|
214
|
-
ready = await self.wait_for_ready()
|
215
|
-
if not ready.success:
|
216
|
-
yield StreamEvent(type="error", data=ready.error or "Sandbox not ready")
|
217
|
-
return
|
218
|
-
|
219
|
-
headers = {
|
220
|
-
"Authorization": f"Bearer {self.token}",
|
221
|
-
"Content-Type": "application/json"
|
222
|
-
}
|
225
|
+
return await self.code.execute(code, timeout)
|
226
|
+
|
227
|
+
async def execute_code_stream(self, code: str, timeout: int = 30) -> AsyncGenerator[StreamEvent, None]:
|
228
|
+
"""Execute Python code in the sandbox with streaming output.
|
223
229
|
|
224
|
-
|
225
|
-
"code": code,
|
226
|
-
"timeout": timeout
|
227
|
-
}
|
230
|
+
This is a backward compatibility method that delegates to the code submodule.
|
228
231
|
|
229
|
-
|
230
|
-
|
231
|
-
|
232
|
-
total=timeout + 30, # Add buffer for connection overhead
|
233
|
-
connect=30,
|
234
|
-
sock_connect=30,
|
235
|
-
sock_read=timeout + 30
|
236
|
-
)
|
232
|
+
Args:
|
233
|
+
code: The Python code to execute
|
234
|
+
timeout: Maximum execution time in seconds
|
237
235
|
|
238
|
-
|
239
|
-
|
240
|
-
|
241
|
-
|
242
|
-
|
243
|
-
) as response:
|
244
|
-
if response.status != 200:
|
245
|
-
error_text = await response.text()
|
246
|
-
yield StreamEvent(type="error", data=error_text)
|
247
|
-
return
|
248
|
-
|
249
|
-
# Process the streaming response
|
250
|
-
async for line in response.content:
|
251
|
-
if line:
|
252
|
-
try:
|
253
|
-
event = json.loads(line.decode())
|
254
|
-
yield StreamEvent(type=event['type'], data=event['data'])
|
255
|
-
|
256
|
-
# Stop if we receive an error or completed event
|
257
|
-
if event['type'] in ['error', 'completed']:
|
258
|
-
break
|
259
|
-
except json.JSONDecodeError as e:
|
260
|
-
yield StreamEvent(type="error", data=f"Failed to parse event: {str(e)}")
|
261
|
-
break
|
262
|
-
|
263
|
-
except Exception as e:
|
264
|
-
yield StreamEvent(type="error", data=f"Connection error: {str(e)}")
|
236
|
+
Yields:
|
237
|
+
StreamEvent objects containing execution output
|
238
|
+
"""
|
239
|
+
async for event in self.code.execute_stream(code, timeout):
|
240
|
+
yield event
|
265
241
|
|
266
|
-
async def execute_shell(
|
267
|
-
self,
|
268
|
-
command: str,
|
269
|
-
args: Optional[List[str]] = None,
|
270
|
-
timeout: int = 30
|
271
|
-
) -> SandboxResponse:
|
242
|
+
async def execute_shell(self, command: str, args: Optional[List[str]] = None, timeout: int = 30) -> SandboxResponse:
|
272
243
|
"""Execute a shell command in the sandbox.
|
273
244
|
|
245
|
+
This is a backward compatibility method that delegates to the shell submodule.
|
246
|
+
|
274
247
|
Args:
|
275
248
|
command: The shell command to execute
|
276
249
|
args: Optional list of arguments for the command
|
@@ -279,75 +252,19 @@ class SandboxClient:
|
|
279
252
|
Returns:
|
280
253
|
SandboxResponse containing execution results
|
281
254
|
"""
|
282
|
-
|
283
|
-
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
284
|
-
|
285
|
-
# Ensure sandbox is ready
|
286
|
-
ready = await self.wait_for_ready()
|
287
|
-
if not ready.success:
|
288
|
-
return ready
|
289
|
-
|
290
|
-
headers = {
|
291
|
-
"Authorization": f"Bearer {self.token}",
|
292
|
-
"Content-Type": "application/json"
|
293
|
-
}
|
294
|
-
|
295
|
-
data = {
|
296
|
-
"command": command,
|
297
|
-
"args": args or [],
|
298
|
-
"timeout": timeout
|
299
|
-
}
|
300
|
-
|
301
|
-
try:
|
302
|
-
async with aiohttp.ClientSession() as session:
|
303
|
-
async with session.post(
|
304
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/shell",
|
305
|
-
headers=headers,
|
306
|
-
json=data
|
307
|
-
) as response:
|
308
|
-
if response.status != 200:
|
309
|
-
error_text = await response.text()
|
310
|
-
return SandboxResponse(success=False, error=error_text)
|
311
|
-
|
312
|
-
# Parse the response
|
313
|
-
result = await response.json()
|
314
|
-
return SandboxResponse(success=True, data=result)
|
315
|
-
|
316
|
-
except Exception as e:
|
317
|
-
return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
|
255
|
+
return await self.shell.execute(command, args, timeout)
|
318
256
|
|
319
|
-
async def cleanup(self) -> SandboxResponse:
|
320
|
-
"""Delete the sandbox.
|
321
|
-
|
322
|
-
Returns:
|
323
|
-
SandboxResponse indicating success/failure of cleanup
|
324
|
-
"""
|
325
|
-
if not self.token or not self.sandbox_id:
|
326
|
-
return SandboxResponse(success=True)
|
327
|
-
|
328
|
-
headers = {"Authorization": f"Bearer {self.token}"}
|
329
|
-
async with aiohttp.ClientSession() as session:
|
330
|
-
async with session.delete(
|
331
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}",
|
332
|
-
headers=headers
|
333
|
-
) as response:
|
334
|
-
if response.status == 200:
|
335
|
-
data = await response.json()
|
336
|
-
self.sandbox_id = None
|
337
|
-
return SandboxResponse(success=True, data=data)
|
338
|
-
else:
|
339
|
-
text = await response.text()
|
340
|
-
return SandboxResponse(success=False, error=text)
|
341
|
-
|
342
257
|
async def upload_file(
|
343
258
|
self,
|
344
259
|
file_path: Union[str, Path],
|
345
260
|
destination: str = "/workspace",
|
346
|
-
chunk_size: int = 1024 * 1024,
|
347
|
-
timeout: int = 300
|
261
|
+
chunk_size: int = 1024 * 1024,
|
262
|
+
timeout: int = 300
|
348
263
|
) -> FileOperationResponse:
|
349
264
|
"""Upload a file to the sandbox environment.
|
350
265
|
|
266
|
+
This is a backward compatibility method that delegates to the fs submodule.
|
267
|
+
|
351
268
|
Args:
|
352
269
|
file_path: Path to the file to upload
|
353
270
|
destination: Destination path in the sandbox (absolute path starting with /)
|
@@ -357,222 +274,40 @@ class SandboxClient:
|
|
357
274
|
Returns:
|
358
275
|
FileOperationResponse containing upload results
|
359
276
|
"""
|
360
|
-
|
361
|
-
|
362
|
-
success=False,
|
363
|
-
error="Client not properly initialized. Call setup() first"
|
364
|
-
)
|
365
|
-
|
366
|
-
# Ensure sandbox is ready
|
367
|
-
ready = await self.wait_for_ready()
|
368
|
-
if not ready.success:
|
369
|
-
return FileOperationResponse(
|
370
|
-
success=False,
|
371
|
-
error=ready.error or "Sandbox not ready"
|
372
|
-
)
|
373
|
-
|
374
|
-
# Convert to Path object and validate file
|
375
|
-
file_path = Path(file_path)
|
376
|
-
if not file_path.exists():
|
377
|
-
return FileOperationResponse(
|
378
|
-
success=False,
|
379
|
-
error=f"File not found: {file_path}"
|
380
|
-
)
|
381
|
-
|
382
|
-
if not file_path.is_file():
|
383
|
-
return FileOperationResponse(
|
384
|
-
success=False,
|
385
|
-
error=f"Not a file: {file_path}"
|
386
|
-
)
|
387
|
-
|
388
|
-
# Get file size and validate
|
389
|
-
file_size = file_path.stat().st_size
|
390
|
-
if file_size > 100 * 1024 * 1024: # 100MB limit
|
391
|
-
return FileOperationResponse(
|
392
|
-
success=False,
|
393
|
-
error="File too large. Maximum size is 100MB"
|
394
|
-
)
|
395
|
-
|
396
|
-
try:
|
397
|
-
# Prepare the upload
|
398
|
-
headers = {
|
399
|
-
"Authorization": f"Bearer {self.token}"
|
400
|
-
}
|
401
|
-
|
402
|
-
# Guess content type
|
403
|
-
content_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream'
|
404
|
-
|
405
|
-
# Prepare multipart form data
|
406
|
-
data = aiohttp.FormData()
|
407
|
-
data.add_field('file',
|
408
|
-
open(file_path, 'rb').read(),
|
409
|
-
filename=file_path.name,
|
410
|
-
content_type=content_type)
|
411
|
-
data.add_field('path', destination)
|
412
|
-
|
413
|
-
timeout_settings = aiohttp.ClientTimeout(
|
414
|
-
total=timeout,
|
415
|
-
connect=30,
|
416
|
-
sock_connect=30,
|
417
|
-
sock_read=timeout
|
418
|
-
)
|
419
|
-
|
420
|
-
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
421
|
-
async with session.post(
|
422
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/upload",
|
423
|
-
headers=headers,
|
424
|
-
data=data
|
425
|
-
) as response:
|
426
|
-
if response.status != 200:
|
427
|
-
error_text = await response.text()
|
428
|
-
return FileOperationResponse(
|
429
|
-
success=False,
|
430
|
-
error=f"Upload failed: {error_text}"
|
431
|
-
)
|
432
|
-
|
433
|
-
result = await response.json()
|
434
|
-
return FileOperationResponse(
|
435
|
-
success=True,
|
436
|
-
filename=result.get("filename"),
|
437
|
-
size=result.get("size"),
|
438
|
-
path=result.get("path"),
|
439
|
-
message=result.get("message")
|
440
|
-
)
|
441
|
-
|
442
|
-
except asyncio.TimeoutError:
|
443
|
-
return FileOperationResponse(
|
444
|
-
success=False,
|
445
|
-
error=f"Upload timed out after {timeout} seconds"
|
446
|
-
)
|
447
|
-
except Exception as e:
|
448
|
-
return FileOperationResponse(
|
449
|
-
success=False,
|
450
|
-
error=f"Upload failed: {str(e)}"
|
451
|
-
)
|
452
|
-
|
277
|
+
return await self.fs.upload_file(file_path, destination, chunk_size, timeout)
|
278
|
+
|
453
279
|
async def download_file(
|
454
280
|
self,
|
455
|
-
|
281
|
+
remote_path: str,
|
456
282
|
local_path: Optional[Union[str, Path]] = None,
|
457
|
-
|
458
|
-
timeout: int = 300 # 5 minutes
|
283
|
+
timeout: int = 300
|
459
284
|
) -> FileOperationResponse:
|
460
|
-
"""Download a file from the sandbox
|
285
|
+
"""Download a file from the sandbox.
|
286
|
+
|
287
|
+
This is a backward compatibility method that delegates to the fs submodule.
|
461
288
|
|
462
289
|
Args:
|
463
|
-
|
464
|
-
|
465
|
-
local_path: Local path to save the file (default: current directory with original filename)
|
466
|
-
chunk_size: Size of chunks for downloading large files
|
290
|
+
remote_path: Path to the file in the sandbox
|
291
|
+
local_path: Local path to save the file (if None, uses the filename from remote_path)
|
467
292
|
timeout: Maximum download time in seconds
|
468
293
|
|
469
294
|
Returns:
|
470
295
|
FileOperationResponse containing download results
|
471
296
|
"""
|
472
|
-
|
473
|
-
|
474
|
-
success=False,
|
475
|
-
error="Client not properly initialized. Call setup() first"
|
476
|
-
)
|
477
|
-
|
478
|
-
# Ensure sandbox is ready
|
479
|
-
ready = await self.wait_for_ready()
|
480
|
-
if not ready.success:
|
481
|
-
return FileOperationResponse(
|
482
|
-
success=False,
|
483
|
-
error=ready.error or "Sandbox not ready"
|
484
|
-
)
|
485
|
-
|
486
|
-
# Ensure path is absolute and normalize any double slashes
|
487
|
-
if not sandbox_path.startswith('/'):
|
488
|
-
sandbox_path = f"/{sandbox_path}"
|
489
|
-
clean_path = '/'.join(part for part in sandbox_path.split('/') if part)
|
490
|
-
clean_path = f"/{clean_path}"
|
491
|
-
|
492
|
-
# Determine local path
|
493
|
-
if local_path is None:
|
494
|
-
local_path = Path(os.path.basename(sandbox_path))
|
495
|
-
else:
|
496
|
-
local_path = Path(local_path)
|
497
|
-
|
498
|
-
# Create parent directories if they don't exist
|
499
|
-
local_path.parent.mkdir(parents=True, exist_ok=True)
|
500
|
-
|
501
|
-
try:
|
502
|
-
timeout_settings = aiohttp.ClientTimeout(
|
503
|
-
total=timeout,
|
504
|
-
connect=30,
|
505
|
-
sock_connect=30,
|
506
|
-
sock_read=timeout
|
507
|
-
)
|
508
|
-
|
509
|
-
headers = {
|
510
|
-
"Authorization": f"Bearer {self.token}"
|
511
|
-
}
|
512
|
-
|
513
|
-
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
514
|
-
async with session.get(
|
515
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/download{clean_path}",
|
516
|
-
headers=headers
|
517
|
-
) as response:
|
518
|
-
if response.status != 200:
|
519
|
-
error_text = await response.text()
|
520
|
-
return FileOperationResponse(
|
521
|
-
success=False,
|
522
|
-
error=f"Download failed: {error_text}"
|
523
|
-
)
|
524
|
-
|
525
|
-
# Get content length if available
|
526
|
-
total_size = int(response.headers.get('content-length', 0))
|
527
|
-
|
528
|
-
# Download the file in chunks
|
529
|
-
downloaded_size = 0
|
530
|
-
try:
|
531
|
-
with open(local_path, 'wb') as f:
|
532
|
-
async for chunk in response.content.iter_chunked(chunk_size):
|
533
|
-
f.write(chunk)
|
534
|
-
downloaded_size += len(chunk)
|
535
|
-
|
536
|
-
return FileOperationResponse(
|
537
|
-
success=True,
|
538
|
-
filename=local_path.name,
|
539
|
-
size=downloaded_size or total_size,
|
540
|
-
path=str(local_path.absolute()),
|
541
|
-
message="File downloaded successfully"
|
542
|
-
)
|
543
|
-
except Exception as e:
|
544
|
-
# Clean up partial download
|
545
|
-
if local_path.exists():
|
546
|
-
local_path.unlink()
|
547
|
-
raise e
|
548
|
-
|
549
|
-
except asyncio.TimeoutError:
|
550
|
-
# Clean up partial download
|
551
|
-
if local_path.exists():
|
552
|
-
local_path.unlink()
|
553
|
-
return FileOperationResponse(
|
554
|
-
success=False,
|
555
|
-
error=f"Download timed out after {timeout} seconds"
|
556
|
-
)
|
557
|
-
except Exception as e:
|
558
|
-
# Clean up partial download
|
559
|
-
if local_path.exists():
|
560
|
-
local_path.unlink()
|
561
|
-
return FileOperationResponse(
|
562
|
-
success=False,
|
563
|
-
error=f"Download failed: {str(e)}"
|
564
|
-
)
|
565
|
-
|
297
|
+
return await self.fs.download_file(remote_path, local_path, timeout)
|
298
|
+
|
566
299
|
async def upload_bytes(
|
567
300
|
self,
|
568
301
|
content: Union[bytes, BinaryIO],
|
569
302
|
filename: str,
|
570
303
|
destination: str = "/workspace",
|
571
304
|
content_type: Optional[str] = None,
|
572
|
-
timeout: int = 300
|
305
|
+
timeout: int = 300
|
573
306
|
) -> FileOperationResponse:
|
574
307
|
"""Upload bytes or a file-like object to the sandbox environment.
|
575
308
|
|
309
|
+
This is a backward compatibility method that delegates to the fs submodule.
|
310
|
+
|
576
311
|
Args:
|
577
312
|
content: Bytes or file-like object to upload
|
578
313
|
filename: Name to give the file in the sandbox
|
@@ -583,179 +318,67 @@ class SandboxClient:
|
|
583
318
|
Returns:
|
584
319
|
FileOperationResponse containing upload results
|
585
320
|
"""
|
586
|
-
|
587
|
-
|
588
|
-
|
589
|
-
error="Client not properly initialized. Call setup() first"
|
590
|
-
)
|
591
|
-
|
592
|
-
# Ensure sandbox is ready
|
593
|
-
ready = await self.wait_for_ready()
|
594
|
-
if not ready.success:
|
595
|
-
return FileOperationResponse(
|
596
|
-
success=False,
|
597
|
-
error=ready.error or "Sandbox not ready"
|
598
|
-
)
|
599
|
-
|
600
|
-
try:
|
601
|
-
# Handle both bytes and file-like objects
|
321
|
+
# Create a temporary file with the content
|
322
|
+
import tempfile
|
323
|
+
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
|
602
324
|
if isinstance(content, bytes):
|
603
|
-
|
604
|
-
content_length = len(content)
|
325
|
+
temp_file.write(content)
|
605
326
|
else:
|
606
327
|
# Ensure we're at the start of the file
|
607
328
|
if hasattr(content, 'seek'):
|
608
329
|
content.seek(0)
|
609
|
-
|
610
|
-
|
611
|
-
|
612
|
-
|
613
|
-
|
614
|
-
current_pos = content.tell()
|
615
|
-
content.seek(0, os.SEEK_END)
|
616
|
-
content_length = content.tell()
|
617
|
-
content.seek(current_pos)
|
618
|
-
except (OSError, IOError):
|
619
|
-
pass
|
620
|
-
|
621
|
-
# Validate size if we can determine it
|
622
|
-
if content_length and content_length > 100 * 1024 * 1024: # 100MB limit
|
623
|
-
return FileOperationResponse(
|
624
|
-
success=False,
|
625
|
-
error="Content too large. Maximum size is 100MB"
|
626
|
-
)
|
627
|
-
|
628
|
-
# Prepare the upload
|
629
|
-
headers = {
|
630
|
-
"Authorization": f"Bearer {self.token}"
|
631
|
-
}
|
632
|
-
|
633
|
-
# Guess content type if not provided
|
634
|
-
if not content_type:
|
635
|
-
content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
|
636
|
-
|
637
|
-
# Prepare multipart form data
|
638
|
-
data = aiohttp.FormData()
|
639
|
-
data.add_field('file',
|
640
|
-
file_obj,
|
641
|
-
filename=filename,
|
642
|
-
content_type=content_type)
|
643
|
-
data.add_field('path', destination)
|
644
|
-
|
645
|
-
timeout_settings = aiohttp.ClientTimeout(
|
646
|
-
total=timeout,
|
647
|
-
connect=30,
|
648
|
-
sock_connect=30,
|
649
|
-
sock_read=timeout
|
650
|
-
)
|
651
|
-
|
652
|
-
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
653
|
-
async with session.post(
|
654
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/upload",
|
655
|
-
headers=headers,
|
656
|
-
data=data
|
657
|
-
) as response:
|
658
|
-
if response.status != 200:
|
659
|
-
error_text = await response.text()
|
660
|
-
return FileOperationResponse(
|
661
|
-
success=False,
|
662
|
-
error=f"Upload failed: {error_text}"
|
663
|
-
)
|
664
|
-
|
665
|
-
result = await response.json()
|
666
|
-
return FileOperationResponse(
|
667
|
-
success=True,
|
668
|
-
filename=result.get("filename"),
|
669
|
-
size=result.get("size"),
|
670
|
-
path=result.get("path"),
|
671
|
-
message=result.get("message")
|
672
|
-
)
|
673
|
-
|
674
|
-
except asyncio.TimeoutError:
|
675
|
-
return FileOperationResponse(
|
676
|
-
success=False,
|
677
|
-
error=f"Upload timed out after {timeout} seconds"
|
678
|
-
)
|
679
|
-
except Exception as e:
|
680
|
-
return FileOperationResponse(
|
681
|
-
success=False,
|
682
|
-
error=f"Upload failed: {str(e)}"
|
683
|
-
)
|
684
|
-
|
685
|
-
async def download_bytes(
|
686
|
-
self,
|
687
|
-
sandbox_path: str,
|
688
|
-
chunk_size: int = 8192, # 8KB chunks for download
|
689
|
-
timeout: int = 300 # 5 minutes
|
690
|
-
) -> Union[bytes, FileOperationResponse]:
|
691
|
-
"""Download a file from the sandbox environment into memory.
|
330
|
+
# Read and write in chunks to handle large files
|
331
|
+
chunk = content.read(1024 * 1024) # 1MB chunks
|
332
|
+
while chunk:
|
333
|
+
temp_file.write(chunk)
|
334
|
+
chunk = content.read(1024 * 1024)
|
692
335
|
|
693
|
-
|
694
|
-
|
695
|
-
|
696
|
-
|
697
|
-
|
336
|
+
try:
|
337
|
+
# Upload the temporary file
|
338
|
+
temp_path = Path(temp_file.name)
|
339
|
+
result = await self.fs.upload_file(
|
340
|
+
file_path=temp_path,
|
341
|
+
destination=os.path.join(destination, filename),
|
342
|
+
timeout=timeout
|
343
|
+
)
|
698
344
|
|
699
|
-
|
700
|
-
|
701
|
-
|
345
|
+
# If successful, update the filename in the response
|
346
|
+
if result.success:
|
347
|
+
result.filename = filename
|
348
|
+
|
349
|
+
return result
|
350
|
+
finally:
|
351
|
+
# Clean up the temporary file
|
352
|
+
if os.path.exists(temp_file.name):
|
353
|
+
os.unlink(temp_file.name)
|
354
|
+
|
355
|
+
async def download_bytes(self, remote_path: str, timeout: Optional[float] = None) -> Union[bytes, FileOperationResponse]:
|
702
356
|
"""
|
703
|
-
|
704
|
-
return FileOperationResponse(
|
705
|
-
success=False,
|
706
|
-
error="Client not properly initialized. Call setup() first"
|
707
|
-
)
|
357
|
+
Download a file from the sandbox into memory.
|
708
358
|
|
709
|
-
|
710
|
-
|
711
|
-
|
712
|
-
return FileOperationResponse(
|
713
|
-
success=False,
|
714
|
-
error=ready.error or "Sandbox not ready"
|
715
|
-
)
|
716
|
-
|
717
|
-
# Ensure path is absolute and normalize any double slashes
|
718
|
-
if not sandbox_path.startswith('/'):
|
719
|
-
sandbox_path = f"/{sandbox_path}"
|
720
|
-
clean_path = '/'.join(part for part in sandbox_path.split('/') if part)
|
721
|
-
clean_path = f"/{clean_path}"
|
359
|
+
Args:
|
360
|
+
remote_path: Path to the file in the sandbox.
|
361
|
+
timeout: Timeout in seconds for the operation.
|
722
362
|
|
363
|
+
Returns:
|
364
|
+
bytes: The file contents as bytes if successful.
|
365
|
+
FileOperationResponse: On failure, returns a FileOperationResponse with error details.
|
366
|
+
"""
|
367
|
+
await self.wait_for_ready()
|
368
|
+
|
723
369
|
try:
|
724
|
-
|
725
|
-
|
726
|
-
|
727
|
-
|
728
|
-
|
729
|
-
|
730
|
-
|
731
|
-
|
732
|
-
"Authorization": f"Bearer {self.token}"
|
733
|
-
}
|
734
|
-
|
735
|
-
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
736
|
-
async with session.get(
|
737
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/download{clean_path}",
|
738
|
-
headers=headers
|
739
|
-
) as response:
|
740
|
-
if response.status != 200:
|
741
|
-
error_text = await response.text()
|
742
|
-
return FileOperationResponse(
|
743
|
-
success=False,
|
744
|
-
error=f"Download failed: {error_text}"
|
745
|
-
)
|
746
|
-
|
747
|
-
# Read the entire response into memory
|
748
|
-
content = await response.read()
|
749
|
-
|
750
|
-
return content
|
751
|
-
|
752
|
-
except asyncio.TimeoutError:
|
753
|
-
return FileOperationResponse(
|
754
|
-
success=False,
|
755
|
-
error=f"Download timed out after {timeout} seconds"
|
756
|
-
)
|
370
|
+
response = await self.fs.read_file(remote_path, encoding=None)
|
371
|
+
if response.success:
|
372
|
+
return response.data.get('content')
|
373
|
+
else:
|
374
|
+
return FileOperationResponse(
|
375
|
+
success=False,
|
376
|
+
error=response.error or "Failed to download file"
|
377
|
+
)
|
757
378
|
except Exception as e:
|
758
379
|
return FileOperationResponse(
|
759
380
|
success=False,
|
760
|
-
error=f"
|
761
|
-
)
|
381
|
+
error=f"Error downloading file: {str(e)}"
|
382
|
+
)
|
383
|
+
|
384
|
+
# Additional backward compatibility methods can be added as needed
|