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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ai-computer-client
3
- Version: 0.2.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,4 @@
1
+ from .client import SandboxClient, SandboxResponse, StreamEvent, FileOperationResponse
2
+
3
+ __version__ = "0.3.0"
4
+ __all__ = ["SandboxClient", "SandboxResponse", "StreamEvent", "FileOperationResponse"]
@@ -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
+ )
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "ai-computer-client"
7
- version = "0.2.0"
7
+ version = "0.3.0"
8
8
  description = "Python client for interacting with the AI Computer service"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.7"
@@ -1,4 +0,0 @@
1
- from .client import SandboxClient, SandboxResponse, StreamEvent
2
-
3
- __version__ = "0.1.0"
4
- __all__ = ["SandboxClient", "SandboxResponse", "StreamEvent"]
@@ -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)