blaxel 0.1.22rc70__py3-none-any.whl → 0.2.0__py3-none-any.whl

blaxel/core/sandbox/filesystem.py

@@ -0,0 +1,280 @@
+import asyncio
+import io
+import json
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import httpx
+
+from ..common.settings import settings
+from .action import SandboxAction
+from .client.models import Directory, FileRequest, SuccessResponse
+from .types import CopyResponse, SandboxConfiguration, SandboxFilesystemFile, WatchEvent
+
+
+class SandboxFileSystem(SandboxAction):
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    async def mkdir(self, path: str, permissions: str = "0755") -> SuccessResponse:
+        path = self.format_path(path)
+        body = FileRequest(is_directory=True, permissions=permissions)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.put(f"/filesystem/{path}", json=body.to_dict())
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def write(self, path: str, content: str) -> SuccessResponse:
+        path = self.format_path(path)
+        body = FileRequest(content=content)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.put(f"/filesystem/{path}", json=body.to_dict())
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def write_binary(self, path: str, content: Union[bytes, bytearray]) -> SuccessResponse:
+        """Write binary content to a file."""
+        path = self.format_path(path)
+
+        # Convert bytearray to bytes if necessary
+        if isinstance(content, bytearray):
+            content = bytes(content)
+
+        # Wrap binary content in BytesIO to provide file-like interface
+        binary_file = io.BytesIO(content)
+
+        # Prepare multipart form data
+        files = {
+            "file": ("binary-file.bin", binary_file, "application/octet-stream"),
+        }
+        data = {"permissions": "0644", "path": path}
+
+        # Use the fixed get_client method
+        url = f"{self.url}/filesystem/{path}"
+        headers = {**settings.headers, **self.sandbox_config.headers}
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.put(url, files=files, data=data, headers=headers)
+
+            if not response.is_success:
+                raise Exception(f"Failed to write binary: {response.status_code} {response.text}")
+
+            return SuccessResponse.from_dict(response.json())
+
+    async def write_tree(
+        self,
+        files: List[Union[SandboxFilesystemFile, Dict[str, Any]]],
+        destination_path: Optional[str] = None,
+    ) -> Directory:
+        """Write multiple files in a tree structure."""
+        files_dict = {}
+        for file in files:
+            if isinstance(file, dict):
+                file = SandboxFilesystemFile.from_dict(file)
+            files_dict[file.path] = file.content
+
+        path = destination_path or ""
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.put(
+                f"/filesystem/tree/{path}",
+                json={"files": files_dict},
+                headers={"Content-Type": "application/json"},
+            )
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return Directory.from_dict(response.json())
+
+    async def read(self, path: str) -> str:
+        path = self.format_path(path)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/filesystem/{path}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+
+            data = response.json()
+            if "content" in data:
+                return data["content"]
+            raise Exception("Unsupported file type")
+
+    async def rm(self, path: str, recursive: bool = False) -> SuccessResponse:
+        path = self.format_path(path)
+
+        async with self.get_client() as client_instance:
+            params = {"recursive": "true"} if recursive else {}
+            response = await client_instance.delete(f"/filesystem/{path}", params=params)
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def ls(self, path: str) -> Directory:
+        path = self.format_path(path)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/filesystem/{path}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+
+            data = response.json()
+            if not ("files" in data or "subdirectories" in data):
+                raise Exception('{"error": "Directory not found"}')
+            return Directory.from_dict(data)
+
+    async def cp(self, source: str, destination: str) -> CopyResponse:
+        source = self.format_path(source)
+        destination = self.format_path(destination)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/filesystem/{source}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+
+            data = response.json()
+            if "files" in data or "subdirectories" in data:
+                # Create destination directory
+                await self.mkdir(destination)
+
+                # Process subdirectories in batches of 5
+                subdirectories = data.get("subdirectories", [])
+                for i in range(0, len(subdirectories), 5):
+                    batch = subdirectories[i : i + 5]
+                    await asyncio.gather(
+                        *[
+                            self.cp(
+                                subdir.get("path", f"{source}/{subdir.get('path', '')}"),
+                                f"{destination}/{subdir.get('path', '')}",
+                            )
+                            for subdir in batch
+                        ]
+                    )
+
+                # Process files in batches of 10
+                files = data.get("files", [])
+                for i in range(0, len(files), 10):
+                    batch = files[i : i + 10]
+                    tasks = []
+                    for file in batch:
+                        source_path = file.get("path", f"{source}/{file.get('path', '')}")
+                        dest_path = f"{destination}/{file.get('path', '')}"
+                        tasks.append(self._copy_file(source_path, dest_path))
+                    await asyncio.gather(*tasks)
+
+                return CopyResponse(
+                    message="Directory copied successfully", source=source, destination=destination
+                )
+            elif "content" in data:
+                await self.write(destination, data["content"])
+                return CopyResponse(
+                    message="File copied successfully", source=source, destination=destination
+                )
+
+        raise Exception("Unsupported file type")
+
+    async def _copy_file(self, source_path: str, dest_path: str):
+        """Helper method to copy a single file."""
+        content = await self.read(source_path)
+        await self.write(dest_path, content)
+
+    def watch(
+        self,
+        path: str,
+        callback: Callable[[WatchEvent], None],
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Callable]:
+        """Watch for file system changes."""
+        path = self.format_path(path)
+        closed = False
+
+        if options is None:
+            options = {}
+
+        async def start_watching():
+            nonlocal closed
+
+            params = {}
+            if options.get("ignore"):
+                params["ignore"] = ",".join(options["ignore"])
+
+            url = f"{self.url}/filesystem/{path}/watch"
+            headers = {**settings.headers, **self.sandbox_config.headers}
+
+            async with httpx.AsyncClient() as client_instance:
+                async with client_instance.stream(
+                    "GET", url, params=params, headers=headers
+                ) as response:
+                    if not response.is_success:
+                        raise Exception(f"Failed to start watching: {response.status_code}")
+
+                    buffer = ""
+                    async for chunk in response.aiter_text():
+                        if closed:
+                            break
+
+                        buffer += chunk
+                        lines = buffer.split("\n")
+                        buffer = lines.pop()  # Keep incomplete line in buffer
+
+                        for line in lines:
+                            line = line.strip()
+                            if not line:
+                                continue
+
+                            try:
+                                file_event_data = json.loads(line)
+                                file_event = WatchEvent(
+                                    op=file_event_data.get("op", ""),
+                                    path=file_event_data.get("path", ""),
+                                    name=file_event_data.get("name", ""),
+                                    content=file_event_data.get("content"),
+                                )
+
+                                if options.get("with_content") and file_event.op in [
+                                    "CREATE",
+                                    "WRITE",
+                                ]:
+                                    try:
+                                        file_path = file_event.path
+                                        if file_path.endswith("/"):
+                                            file_path = file_path + file_event.name
+                                        else:
+                                            file_path = file_path + "/" + file_event.name
+
+                                        content = await self.read(file_path)
+                                        file_event.content = content
+                                    except:
+                                        file_event.content = None
+
+                                await asyncio.create_task(asyncio.coroutine(callback)(file_event))
+                            except json.JSONDecodeError:
+                                continue
+                            except Exception as e:
+                                if options.get("on_error"):
+                                    options["on_error"](e)
+
+        # Start watching in the background
+        task = asyncio.create_task(start_watching())
+
+        def close():
+            nonlocal closed
+            closed = True
+            task.cancel()
+
+        return {"close": close}
+
+    def format_path(self, path: str) -> str:
+        if path == "/":
+            return path
+        if path.startswith("/"):
+            path = path[1:]
+        return path

