ai-computer-client 0.3.1__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/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
@@ -63,13 +71,33 @@ class SandboxClient:
63
71
 
64
72
  def __init__(
65
73
  self,
66
- base_url: str = "http://aicomputer.dev",
74
+ base_url: str = "http://api.aicomputer.dev",
67
75
  token: Optional[str] = None
68
76
  ):
69
77
  self.base_url = base_url.rstrip('/')
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, max_retries: int = 30, delay: int = 1) -> SandboxResponse:
111
- """Wait for the sandbox to be in Running state.
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
- max_retries: Maximum number of status check attempts
115
- delay: Delay between retries in seconds
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 _ 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:
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
- if data["status"] == "Running":
164
+ status = data.get("status")
165
+
166
+ if status == "ready":
134
167
  return SandboxResponse(success=True, data=data)
135
- await asyncio.sleep(delay)
136
-
137
- return SandboxResponse(success=False, error="Sandbox failed to become ready")
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 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.
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 containing execution results
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
- # 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
- }
192
+ headers = {"Authorization": f"Bearer {self.token}"}
173
193
 
174
194
  try:
175
195
  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
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
- error_text = await response.text()
183
- return SandboxResponse(success=False, error=error_text)
201
+ text = await response.text()
202
+ return SandboxResponse(success=False, error=text)
184
203
 
185
- # Parse the response
186
- result = await response.json()
187
- return SandboxResponse(success=True, data=result)
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
- 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.
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 method returns an async generator that yields StreamEvent objects containing
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
- Yields:
207
- StreamEvent objects with execution output/events
222
+ Returns:
223
+ SandboxResponse containing execution results
208
224
  """
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
- }
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
- data = {
225
- "code": code,
226
- "timeout": timeout
227
- }
230
+ This is a backward compatibility method that delegates to the code submodule.
228
231
 
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
- )
232
+ Args:
233
+ code: The Python code to execute
234
+ timeout: Maximum execution time in seconds
237
235
 
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)}")
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
- 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)}")
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, # 1MB chunks
347
- timeout: int = 300 # 5 minutes
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
- 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').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
- sandbox_path: str,
281
+ remote_path: str,
456
282
  local_path: Optional[Union[str, Path]] = None,
457
- chunk_size: int = 8192, # 8KB chunks for download
458
- timeout: int = 300 # 5 minutes
283
+ timeout: int = 300
459
284
  ) -> FileOperationResponse:
460
- """Download a file from the sandbox environment.
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
- sandbox_path: Path to the file in the sandbox (must be an absolute path starting with /).
464
- Any double slashes in the path will be normalized.
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
- if not self.token or not self.sandbox_id:
473
- return FileOperationResponse(
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 # 5 minutes
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
- if not self.token or not self.sandbox_id:
587
- return FileOperationResponse(
588
- success=False,
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
- file_obj = content
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
- file_obj = content
610
- # Try to get content length if possible
611
- content_length = None
612
- if hasattr(content, 'seek') and hasattr(content, 'tell'):
613
- try:
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
- Args:
694
- sandbox_path: Path to the file in the sandbox (must be an absolute path starting with /).
695
- Any double slashes in the path will be normalized.
696
- chunk_size: Size of chunks for downloading large files
697
- timeout: Maximum download time in seconds
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
- Returns:
700
- On success: The file contents as bytes
701
- On failure: FileOperationResponse with error details
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
- if not self.token or not self.sandbox_id:
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
- # Ensure sandbox is ready
710
- ready = await self.wait_for_ready()
711
- if not ready.success:
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
- timeout_settings = aiohttp.ClientTimeout(
725
- total=timeout,
726
- connect=30,
727
- sock_connect=30,
728
- sock_read=timeout
729
- )
730
-
731
- headers = {
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"Download failed: {str(e)}"
761
- )
381
+ error=f"Error downloading file: {str(e)}"
382
+ )
383
+
384
+ # Additional backward compatibility methods can be added as needed