ai-computer-client 0.3.2__tar.gz → 0.3.3__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.3.2
3
+ Version: 0.3.3
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,14 @@
1
+ from .client import SandboxClient
2
+ from .models import SandboxResponse, StreamEvent, FileOperationResponse
3
+ from .submodules import FileSystemModule, ShellModule, CodeModule
4
+
5
+ __version__ = "0.3.3"
6
+ __all__ = [
7
+ "SandboxClient",
8
+ "SandboxResponse",
9
+ "StreamEvent",
10
+ "FileOperationResponse",
11
+ "FileSystemModule",
12
+ "ShellModule",
13
+ "CodeModule"
14
+ ]
@@ -0,0 +1,384 @@
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
+ from .models import SandboxResponse, StreamEvent, FileOperationResponse
11
+ from .submodules import FileSystemModule, ShellModule, CodeModule
12
+
13
+ @dataclass
14
+ class SandboxResponse:
15
+ """Response from sandbox operations.
16
+
17
+ Attributes:
18
+ success: Whether the operation was successful
19
+ data: Optional response data
20
+ error: Optional error message if operation failed
21
+ """
22
+ success: bool
23
+ data: Optional[Dict] = None
24
+ error: Optional[str] = None
25
+
26
+ @dataclass
27
+ class StreamEvent:
28
+ """Event from streaming code execution.
29
+
30
+ Attributes:
31
+ type: Type of event ('stdout', 'stderr', 'info', 'error', 'completed', 'keepalive')
32
+ data: Event data
33
+ """
34
+ type: str
35
+ data: str
36
+
37
+ @dataclass
38
+ class FileOperationResponse:
39
+ """Response from file operations.
40
+
41
+ Attributes:
42
+ success: Whether the operation was successful
43
+ filename: Name of the file
44
+ size: Size of the file in bytes
45
+ path: Path where the file was saved
46
+ message: Optional status message
47
+ error: Optional error message if operation failed
48
+ """
49
+ success: bool
50
+ filename: Optional[str] = None
51
+ size: Optional[int] = None
52
+ path: Optional[str] = None
53
+ message: Optional[str] = None
54
+ error: Optional[str] = None
55
+
56
+ class SandboxClient:
57
+ """Client for interacting with the AI Sandbox service.
58
+
59
+ This client provides methods to execute Python code in an isolated sandbox environment.
60
+ It handles authentication, sandbox creation/deletion, and code execution.
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
+
67
+ Args:
68
+ base_url: The base URL of the sandbox service
69
+ token: Optional pre-existing authentication token
70
+ """
71
+
72
+ def __init__(
73
+ self,
74
+ base_url: str = "http://api.aicomputer.dev",
75
+ token: Optional[str] = None
76
+ ):
77
+ self.base_url = base_url.rstrip('/')
78
+ self.token = token
79
+ self.sandbox_id = None
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
+
101
+ async def setup(self) -> SandboxResponse:
102
+ """Initialize the client and create a sandbox.
103
+
104
+ This method:
105
+ 1. Gets a development token (if not provided)
106
+ 2. Creates a new sandbox
107
+ 3. Waits for the sandbox to be ready
108
+
109
+ Returns:
110
+ SandboxResponse indicating success/failure
111
+ """
112
+ async with aiohttp.ClientSession() as session:
113
+ # Get development token if not provided
114
+ if not self.token:
115
+ async with session.post(f"{self.base_url}/dev/token") as response:
116
+ if response.status == 200:
117
+ data = await response.json()
118
+ self.token = data["access_token"]
119
+ else:
120
+ text = await response.text()
121
+ return SandboxResponse(success=False, error=text)
122
+
123
+ # Create sandbox
124
+ headers = {"Authorization": f"Bearer {self.token}"}
125
+ async with session.post(f"{self.base_url}/api/v1/sandbox/create", headers=headers) as response:
126
+ if response.status == 200:
127
+ data = await response.json()
128
+ self.sandbox_id = data["sandbox_id"]
129
+ else:
130
+ text = await response.text()
131
+ return SandboxResponse(success=False, error=text)
132
+
133
+ # Wait for sandbox to be ready
134
+ return await self.wait_for_ready()
135
+
136
+ async def wait_for_ready(self, max_attempts: int = 10, delay: float = 1.0) -> SandboxResponse:
137
+ """Wait for the sandbox to be ready.
138
+
139
+ Args:
140
+ max_attempts: Maximum number of attempts to check if sandbox is ready
141
+ delay: Delay between attempts in seconds
142
+
143
+ Returns:
144
+ SandboxResponse indicating if the sandbox is ready
145
+ """
146
+ if not self.token or not self.sandbox_id:
147
+ return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
148
+
149
+ headers = {"Authorization": f"Bearer {self.token}"}
150
+
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
+
163
+ data = await response.json()
164
+ status = data.get("status")
165
+
166
+ if status == "ready":
167
+ return SandboxResponse(success=True, data=data)
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")
182
+
183
+ async def cleanup(self) -> SandboxResponse:
184
+ """Delete the sandbox.
185
+
186
+ Returns:
187
+ SandboxResponse indicating success/failure
188
+ """
189
+ if not self.token or not self.sandbox_id:
190
+ return SandboxResponse(success=False, error="Client not properly initialized. Call setup() first")
191
+
192
+ headers = {"Authorization": f"Bearer {self.token}"}
193
+
194
+ try:
195
+ async with aiohttp.ClientSession() as session:
196
+ async with session.delete(
197
+ f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}",
198
+ headers=headers
199
+ ) as response:
200
+ if response.status != 200:
201
+ text = await response.text()
202
+ return SandboxResponse(success=False, error=text)
203
+
204
+ # Reset sandbox ID
205
+ self.sandbox_id = None
206
+ return SandboxResponse(success=True)
207
+
208
+ except Exception as e:
209
+ return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
210
+
211
+ # Backward compatibility methods
212
+
213
+ async def execute_code(self, code: str, timeout: int = 30) -> SandboxResponse:
214
+ """Execute Python code in the sandbox.
215
+
216
+ This is a backward compatibility method that delegates to the code submodule.
217
+
218
+ Args:
219
+ code: The Python code to execute
220
+ timeout: Maximum execution time in seconds
221
+
222
+ Returns:
223
+ SandboxResponse containing execution results
224
+ """
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.
229
+
230
+ This is a backward compatibility method that delegates to the code submodule.
231
+
232
+ Args:
233
+ code: The Python code to execute
234
+ timeout: Maximum execution time in seconds
235
+
236
+ Yields:
237
+ StreamEvent objects containing execution output
238
+ """
239
+ async for event in self.code.execute_stream(code, timeout):
240
+ yield event
241
+
242
+ async def execute_shell(self, command: str, args: Optional[List[str]] = None, timeout: int = 30) -> SandboxResponse:
243
+ """Execute a shell command in the sandbox.
244
+
245
+ This is a backward compatibility method that delegates to the shell submodule.
246
+
247
+ Args:
248
+ command: The shell command to execute
249
+ args: Optional list of arguments for the command
250
+ timeout: Maximum execution time in seconds
251
+
252
+ Returns:
253
+ SandboxResponse containing execution results
254
+ """
255
+ return await self.shell.execute(command, args, timeout)
256
+
257
+ async def upload_file(
258
+ self,
259
+ file_path: Union[str, Path],
260
+ destination: str = "/workspace",
261
+ chunk_size: int = 1024 * 1024,
262
+ timeout: int = 300
263
+ ) -> FileOperationResponse:
264
+ """Upload a file to the sandbox environment.
265
+
266
+ This is a backward compatibility method that delegates to the fs submodule.
267
+
268
+ Args:
269
+ file_path: Path to the file to upload
270
+ destination: Destination path in the sandbox (absolute path starting with /)
271
+ chunk_size: Size of chunks for reading large files
272
+ timeout: Maximum upload time in seconds
273
+
274
+ Returns:
275
+ FileOperationResponse containing upload results
276
+ """
277
+ return await self.fs.upload_file(file_path, destination, chunk_size, timeout)
278
+
279
+ async def download_file(
280
+ self,
281
+ remote_path: str,
282
+ local_path: Optional[Union[str, Path]] = None,
283
+ timeout: int = 300
284
+ ) -> FileOperationResponse:
285
+ """Download a file from the sandbox.
286
+
287
+ This is a backward compatibility method that delegates to the fs submodule.
288
+
289
+ Args:
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)
292
+ timeout: Maximum download time in seconds
293
+
294
+ Returns:
295
+ FileOperationResponse containing download results
296
+ """
297
+ return await self.fs.download_file(remote_path, local_path, timeout)
298
+
299
+ async def upload_bytes(
300
+ self,
301
+ content: Union[bytes, BinaryIO],
302
+ filename: str,
303
+ destination: str = "/workspace",
304
+ content_type: Optional[str] = None,
305
+ timeout: int = 300
306
+ ) -> FileOperationResponse:
307
+ """Upload bytes or a file-like object to the sandbox environment.
308
+
309
+ This is a backward compatibility method that delegates to the fs submodule.
310
+
311
+ Args:
312
+ content: Bytes or file-like object to upload
313
+ filename: Name to give the file in the sandbox
314
+ destination: Destination path in the sandbox (absolute path starting with /)
315
+ content_type: Optional MIME type (will be guessed from filename if not provided)
316
+ timeout: Maximum upload time in seconds
317
+
318
+ Returns:
319
+ FileOperationResponse containing upload results
320
+ """
321
+ # Create a temporary file with the content
322
+ import tempfile
323
+ with tempfile.NamedTemporaryFile(delete=False) as temp_file:
324
+ if isinstance(content, bytes):
325
+ temp_file.write(content)
326
+ else:
327
+ # Ensure we're at the start of the file
328
+ if hasattr(content, 'seek'):
329
+ content.seek(0)
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)
335
+
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
+ )
344
+
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]:
356
+ """
357
+ Download a file from the sandbox into memory.
358
+
359
+ Args:
360
+ remote_path: Path to the file in the sandbox.
361
+ timeout: Timeout in seconds for the operation.
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
+
369
+ try:
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
+ )
378
+ except Exception as e:
379
+ return FileOperationResponse(
380
+ success=False,
381
+ error=f"Error downloading file: {str(e)}"
382
+ )
383
+
384
+ # Additional backward compatibility methods can be added as needed
@@ -0,0 +1,45 @@
1
+ from dataclasses import dataclass
2
+ from typing import Optional, Dict
3
+
4
+ @dataclass
5
+ class SandboxResponse:
6
+ """Response from sandbox operations.
7
+
8
+ Attributes:
9
+ success: Whether the operation was successful
10
+ data: Optional response data
11
+ error: Optional error message if operation failed
12
+ """
13
+ success: bool
14
+ data: Optional[Dict] = None
15
+ error: Optional[str] = None
16
+
17
+ @dataclass
18
+ class StreamEvent:
19
+ """Event from streaming code execution.
20
+
21
+ Attributes:
22
+ type: Type of event ('stdout', 'stderr', 'info', 'error', 'completed', 'keepalive')
23
+ data: Event data
24
+ """
25
+ type: str
26
+ data: str
27
+
28
+ @dataclass
29
+ class FileOperationResponse:
30
+ """Response from file operations.
31
+
32
+ Attributes:
33
+ success: Whether the operation was successful
34
+ filename: Name of the file
35
+ size: Size of the file in bytes
36
+ path: Path where the file was saved
37
+ message: Optional status message
38
+ error: Optional error message if operation failed
39
+ """
40
+ success: bool
41
+ filename: Optional[str] = None
42
+ size: Optional[int] = None
43
+ path: Optional[str] = None
44
+ message: Optional[str] = None
45
+ error: Optional[str] = None
@@ -0,0 +1,5 @@
1
+ from .filesystem import FileSystemModule
2
+ from .shell import ShellModule
3
+ from .code import CodeModule
4
+
5
+ __all__ = ["FileSystemModule", "ShellModule", "CodeModule"]
@@ -0,0 +1,81 @@
1
+ from typing import Optional, Dict, Any
2
+ import aiohttp
3
+ from ..models import SandboxResponse
4
+
5
+ class BaseSubmodule:
6
+ """Base class for all submodules.
7
+
8
+ This class provides common functionality for all submodules, including
9
+ access to the parent client's authentication token and sandbox ID.
10
+
11
+ Attributes:
12
+ _client: Reference to the parent SandboxClient
13
+ """
14
+
15
+ def __init__(self, client):
16
+ """Initialize the submodule.
17
+
18
+ Args:
19
+ client: The parent SandboxClient instance
20
+ """
21
+ self._client = client
22
+
23
+ @property
24
+ def base_url(self) -> str:
25
+ """Get the base URL from the parent client."""
26
+ return self._client.base_url
27
+
28
+ @property
29
+ def token(self) -> Optional[str]:
30
+ """Get the authentication token from the parent client."""
31
+ return self._client.token
32
+
33
+ @property
34
+ def sandbox_id(self) -> Optional[str]:
35
+ """Get the sandbox ID from the parent client."""
36
+ return self._client.sandbox_id
37
+
38
+ async def _ensure_ready(self) -> SandboxResponse:
39
+ """Ensure the sandbox is ready for operations.
40
+
41
+ Returns:
42
+ SandboxResponse indicating if the sandbox is ready
43
+ """
44
+ if not self.token or not self.sandbox_id:
45
+ return SandboxResponse(
46
+ success=False,
47
+ error="Client not properly initialized. Call setup() first"
48
+ )
49
+
50
+ # Ensure sandbox is ready
51
+ return await self._client.wait_for_ready()
52
+
53
+ def _get_headers(self, content_type: str = "application/json") -> Dict[str, str]:
54
+ """Get the headers for API requests.
55
+
56
+ Args:
57
+ content_type: The content type for the request
58
+
59
+ Returns:
60
+ Dictionary of headers
61
+ """
62
+ return {
63
+ "Authorization": f"Bearer {self.token}",
64
+ "Content-Type": content_type
65
+ }
66
+
67
+ async def _handle_response(self, response: aiohttp.ClientResponse) -> SandboxResponse:
68
+ """Handle the API response.
69
+
70
+ Args:
71
+ response: The aiohttp response object
72
+
73
+ Returns:
74
+ SandboxResponse with the parsed response data
75
+ """
76
+ if response.status != 200:
77
+ error_text = await response.text()
78
+ return SandboxResponse(success=False, error=error_text)
79
+
80
+ result = await response.json()
81
+ return SandboxResponse(success=True, data=result)