blaxel/core/sandbox/network.py

@@ -0,0 +1,10 @@
+from .action import SandboxAction
+from .types import SandboxConfiguration
+
+
+class SandboxNetwork(SandboxAction):
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    # Network functionality can be expanded here in the future
+    # Currently this is a placeholder matching the TypeScript implementation

blaxel/{sandbox → core/sandbox}/preview.py

@@ -1,25 +1,38 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from ..client.api.compute.create_sandbox_preview import asyncio as create_sandbox_preview
+from ..client.api.compute.create_sandbox_preview import (
+    asyncio as create_sandbox_preview,
+)
 from ..client.api.compute.create_sandbox_preview_token import (
     asyncio as create_sandbox_preview_token,
 )
-from ..client.api.compute.delete_sandbox_preview import asyncio as delete_sandbox_preview
+from ..client.api.compute.delete_sandbox_preview import (
+    asyncio as delete_sandbox_preview,
+)
 from ..client.api.compute.delete_sandbox_preview_token import (
     asyncio as delete_sandbox_preview_token,
 )
 from ..client.api.compute.get_sandbox_preview import asyncio as get_sandbox_preview
-from ..client.api.compute.list_sandbox_preview_tokens import asyncio as list_sandbox_preview_tokens
+from ..client.api.compute.list_sandbox_preview_tokens import (
+    asyncio as list_sandbox_preview_tokens,
+)
 from ..client.api.compute.list_sandbox_previews import asyncio as list_sandbox_previews
 from ..client.client import client
