ai-computer-client 0.2.0__tar.gz → 0.3.0__tar.gz
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_client-0.2.0 → ai_computer_client-0.3.0}/PKG-INFO +1 -1
- ai_computer_client-0.3.0/ai_computer/__init__.py +4 -0
- ai_computer_client-0.3.0/ai_computer/client.py +753 -0
- {ai_computer_client-0.2.0 → ai_computer_client-0.3.0}/pyproject.toml +1 -1
- ai_computer_client-0.2.0/ai_computer/__init__.py +0 -4
- ai_computer_client-0.2.0/ai_computer/client.py +0 -318
- {ai_computer_client-0.2.0 → ai_computer_client-0.3.0}/.gitignore +0 -0
- {ai_computer_client-0.2.0 → ai_computer_client-0.3.0}/LICENSE +0 -0
- {ai_computer_client-0.2.0 → ai_computer_client-0.3.0}/README.md +0 -0
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: ai-computer-client
|
3
|
-
Version: 0.
|
3
|
+
Version: 0.3.0
|
4
4
|
Summary: Python client for interacting with the AI Computer service
|
5
5
|
Project-URL: Homepage, https://github.com/ColeMurray/ai-computer-client-python
|
6
6
|
Project-URL: Documentation, https://github.com/ColeMurray/ai-computer-client-python#readme
|
@@ -0,0 +1,753 @@
|
|
1
|
+
import aiohttp
|
2
|
+
import json
|
3
|
+
import asyncio
|
4
|
+
from typing import Optional, Dict, AsyncGenerator, Union, List, BinaryIO
|
5
|
+
from dataclasses import dataclass
|
6
|
+
import os
|
7
|
+
import mimetypes
|
8
|
+
from pathlib import Path
|
9
|
+
|
10
|
+
@dataclass
|
11
|
+
class SandboxResponse:
|
12
|
+
"""Response from sandbox operations.
|
13
|
+
|
14
|
+
Attributes:
|
15
|
+
success: Whether the operation was successful
|
16
|
+
data: Optional response data
|
17
|
+
error: Optional error message if operation failed
|
18
|
+
"""
|
19
|
+
success: bool
|
20
|
+
data: Optional[Dict] = None
|
21
|
+
error: Optional[str] = None
|
22
|
+
|
23
|
+
@dataclass
|
24
|
+
class StreamEvent:
|
25
|
+
"""Event from streaming code execution.
|
26
|
+
|
27
|
+
Attributes:
|
28
|
+
type: Type of event ('stdout', 'stderr', 'info', 'error', 'completed', 'keepalive')
|
29
|
+
data: Event data
|
30
|
+
"""
|
31
|
+
type: str
|
32
|
+
data: str
|
33
|
+
|
34
|
+
@dataclass
|
35
|
+
class FileOperationResponse:
|
36
|
+
"""Response from file operations.
|
37
|
+
|
38
|
+
Attributes:
|
39
|
+
success: Whether the operation was successful
|
40
|
+
filename: Name of the file
|
41
|
+
size: Size of the file in bytes
|
42
|
+
path: Path where the file was saved
|
43
|
+
message: Optional status message
|
44
|
+
error: Optional error message if operation failed
|
45
|
+
"""
|
46
|
+
success: bool
|
47
|
+
filename: Optional[str] = None
|
48
|
+
size: Optional[int] = None
|
49
|
+
path: Optional[str] = None
|
50
|
+
message: Optional[str] = None
|
51
|
+
error: Optional[str] = None
|
52
|
+
|
53
|
+
class SandboxClient:
|
54
|
+
"""Client for interacting with the AI Sandbox service.
|
55
|
+
|
56
|
+
This client provides methods to execute Python code in an isolated sandbox environment.
|
57
|
+
It handles authentication, sandbox creation/deletion, and code execution.
|
58
|
+
|
59
|
+
Args:
|
60
|
+
base_url: The base URL of the sandbox service
|
61
|
+
token: Optional pre-existing authentication token
|
62
|
+
"""
|
63
|
+
|
64
|
+
def __init__(
|
65
|
+
self,
|
66
|
+
base_url: str = "http://aicomputer.dev",
|
67
|
+
token: Optional[str] = None
|
68
|
+
):
|
69
|
+
self.base_url = base_url.rstrip('/')
|
70
|
+
self.token = token
|
71
|
+
self.sandbox_id = None
|
72
|
+
|
73
|
+
async def setup(self) -> SandboxResponse:
|
74
|
+
"""Initialize the client and create a sandbox.
|
75
|
+
|
76
|
+
This method:
|
77
|
+
1. Gets a development token (if not provided)
|
78
|
+
2. Creates a new sandbox
|
79
|
+
3. Waits for the sandbox to be ready
|
80
|
+
|
81
|
+
Returns:
|
82
|
+
SandboxResponse indicating success/failure
|
83
|
+
"""
|
84
|
+
async with aiohttp.ClientSession() as session:
|
85
|
+
# Get development token if not provided
|
86
|
+
if not self.token:
|
87
|
+
async with session.post(f"{self.base_url}/dev/token") as response:
|
88
|
+
if response.status == 200:
|
89
|
+
data = await response.json()
|
90
|
+
self.token = data["access_token"]
|
91
|
+
else:
|
92
|
+
text = await response.text()
|
93
|
+
return SandboxResponse(success=False, error=text)
|
94
|
+
|
95
|
+
# Create sandbox
|
96
|
+
headers = {"Authorization": f"Bearer {self.token}"}
|
97
|
+
async with session.post(f"{self.base_url}/api/v1/sandbox/create", headers=headers) as response:
|
98
|
+
if response.status == 200:
|
99
|
+
data = await response.json()
|
100
|
+
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
|
+
else:
|
107
|
+
text = await response.text()
|
108
|
+
return SandboxResponse(success=False, error=text)
|
109
|
+
|
110
|
+
async def wait_for_ready(self, max_retries: int = 30, delay: int = 1) -> SandboxResponse:
|
111
|
+
"""Wait for the sandbox to be in Running state.
|
112
|
+
|
113
|
+
Args:
|
114
|
+
max_retries: Maximum number of status check attempts
|
115
|
+
delay: Delay between retries in seconds
|
116
|
+
|
117
|
+
Returns:
|
118
|
+
SandboxResponse indicating if sandbox is ready
|
119
|
+
"""
|
120
|
+
if not self.token or not self.sandbox_id:
|
121
|
+
return SandboxResponse(success=False, error="Client not properly initialized")
|
122
|
+
|
123
|
+
headers = {"Authorization": f"Bearer {self.token}"}
|
124
|
+
|
125
|
+
for _ in range(max_retries):
|
126
|
+
async with aiohttp.ClientSession() as session:
|
127
|
+
async with session.get(
|
128
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/status",
|
129
|
+
headers=headers
|
130
|
+
) as response:
|
131
|
+
if response.status == 200:
|
132
|
+
data = await response.json()
|
133
|
+
if data["status"] == "Running":
|
134
|
+
return SandboxResponse(success=True, data=data)
|
135
|
+
await asyncio.sleep(delay)
|
136
|
+
|
137
|
+
return SandboxResponse(success=False, error="Sandbox failed to become ready")
|
138
|
+
|
139
|
+
async def execute_code(
|
140
|
+
self,
|
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.
|
148
|
+
|
149
|
+
Args:
|
150
|
+
code: Python code to execute
|
151
|
+
timeout: Maximum execution time in seconds
|
152
|
+
|
153
|
+
Returns:
|
154
|
+
SandboxResponse containing execution results
|
155
|
+
"""
|
156
|
+
if not self.token or not self.sandbox_id:
|
157
|
+
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
158
|
+
|
159
|
+
# Ensure sandbox is ready
|
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
|
+
}
|
173
|
+
|
174
|
+
try:
|
175
|
+
async with aiohttp.ClientSession() as session:
|
176
|
+
async with session.post(
|
177
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute",
|
178
|
+
headers=headers,
|
179
|
+
json=data
|
180
|
+
) as response:
|
181
|
+
if response.status != 200:
|
182
|
+
error_text = await response.text()
|
183
|
+
return SandboxResponse(success=False, error=error_text)
|
184
|
+
|
185
|
+
# Parse the response
|
186
|
+
result = await response.json()
|
187
|
+
return SandboxResponse(success=True, data=result)
|
188
|
+
|
189
|
+
except Exception as e:
|
190
|
+
return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
|
191
|
+
|
192
|
+
async def execute_code_stream(
|
193
|
+
self,
|
194
|
+
code: Union[str, bytes],
|
195
|
+
timeout: int = 30
|
196
|
+
) -> AsyncGenerator[StreamEvent, None]:
|
197
|
+
"""Execute Python code in the sandbox and stream the output.
|
198
|
+
|
199
|
+
This method returns an async generator that yields StreamEvent objects containing
|
200
|
+
the type of event and the associated data.
|
201
|
+
|
202
|
+
Args:
|
203
|
+
code: Python code to execute
|
204
|
+
timeout: Maximum execution time in seconds
|
205
|
+
|
206
|
+
Yields:
|
207
|
+
StreamEvent objects with execution output/events
|
208
|
+
"""
|
209
|
+
if not self.token or not self.sandbox_id:
|
210
|
+
yield StreamEvent(type="error", data="Client not properly initialized. Call setup() first")
|
211
|
+
return
|
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
|
+
}
|
223
|
+
|
224
|
+
data = {
|
225
|
+
"code": code,
|
226
|
+
"timeout": timeout
|
227
|
+
}
|
228
|
+
|
229
|
+
try:
|
230
|
+
# Create a ClientTimeout object with all timeout settings
|
231
|
+
timeout_settings = aiohttp.ClientTimeout(
|
232
|
+
total=timeout + 30, # Add buffer for connection overhead
|
233
|
+
connect=30,
|
234
|
+
sock_connect=30,
|
235
|
+
sock_read=timeout + 30
|
236
|
+
)
|
237
|
+
|
238
|
+
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
239
|
+
async with session.post(
|
240
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/stream",
|
241
|
+
headers=headers,
|
242
|
+
json=data
|
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)}")
|
265
|
+
|
266
|
+
async def execute_shell(
|
267
|
+
self,
|
268
|
+
command: str,
|
269
|
+
args: Optional[List[str]] = None,
|
270
|
+
timeout: int = 30
|
271
|
+
) -> SandboxResponse:
|
272
|
+
"""Execute a shell command in the sandbox.
|
273
|
+
|
274
|
+
Args:
|
275
|
+
command: The shell command to execute
|
276
|
+
args: Optional list of arguments for the command
|
277
|
+
timeout: Maximum execution time in seconds
|
278
|
+
|
279
|
+
Returns:
|
280
|
+
SandboxResponse containing execution results
|
281
|
+
"""
|
282
|
+
if not self.token or not self.sandbox_id:
|
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)}")
|
318
|
+
|
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
|
+
async def upload_file(
|
343
|
+
self,
|
344
|
+
file_path: Union[str, Path],
|
345
|
+
destination: str = "/workspace",
|
346
|
+
chunk_size: int = 1024 * 1024, # 1MB chunks
|
347
|
+
timeout: int = 300 # 5 minutes
|
348
|
+
) -> FileOperationResponse:
|
349
|
+
"""Upload a file to the sandbox environment.
|
350
|
+
|
351
|
+
Args:
|
352
|
+
file_path: Path to the file to upload
|
353
|
+
destination: Destination path in the sandbox (default: /workspace)
|
354
|
+
chunk_size: Size of chunks for reading large files
|
355
|
+
timeout: Maximum upload time in seconds
|
356
|
+
|
357
|
+
Returns:
|
358
|
+
FileOperationResponse containing upload results
|
359
|
+
"""
|
360
|
+
if not self.token or not self.sandbox_id:
|
361
|
+
return FileOperationResponse(
|
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'),
|
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
|
+
|
453
|
+
async def download_file(
|
454
|
+
self,
|
455
|
+
sandbox_path: str,
|
456
|
+
local_path: Optional[Union[str, Path]] = None,
|
457
|
+
chunk_size: int = 8192, # 8KB chunks for download
|
458
|
+
timeout: int = 300 # 5 minutes
|
459
|
+
) -> FileOperationResponse:
|
460
|
+
"""Download a file from the sandbox environment.
|
461
|
+
|
462
|
+
Args:
|
463
|
+
sandbox_path: Path to the file in the sandbox
|
464
|
+
local_path: Local path to save the file (default: current directory with original filename)
|
465
|
+
chunk_size: Size of chunks for downloading large files
|
466
|
+
timeout: Maximum download time in seconds
|
467
|
+
|
468
|
+
Returns:
|
469
|
+
FileOperationResponse containing download results
|
470
|
+
"""
|
471
|
+
if not self.token or not self.sandbox_id:
|
472
|
+
return FileOperationResponse(
|
473
|
+
success=False,
|
474
|
+
error="Client not properly initialized. Call setup() first"
|
475
|
+
)
|
476
|
+
|
477
|
+
# Ensure sandbox is ready
|
478
|
+
ready = await self.wait_for_ready()
|
479
|
+
if not ready.success:
|
480
|
+
return FileOperationResponse(
|
481
|
+
success=False,
|
482
|
+
error=ready.error or "Sandbox not ready"
|
483
|
+
)
|
484
|
+
|
485
|
+
# Normalize sandbox path
|
486
|
+
sandbox_path = sandbox_path.lstrip('/')
|
487
|
+
|
488
|
+
# Determine local path
|
489
|
+
if local_path is None:
|
490
|
+
local_path = Path(os.path.basename(sandbox_path))
|
491
|
+
else:
|
492
|
+
local_path = Path(local_path)
|
493
|
+
|
494
|
+
# Create parent directories if they don't exist
|
495
|
+
local_path.parent.mkdir(parents=True, exist_ok=True)
|
496
|
+
|
497
|
+
try:
|
498
|
+
timeout_settings = aiohttp.ClientTimeout(
|
499
|
+
total=timeout,
|
500
|
+
connect=30,
|
501
|
+
sock_connect=30,
|
502
|
+
sock_read=timeout
|
503
|
+
)
|
504
|
+
|
505
|
+
headers = {
|
506
|
+
"Authorization": f"Bearer {self.token}"
|
507
|
+
}
|
508
|
+
|
509
|
+
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
510
|
+
async with session.get(
|
511
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/download/{sandbox_path}",
|
512
|
+
headers=headers
|
513
|
+
) as response:
|
514
|
+
if response.status != 200:
|
515
|
+
error_text = await response.text()
|
516
|
+
return FileOperationResponse(
|
517
|
+
success=False,
|
518
|
+
error=f"Download failed: {error_text}"
|
519
|
+
)
|
520
|
+
|
521
|
+
# Get content length if available
|
522
|
+
total_size = int(response.headers.get('content-length', 0))
|
523
|
+
|
524
|
+
# Download the file in chunks
|
525
|
+
downloaded_size = 0
|
526
|
+
try:
|
527
|
+
with open(local_path, 'wb') as f:
|
528
|
+
async for chunk in response.content.iter_chunked(chunk_size):
|
529
|
+
f.write(chunk)
|
530
|
+
downloaded_size += len(chunk)
|
531
|
+
|
532
|
+
return FileOperationResponse(
|
533
|
+
success=True,
|
534
|
+
filename=local_path.name,
|
535
|
+
size=downloaded_size or total_size,
|
536
|
+
path=str(local_path.absolute()),
|
537
|
+
message="File downloaded successfully"
|
538
|
+
)
|
539
|
+
except Exception as e:
|
540
|
+
# Clean up partial download
|
541
|
+
if local_path.exists():
|
542
|
+
local_path.unlink()
|
543
|
+
raise e
|
544
|
+
|
545
|
+
except asyncio.TimeoutError:
|
546
|
+
# Clean up partial download
|
547
|
+
if local_path.exists():
|
548
|
+
local_path.unlink()
|
549
|
+
return FileOperationResponse(
|
550
|
+
success=False,
|
551
|
+
error=f"Download timed out after {timeout} seconds"
|
552
|
+
)
|
553
|
+
except Exception as e:
|
554
|
+
# Clean up partial download
|
555
|
+
if local_path.exists():
|
556
|
+
local_path.unlink()
|
557
|
+
return FileOperationResponse(
|
558
|
+
success=False,
|
559
|
+
error=f"Download failed: {str(e)}"
|
560
|
+
)
|
561
|
+
|
562
|
+
async def upload_bytes(
|
563
|
+
self,
|
564
|
+
content: Union[bytes, BinaryIO],
|
565
|
+
filename: str,
|
566
|
+
destination: str = "/workspace",
|
567
|
+
content_type: Optional[str] = None,
|
568
|
+
timeout: int = 300 # 5 minutes
|
569
|
+
) -> FileOperationResponse:
|
570
|
+
"""Upload bytes or a file-like object to the sandbox environment.
|
571
|
+
|
572
|
+
Args:
|
573
|
+
content: Bytes or file-like object to upload
|
574
|
+
filename: Name to give the file in the sandbox
|
575
|
+
destination: Destination path in the sandbox (default: /workspace)
|
576
|
+
content_type: Optional MIME type (will be guessed from filename if not provided)
|
577
|
+
timeout: Maximum upload time in seconds
|
578
|
+
|
579
|
+
Returns:
|
580
|
+
FileOperationResponse containing upload results
|
581
|
+
"""
|
582
|
+
if not self.token or not self.sandbox_id:
|
583
|
+
return FileOperationResponse(
|
584
|
+
success=False,
|
585
|
+
error="Client not properly initialized. Call setup() first"
|
586
|
+
)
|
587
|
+
|
588
|
+
# Ensure sandbox is ready
|
589
|
+
ready = await self.wait_for_ready()
|
590
|
+
if not ready.success:
|
591
|
+
return FileOperationResponse(
|
592
|
+
success=False,
|
593
|
+
error=ready.error or "Sandbox not ready"
|
594
|
+
)
|
595
|
+
|
596
|
+
try:
|
597
|
+
# Handle both bytes and file-like objects
|
598
|
+
if isinstance(content, bytes):
|
599
|
+
file_obj = content
|
600
|
+
content_length = len(content)
|
601
|
+
else:
|
602
|
+
# Ensure we're at the start of the file
|
603
|
+
if hasattr(content, 'seek'):
|
604
|
+
content.seek(0)
|
605
|
+
file_obj = content
|
606
|
+
# Try to get content length if possible
|
607
|
+
content_length = None
|
608
|
+
if hasattr(content, 'seek') and hasattr(content, 'tell'):
|
609
|
+
try:
|
610
|
+
current_pos = content.tell()
|
611
|
+
content.seek(0, os.SEEK_END)
|
612
|
+
content_length = content.tell()
|
613
|
+
content.seek(current_pos)
|
614
|
+
except (OSError, IOError):
|
615
|
+
pass
|
616
|
+
|
617
|
+
# Validate size if we can determine it
|
618
|
+
if content_length and content_length > 100 * 1024 * 1024: # 100MB limit
|
619
|
+
return FileOperationResponse(
|
620
|
+
success=False,
|
621
|
+
error="Content too large. Maximum size is 100MB"
|
622
|
+
)
|
623
|
+
|
624
|
+
# Prepare the upload
|
625
|
+
headers = {
|
626
|
+
"Authorization": f"Bearer {self.token}"
|
627
|
+
}
|
628
|
+
|
629
|
+
# Guess content type if not provided
|
630
|
+
if not content_type:
|
631
|
+
content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
|
632
|
+
|
633
|
+
# Prepare multipart form data
|
634
|
+
data = aiohttp.FormData()
|
635
|
+
data.add_field('file',
|
636
|
+
file_obj,
|
637
|
+
filename=filename,
|
638
|
+
content_type=content_type)
|
639
|
+
data.add_field('path', destination)
|
640
|
+
|
641
|
+
timeout_settings = aiohttp.ClientTimeout(
|
642
|
+
total=timeout,
|
643
|
+
connect=30,
|
644
|
+
sock_connect=30,
|
645
|
+
sock_read=timeout
|
646
|
+
)
|
647
|
+
|
648
|
+
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
649
|
+
async with session.post(
|
650
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/upload",
|
651
|
+
headers=headers,
|
652
|
+
data=data
|
653
|
+
) as response:
|
654
|
+
if response.status != 200:
|
655
|
+
error_text = await response.text()
|
656
|
+
return FileOperationResponse(
|
657
|
+
success=False,
|
658
|
+
error=f"Upload failed: {error_text}"
|
659
|
+
)
|
660
|
+
|
661
|
+
result = await response.json()
|
662
|
+
return FileOperationResponse(
|
663
|
+
success=True,
|
664
|
+
filename=result.get("filename"),
|
665
|
+
size=result.get("size"),
|
666
|
+
path=result.get("path"),
|
667
|
+
message=result.get("message")
|
668
|
+
)
|
669
|
+
|
670
|
+
except asyncio.TimeoutError:
|
671
|
+
return FileOperationResponse(
|
672
|
+
success=False,
|
673
|
+
error=f"Upload timed out after {timeout} seconds"
|
674
|
+
)
|
675
|
+
except Exception as e:
|
676
|
+
return FileOperationResponse(
|
677
|
+
success=False,
|
678
|
+
error=f"Upload failed: {str(e)}"
|
679
|
+
)
|
680
|
+
|
681
|
+
async def download_bytes(
|
682
|
+
self,
|
683
|
+
sandbox_path: str,
|
684
|
+
chunk_size: int = 8192, # 8KB chunks for download
|
685
|
+
timeout: int = 300 # 5 minutes
|
686
|
+
) -> Union[bytes, FileOperationResponse]:
|
687
|
+
"""Download a file from the sandbox environment into memory.
|
688
|
+
|
689
|
+
Args:
|
690
|
+
sandbox_path: Path to the file in the sandbox
|
691
|
+
chunk_size: Size of chunks for downloading large files
|
692
|
+
timeout: Maximum download time in seconds
|
693
|
+
|
694
|
+
Returns:
|
695
|
+
On success: The file contents as bytes
|
696
|
+
On failure: FileOperationResponse with error details
|
697
|
+
"""
|
698
|
+
if not self.token or not self.sandbox_id:
|
699
|
+
return FileOperationResponse(
|
700
|
+
success=False,
|
701
|
+
error="Client not properly initialized. Call setup() first"
|
702
|
+
)
|
703
|
+
|
704
|
+
# Ensure sandbox is ready
|
705
|
+
ready = await self.wait_for_ready()
|
706
|
+
if not ready.success:
|
707
|
+
return FileOperationResponse(
|
708
|
+
success=False,
|
709
|
+
error=ready.error or "Sandbox not ready"
|
710
|
+
)
|
711
|
+
|
712
|
+
# Normalize sandbox path
|
713
|
+
sandbox_path = sandbox_path.lstrip('/')
|
714
|
+
|
715
|
+
try:
|
716
|
+
timeout_settings = aiohttp.ClientTimeout(
|
717
|
+
total=timeout,
|
718
|
+
connect=30,
|
719
|
+
sock_connect=30,
|
720
|
+
sock_read=timeout
|
721
|
+
)
|
722
|
+
|
723
|
+
headers = {
|
724
|
+
"Authorization": f"Bearer {self.token}"
|
725
|
+
}
|
726
|
+
|
727
|
+
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
728
|
+
async with session.get(
|
729
|
+
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/files/download/{sandbox_path}",
|
730
|
+
headers=headers
|
731
|
+
) as response:
|
732
|
+
if response.status != 200:
|
733
|
+
error_text = await response.text()
|
734
|
+
return FileOperationResponse(
|
735
|
+
success=False,
|
736
|
+
error=f"Download failed: {error_text}"
|
737
|
+
)
|
738
|
+
|
739
|
+
# Read the entire response into memory
|
740
|
+
content = await response.read()
|
741
|
+
|
742
|
+
return content
|
743
|
+
|
744
|
+
except asyncio.TimeoutError:
|
745
|
+
return FileOperationResponse(
|
746
|
+
success=False,
|
747
|
+
error=f"Download timed out after {timeout} seconds"
|
748
|
+
)
|
749
|
+
except Exception as e:
|
750
|
+
return FileOperationResponse(
|
751
|
+
success=False,
|
752
|
+
error=f"Download failed: {str(e)}"
|
753
|
+
)
|
@@ -1,318 +0,0 @@
|
|
1
|
-
import aiohttp
|
2
|
-
import json
|
3
|
-
import asyncio
|
4
|
-
from typing import Optional, Dict, AsyncGenerator, Union, List
|
5
|
-
from dataclasses import dataclass
|
6
|
-
|
7
|
-
@dataclass
|
8
|
-
class SandboxResponse:
|
9
|
-
"""Response from sandbox operations.
|
10
|
-
|
11
|
-
Attributes:
|
12
|
-
success: Whether the operation was successful
|
13
|
-
data: Optional response data
|
14
|
-
error: Optional error message if operation failed
|
15
|
-
"""
|
16
|
-
success: bool
|
17
|
-
data: Optional[Dict] = None
|
18
|
-
error: Optional[str] = None
|
19
|
-
|
20
|
-
@dataclass
|
21
|
-
class StreamEvent:
|
22
|
-
"""Event from streaming code execution.
|
23
|
-
|
24
|
-
Attributes:
|
25
|
-
type: Type of event ('stdout', 'stderr', 'info', 'error', 'completed', 'keepalive')
|
26
|
-
data: Event data
|
27
|
-
"""
|
28
|
-
type: str
|
29
|
-
data: str
|
30
|
-
|
31
|
-
class SandboxClient:
|
32
|
-
"""Client for interacting with the AI Sandbox service.
|
33
|
-
|
34
|
-
This client provides methods to execute Python code in an isolated sandbox environment.
|
35
|
-
It handles authentication, sandbox creation/deletion, and code execution.
|
36
|
-
|
37
|
-
Args:
|
38
|
-
base_url: The base URL of the sandbox service
|
39
|
-
token: Optional pre-existing authentication token
|
40
|
-
"""
|
41
|
-
|
42
|
-
def __init__(
|
43
|
-
self,
|
44
|
-
base_url: str = "http://aicomputer.dev",
|
45
|
-
token: Optional[str] = None
|
46
|
-
):
|
47
|
-
self.base_url = base_url.rstrip('/')
|
48
|
-
self.token = token
|
49
|
-
self.sandbox_id = None
|
50
|
-
|
51
|
-
async def setup(self) -> SandboxResponse:
|
52
|
-
"""Initialize the client and create a sandbox.
|
53
|
-
|
54
|
-
This method:
|
55
|
-
1. Gets a development token (if not provided)
|
56
|
-
2. Creates a new sandbox
|
57
|
-
3. Waits for the sandbox to be ready
|
58
|
-
|
59
|
-
Returns:
|
60
|
-
SandboxResponse indicating success/failure
|
61
|
-
"""
|
62
|
-
async with aiohttp.ClientSession() as session:
|
63
|
-
# Get development token if not provided
|
64
|
-
if not self.token:
|
65
|
-
async with session.post(f"{self.base_url}/dev/token") as response:
|
66
|
-
if response.status == 200:
|
67
|
-
data = await response.json()
|
68
|
-
self.token = data["access_token"]
|
69
|
-
else:
|
70
|
-
text = await response.text()
|
71
|
-
return SandboxResponse(success=False, error=text)
|
72
|
-
|
73
|
-
# Create sandbox
|
74
|
-
headers = {"Authorization": f"Bearer {self.token}"}
|
75
|
-
async with session.post(f"{self.base_url}/api/v1/sandbox/create", headers=headers) as response:
|
76
|
-
if response.status == 200:
|
77
|
-
data = await response.json()
|
78
|
-
self.sandbox_id = data["sandbox_id"]
|
79
|
-
# Wait for sandbox to be ready
|
80
|
-
ready = await self.wait_for_ready()
|
81
|
-
if not ready.success:
|
82
|
-
return ready
|
83
|
-
return SandboxResponse(success=True, data=data)
|
84
|
-
else:
|
85
|
-
text = await response.text()
|
86
|
-
return SandboxResponse(success=False, error=text)
|
87
|
-
|
88
|
-
async def wait_for_ready(self, max_retries: int = 30, delay: int = 1) -> SandboxResponse:
|
89
|
-
"""Wait for the sandbox to be in Running state.
|
90
|
-
|
91
|
-
Args:
|
92
|
-
max_retries: Maximum number of status check attempts
|
93
|
-
delay: Delay between retries in seconds
|
94
|
-
|
95
|
-
Returns:
|
96
|
-
SandboxResponse indicating if sandbox is ready
|
97
|
-
"""
|
98
|
-
if not self.token or not self.sandbox_id:
|
99
|
-
return SandboxResponse(success=False, error="Client not properly initialized")
|
100
|
-
|
101
|
-
headers = {"Authorization": f"Bearer {self.token}"}
|
102
|
-
|
103
|
-
for _ in range(max_retries):
|
104
|
-
async with aiohttp.ClientSession() as session:
|
105
|
-
async with session.get(
|
106
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/status",
|
107
|
-
headers=headers
|
108
|
-
) as response:
|
109
|
-
if response.status == 200:
|
110
|
-
data = await response.json()
|
111
|
-
if data["status"] == "Running":
|
112
|
-
return SandboxResponse(success=True, data=data)
|
113
|
-
await asyncio.sleep(delay)
|
114
|
-
|
115
|
-
return SandboxResponse(success=False, error="Sandbox failed to become ready")
|
116
|
-
|
117
|
-
async def execute_code(
|
118
|
-
self,
|
119
|
-
code: Union[str, bytes],
|
120
|
-
timeout: int = 30
|
121
|
-
) -> SandboxResponse:
|
122
|
-
"""Execute Python code in the sandbox and return the combined output.
|
123
|
-
|
124
|
-
This method collects all output from the streaming response and returns it as a single result.
|
125
|
-
It captures both stdout and stderr, and handles any errors during execution.
|
126
|
-
|
127
|
-
Args:
|
128
|
-
code: Python code to execute
|
129
|
-
timeout: Maximum execution time in seconds
|
130
|
-
|
131
|
-
Returns:
|
132
|
-
SandboxResponse containing execution results
|
133
|
-
"""
|
134
|
-
if not self.token or not self.sandbox_id:
|
135
|
-
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
136
|
-
|
137
|
-
# Ensure sandbox is ready
|
138
|
-
ready = await self.wait_for_ready()
|
139
|
-
if not ready.success:
|
140
|
-
return ready
|
141
|
-
|
142
|
-
headers = {
|
143
|
-
"Authorization": f"Bearer {self.token}",
|
144
|
-
"Content-Type": "application/json"
|
145
|
-
}
|
146
|
-
|
147
|
-
data = {
|
148
|
-
"code": code,
|
149
|
-
"timeout": timeout
|
150
|
-
}
|
151
|
-
|
152
|
-
try:
|
153
|
-
async with aiohttp.ClientSession() as session:
|
154
|
-
async with session.post(
|
155
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute",
|
156
|
-
headers=headers,
|
157
|
-
json=data
|
158
|
-
) as response:
|
159
|
-
if response.status != 200:
|
160
|
-
error_text = await response.text()
|
161
|
-
return SandboxResponse(success=False, error=error_text)
|
162
|
-
|
163
|
-
# Parse the response
|
164
|
-
result = await response.json()
|
165
|
-
return SandboxResponse(success=True, data=result)
|
166
|
-
|
167
|
-
except Exception as e:
|
168
|
-
return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
|
169
|
-
|
170
|
-
async def execute_code_stream(
|
171
|
-
self,
|
172
|
-
code: Union[str, bytes],
|
173
|
-
timeout: int = 30
|
174
|
-
) -> AsyncGenerator[StreamEvent, None]:
|
175
|
-
"""Execute Python code in the sandbox and stream the output.
|
176
|
-
|
177
|
-
This method returns an async generator that yields StreamEvent objects containing
|
178
|
-
the type of event and the associated data.
|
179
|
-
|
180
|
-
Args:
|
181
|
-
code: Python code to execute
|
182
|
-
timeout: Maximum execution time in seconds
|
183
|
-
|
184
|
-
Yields:
|
185
|
-
StreamEvent objects with execution output/events
|
186
|
-
"""
|
187
|
-
if not self.token or not self.sandbox_id:
|
188
|
-
yield StreamEvent(type="error", data="Client not properly initialized. Call setup() first")
|
189
|
-
return
|
190
|
-
|
191
|
-
# Ensure sandbox is ready
|
192
|
-
ready = await self.wait_for_ready()
|
193
|
-
if not ready.success:
|
194
|
-
yield StreamEvent(type="error", data=ready.error or "Sandbox not ready")
|
195
|
-
return
|
196
|
-
|
197
|
-
headers = {
|
198
|
-
"Authorization": f"Bearer {self.token}",
|
199
|
-
"Content-Type": "application/json"
|
200
|
-
}
|
201
|
-
|
202
|
-
data = {
|
203
|
-
"code": code,
|
204
|
-
"timeout": timeout
|
205
|
-
}
|
206
|
-
|
207
|
-
try:
|
208
|
-
# Create a ClientTimeout object with all timeout settings
|
209
|
-
timeout_settings = aiohttp.ClientTimeout(
|
210
|
-
total=timeout + 30, # Add buffer for connection overhead
|
211
|
-
connect=30,
|
212
|
-
sock_connect=30,
|
213
|
-
sock_read=timeout + 30
|
214
|
-
)
|
215
|
-
|
216
|
-
async with aiohttp.ClientSession(timeout=timeout_settings) as session:
|
217
|
-
async with session.post(
|
218
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/stream",
|
219
|
-
headers=headers,
|
220
|
-
json=data
|
221
|
-
) as response:
|
222
|
-
if response.status != 200:
|
223
|
-
error_text = await response.text()
|
224
|
-
yield StreamEvent(type="error", data=error_text)
|
225
|
-
return
|
226
|
-
|
227
|
-
# Process the streaming response
|
228
|
-
async for line in response.content:
|
229
|
-
if line:
|
230
|
-
try:
|
231
|
-
event = json.loads(line.decode())
|
232
|
-
yield StreamEvent(type=event['type'], data=event['data'])
|
233
|
-
|
234
|
-
# Stop if we receive an error or completed event
|
235
|
-
if event['type'] in ['error', 'completed']:
|
236
|
-
break
|
237
|
-
except json.JSONDecodeError as e:
|
238
|
-
yield StreamEvent(type="error", data=f"Failed to parse event: {str(e)}")
|
239
|
-
break
|
240
|
-
|
241
|
-
except Exception as e:
|
242
|
-
yield StreamEvent(type="error", data=f"Connection error: {str(e)}")
|
243
|
-
|
244
|
-
async def execute_shell(
|
245
|
-
self,
|
246
|
-
command: str,
|
247
|
-
args: Optional[List[str]] = None,
|
248
|
-
timeout: int = 30
|
249
|
-
) -> SandboxResponse:
|
250
|
-
"""Execute a shell command in the sandbox.
|
251
|
-
|
252
|
-
Args:
|
253
|
-
command: The shell command to execute
|
254
|
-
args: Optional list of arguments for the command
|
255
|
-
timeout: Maximum execution time in seconds
|
256
|
-
|
257
|
-
Returns:
|
258
|
-
SandboxResponse containing execution results
|
259
|
-
"""
|
260
|
-
if not self.token or not self.sandbox_id:
|
261
|
-
return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
|
262
|
-
|
263
|
-
# Ensure sandbox is ready
|
264
|
-
ready = await self.wait_for_ready()
|
265
|
-
if not ready.success:
|
266
|
-
return ready
|
267
|
-
|
268
|
-
headers = {
|
269
|
-
"Authorization": f"Bearer {self.token}",
|
270
|
-
"Content-Type": "application/json"
|
271
|
-
}
|
272
|
-
|
273
|
-
data = {
|
274
|
-
"command": command,
|
275
|
-
"args": args or [],
|
276
|
-
"timeout": timeout
|
277
|
-
}
|
278
|
-
|
279
|
-
try:
|
280
|
-
async with aiohttp.ClientSession() as session:
|
281
|
-
async with session.post(
|
282
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/shell",
|
283
|
-
headers=headers,
|
284
|
-
json=data
|
285
|
-
) as response:
|
286
|
-
if response.status != 200:
|
287
|
-
error_text = await response.text()
|
288
|
-
return SandboxResponse(success=False, error=error_text)
|
289
|
-
|
290
|
-
# Parse the response
|
291
|
-
result = await response.json()
|
292
|
-
return SandboxResponse(success=True, data=result)
|
293
|
-
|
294
|
-
except Exception as e:
|
295
|
-
return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
|
296
|
-
|
297
|
-
async def cleanup(self) -> SandboxResponse:
|
298
|
-
"""Delete the sandbox.
|
299
|
-
|
300
|
-
Returns:
|
301
|
-
SandboxResponse indicating success/failure of cleanup
|
302
|
-
"""
|
303
|
-
if not self.token or not self.sandbox_id:
|
304
|
-
return SandboxResponse(success=True)
|
305
|
-
|
306
|
-
headers = {"Authorization": f"Bearer {self.token}"}
|
307
|
-
async with aiohttp.ClientSession() as session:
|
308
|
-
async with session.delete(
|
309
|
-
f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}",
|
310
|
-
headers=headers
|
311
|
-
) as response:
|
312
|
-
if response.status == 200:
|
313
|
-
data = await response.json()
|
314
|
-
self.sandbox_id = None
|
315
|
-
return SandboxResponse(success=True, data=data)
|
316
|
-
else:
|
317
|
-
text = await response.text()
|
318
|
-
return SandboxResponse(success=False, error=text)
|
File without changes
|
File without changes
|
File without changes
|