-from ..client.models import Preview, PreviewSpec, PreviewToken, PreviewTokenSpec, Sandbox
+from ..client.models import (
+    Preview,
+    PreviewSpec,
+    PreviewToken,
+    PreviewTokenSpec,
+    Sandbox,
+)
 
 
 @dataclass
 class SandboxPreviewToken:
     """Represents a preview token with its value and expiration."""
+
     preview_token: PreviewToken
 
     @property
@@ -30,23 +43,22 @@ class SandboxPreviewToken:
     def expires_at(self) -> datetime:
         return self.preview_token.spec.expires_at if self.preview_token.spec else datetime.now()
 
+
 class SandboxPreviewTokens:
     """Manages preview tokens for a sandbox preview."""
-    def __init__(self, preview: Preview):
+
+    def __init__(self, preview: Preview, sandbox_name: str):
         self.preview = preview
+        self.sandbox_name = sandbox_name
 
     @property
     def preview_name(self) -> str:
         return self.preview.metadata.name if self.preview.metadata else ""
 
-    @property
-    def resource_name(self) -> str:
-        return self.preview.metadata.resource_name if self.preview.metadata else ""
-
     async def create(self, expires_at: datetime) -> SandboxPreviewToken:
         """Create a new preview token."""
         response: PreviewToken = await create_sandbox_preview_token(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             body=PreviewToken(
                 spec=PreviewTokenSpec(
@@ -60,7 +72,7 @@ class SandboxPreviewTokens:
     async def list(self) -> List[SandboxPreviewToken]:
         """List all preview tokens."""
         response: List[PreviewToken] = await list_sandbox_preview_tokens(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             client=client,
         )
@@ -69,18 +81,21 @@ class SandboxPreviewTokens:
     async def delete(self, token_name: str) -> dict:
         """Delete a preview token."""
         response: PreviewToken = await delete_sandbox_preview_token(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             token_name,
             client=client,
         )
         return response
 
+
 class SandboxPreview:
     """Represents a sandbox preview with its metadata and tokens."""
-    def __init__(self, preview: Preview):
+
+    def __init__(self, preview: Preview, sandbox_name: str = ""):
         self.preview = preview
-        self.tokens = SandboxPreviewTokens(self)
+        self.sandbox_name = sandbox_name
+        self.tokens = SandboxPreviewTokens(preview, sandbox_name)
 
     @property
     def name(self) -> str:
@@ -94,8 +109,10 @@ class SandboxPreview:
     def spec(self) -> Optional[PreviewSpec]:
         return self.preview.spec
 
+
 class SandboxPreviews:
     """Manages sandbox previews."""
+
     def __init__(self, sandbox: Sandbox):
         self.sandbox = sandbox
 
@@ -111,8 +128,11 @@ class SandboxPreviews:
         )
         return [SandboxPreview(preview) for preview in response]
 
-    async def create(self, preview: Preview) -> SandboxPreview:
+    async def create(self, preview: Union[Preview, Dict[str, Any]]) -> SandboxPreview:
         """Create a new preview."""
+        if isinstance(preview, dict):
+            preview = Preview.from_dict(preview)
+
         response: Preview = await create_sandbox_preview(
             self.sandbox_name,
             body=preview,
@@ -138,5 +158,13 @@ class SandboxPreviews:
         )
         return response
 
+
 def to_utc_z(dt: datetime) -> str:
-    return dt.isoformat(timespec="milliseconds").replace("+00:00", "Z")
+    """Convert datetime to UTC Z format string."""
+    # Simple approach: format as ISO string and add Z
+    iso_string = dt.isoformat()
+    if iso_string.endswith("+00:00"):
+        return iso_string.replace("+00:00", "Z")
+    elif "T" in iso_string and not iso_string.endswith("Z"):
+        return iso_string + "Z"
+    return iso_string

blaxel/core/sandbox/process.py

@@ -0,0 +1,159 @@
+import asyncio
+from typing import Any, Callable, Dict, Literal, Optional, Union
+
+import httpx
+
+from ..common.settings import settings
+from .action import SandboxAction
+from .client.models import ProcessResponse, SuccessResponse
+from .client.models.process_request import ProcessRequest
+from .types import SandboxConfiguration
+
+
+class SandboxProcess(SandboxAction):
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    def stream_logs(
+        self, identifier: str, options: Optional[Dict[str, Callable[[str], None]]] = None
+    ) -> Dict[str, Callable[[], None]]:
+        """Stream logs from a process with callbacks for different output types."""
+        if options is None:
+            options = {}
+
+        closed = False
+
+        async def start_streaming():
+            nonlocal closed
+
+            url = f"{self.url}/process/{identifier}/logs/stream"
+            headers = {**settings.headers, **self.sandbox_config.headers}
+
+            try:
+                async with httpx.AsyncClient() as client_instance:
+                    async with client_instance.stream("GET", url, headers=headers) as response:
+                        if response.status_code != 200:
+                            raise Exception(f"Failed to stream logs: {await response.aread()}")
+
+                        buffer = ""
+                        async for chunk in response.aiter_text():
+                            if closed:
+                                break
+
+                            buffer += chunk
+                            lines = buffer.split("\n")
+                            buffer = lines.pop()  # Keep incomplete line in buffer
+
+                            for line in lines:
+                                if line.startswith("stdout:"):
+                                    content = line[7:]  # Remove 'stdout:' prefix
+                                    if options.get("on_stdout"):
+                                        options["on_stdout"](content)
+                                    if options.get("on_log"):
+                                        options["on_log"](content)
+                                elif line.startswith("stderr:"):
+                                    content = line[7:]  # Remove 'stderr:' prefix
+                                    if options.get("on_stderr"):
+                                        options["on_stderr"](content)
+                                    if options.get("on_log"):
+                                        options["on_log"](content)
+                                else:
+                                    if options.get("on_log"):
+                                        options["on_log"](line)
+            except Exception as e:
+                # Suppress AbortError when closing
+                if not (hasattr(e, "name") and e.name == "AbortError"):
+                    raise e
+
+        # Start streaming in the background
+        task = asyncio.create_task(start_streaming())
+
+        def close():
+            nonlocal closed
+            closed = True
+            task.cancel()
+
+        return {"close": close}
+
+    async def exec(self, process: Union[ProcessRequest, Dict[str, Any]]) -> ProcessResponse:
+        if isinstance(process, dict):
+            process = ProcessRequest.from_dict(process)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.post("/process", json=process.to_dict())
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return ProcessResponse.from_dict(response.json())
+
+    async def wait(
+        self, identifier: str, max_wait: int = 60000, interval: int = 1000
+    ) -> ProcessResponse:
+        """Wait for a process to complete."""
+        start_time = asyncio.get_event_loop().time() * 1000  # Convert to milliseconds
+        status = "running"
+        data = await self.get(identifier)
+
+        while status == "running":
+            await asyncio.sleep(interval / 1000)  # Convert to seconds
+            try:
+                data = await self.get(identifier)
+                status = data.status or "running"
+            except:
+                break
+
+            if (asyncio.get_event_loop().time() * 1000) - start_time > max_wait:
+                raise Exception("Process did not finish in time")
+
+        return data
+
+    async def get(self, identifier: str) -> ProcessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/process/{identifier}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return ProcessResponse.from_dict(response.json())
+
+    async def list(self) -> list[ProcessResponse]:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get("/process")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return [ProcessResponse.from_dict(item) for item in response.json()]
+
+    async def stop(self, identifier: str) -> SuccessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.delete(f"/process/{identifier}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def kill(self, identifier: str) -> SuccessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.delete(f"/process/{identifier}/kill")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def logs(
+        self, identifier: str, log_type: Literal["stdout", "stderr", "all"] = "all"
+    ) -> str:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/process/{identifier}/logs")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+
+            data = response.json()
+            if log_type == "all":
+                return data.get("logs", "")
+            elif log_type == "stdout":
+                return data.get("stdout", "")
+            elif log_type == "stderr":
+                return data.get("stderr", "")
+
+            raise Exception("Unsupported log type")