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

blaxel/{mcp → core/mcp}/client.py

@@ -7,11 +7,12 @@ import anyio
 import mcp.types as types
 from anyio.abc import TaskStatus
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from websockets.client import WebSocketClientProtocol
-from websockets.client import connect as ws_connect
+from websockets.asyncio.client import ClientConnection
+from websockets.asyncio.client import connect as ws_connect
 
 logger = logging.getLogger(__name__)
 
+
 def remove_request_params(url: str) -> str:
     return urljoin(url, urlparse(url).path)
 
@@ -44,11 +45,14 @@ async def websocket_client(
                 ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
                 logger.debug(f"Connecting to WebSocket endpoint: {remove_request_params(ws_url)}")
 
-                async with ws_connect(ws_url, extra_headers=headers, open_timeout=timeout) as websocket:
+                # Use different parameters based on websockets version
+                connection_kwargs = {"additional_headers": headers, "open_timeout": timeout}
+
+                async with ws_connect(ws_url, **connection_kwargs) as websocket:
                     logger.debug("WebSocket connection established")
 
                     async def ws_reader(
-                        websocket: WebSocketClientProtocol,
+                        websocket: ClientConnection,
                         task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
                     ):
                         try:
@@ -56,7 +60,9 @@ async def websocket_client(
                             async for message in websocket:
                                 logger.debug(f"Received WebSocket message: {message}")
                                 try:
-                                    parsed_message = types.JSONRPCMessage.model_validate_json(message)
+                                    parsed_message = types.JSONRPCMessage.model_validate_json(
+                                        message
+                                    )
                                     logger.debug(f"Received server message: {parsed_message}")
                                     await read_stream_writer.send(parsed_message)
                                 except Exception as exc:
@@ -68,7 +74,7 @@ async def websocket_client(
                         finally:
                             await read_stream_writer.aclose()
 
-                    async def ws_writer(websocket: WebSocketClientProtocol):
+                    async def ws_writer(websocket: ClientConnection):
                         try:
                             async with write_stream_reader:
                                 async for message in write_stream_reader:
@@ -107,4 +113,4 @@ async def websocket_client(
     finally:
         # Ensure streams are closed even if task group setup fails or GeneratorExit is caught
         await read_stream_writer.aclose()
-        await write_stream.aclose()
+        await write_stream.aclose()

blaxel/{mcp → core/mcp}/server.py

@@ -2,24 +2,21 @@ import logging
 import traceback
 import uuid
 from contextlib import asynccontextmanager
-from typing import Dict, Literal
+from typing import Literal
 
 import anyio
 import mcp.types as types
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp.server.fastmcp import FastMCP as FastMCPBase
-from opentelemetry.trace import Span, StatusCode
-from websockets.server import WebSocketServerProtocol, serve
+from websockets.asyncio.server import ServerConnection, serve
 
 from ..common.env import env
-from ..instrumentation.span import SpanManager
 
 logger = logging.getLogger(__name__)
 
 
 class BlaxelMcpServerTransport:
     """WebSocket server transport for MCP."""
-    spans: Dict[str, Span] = {}
 
     def __init__(self, port: int = 8080):
         """Initialize the WebSocket server transport.
@@ -33,7 +30,6 @@ class BlaxelMcpServerTransport:
             self.port = port
         self.clients = {}
         self.server = None
-        self.span_manager = SpanManager("blaxel-tracer")
 
     @asynccontextmanager
     async def websocket_server(self):
@@ -47,33 +43,21 @@ class BlaxelMcpServerTransport:
         read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
         write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
-        async def handler(websocket: WebSocketServerProtocol):
+        async def handler(websocket: ServerConnection):
             client_id = str(uuid.uuid4())
             self.clients[client_id] = websocket
             logger.info(f"Client connected: {client_id}")
 
             try:
                 async for message in websocket:
-                    span = self.span_manager.create_span("message", {"mcp.client.id": client_id})
                     try:
                         msg = types.JSONRPCMessage.model_validate_json(message)
                         # Modify message ID to include client ID
                         if hasattr(msg, "id") and msg.id is not None:
                             original_id = msg.id
                             msg.id = f"{client_id}:{original_id}"
-                            span.set_attributes({
-                                "mcp.message.parsed": True,
-                                "mcp.method": getattr(msg, "method", None),
-                                "mcp.messageId": getattr(msg, "id", None),
-                                "mcp.toolName": getattr(getattr(msg, "params", None), "name", None),
-                                "span.type": "mcp.message",
-                            })
-                            self.spans[client_id+":"+msg.id] =  span
                         await read_stream_writer.send(msg)
                     except Exception as exc:
-                        span.set_status(StatusCode.ERROR)
-                        span.record_exception(exc)
-                        span.end()
                         logger.error(f"Failed to parse message: {exc}\n{traceback.format_exc()}")
                         await read_stream_writer.send(exc)
             except Exception as e:
@@ -102,23 +86,12 @@ class BlaxelMcpServerTransport:
                     if client_id and client_id in self.clients:
                         # Send to specific client
                         websocket = self.clients[client_id]
-                        span = self.spans.get(client_id+":"+msg_id)
                         try:
                             await websocket.send(data)
-                            if span:
-                                span.set_attributes({
-                                    "mcp.message.response_sent": True,
-                                })
                         except Exception as e:
-                            if span:
-                                span.set_status(StatusCode.ERROR)
-                                span.record_exception(e)
                             logger.error(f"Failed to send message to client {client_id}: {e}")
                             if client_id in self.clients:
                                 del self.clients[client_id]
-                        finally:
-                            if span:
-                                span.end()
                     else:
                         # Broadcast to all clients
                         dead_clients = []
@@ -140,6 +113,7 @@ class BlaxelMcpServerTransport:
                 tg.start_soon(message_sender)
                 yield read_stream, write_stream
 
+
 class FastMCP(FastMCPBase):
     def run(self, transport: Literal["stdio", "sse", "ws"] = "stdio") -> None:
         """Run the FastMCP server. Note this is a synchronous function.

blaxel/core/models/__init__.py

@@ -0,0 +1,52 @@
+
+from ..cache import find_from_cache
+from ..client import client
+from ..client.api.models import get_model
+from ..client.models import Model
+from ..common import settings
+
+
+class BLModel:
+    models = {}
+
+    def __init__(self, model_name, **kwargs):
+        self.model_name = model_name
+        self.kwargs = kwargs
+
+    async def get_parameters(self) -> tuple[str, str, str]:
+        if self.model_name in self.models:
+            # We get the headers in case we need to refresh the token
+            settings.auth.get_headers()
+            model = self.models[self.model_name]
+            return model["url"], model["type"], model["model"]
+        url = f"{settings.run_url}/{settings.auth.workspace_name}/models/{self.model_name}"
+        model_data = await self._get_model_metadata()
+        if not model_data:
+            raise Exception(f"Model {self.model_name} not found")
+        runtime = (model_data.spec and model_data.spec.runtime)
+        if not runtime:
+            raise Exception(f"Model {self.model_name} has no runtime")
+
+        type = runtime.type_ or 'openai'
+        model = runtime.model
+        self.models[self.model_name] = {
+            "url": url,
+            "type": type,
+            "model": model
+        }
+        return url, type, model
+
+    async def _get_model_metadata(self) -> Model | None:
+        cache_data = await find_from_cache('Model', self.model_name)
+        if cache_data:
+            return Model.from_dict(cache_data)
+
+        try:
+            return await get_model.asyncio(client=client, model_name=self.model_name)
+        except Exception:
+            return None
+
+def bl_model(model_name, **kwargs):
+    return BLModel(model_name, **kwargs)
+
+__all__ = ["bl_model"]

blaxel/core/sandbox/__init__.py

@@ -0,0 +1,29 @@
+from .sandbox import (
+    Sandbox,
+    SandboxFileSystem,
+    SandboxInstance,
+    SandboxPreviews,
+    SandboxProcess,
+)
+from .types import (
+    CopyResponse,
+    SandboxConfiguration,
+    SandboxFilesystemFile,
+    SessionCreateOptions,
+    SessionWithToken,
+    WatchEvent,
+)
+
+__all__ = [
+    "SandboxInstance",
+    "SessionCreateOptions",
+    "SessionWithToken",
+    "SandboxConfiguration",
+    "WatchEvent",
+    "SandboxFilesystemFile",
+    "CopyResponse",
+    "Sandbox",
+    "SandboxFileSystem",
+    "SandboxPreviews",
+    "SandboxProcess",
+]

blaxel/core/sandbox/action.py

@@ -0,0 +1,79 @@
+from typing import Any, Optional
+
+import httpx
+
+from ..common.internal import get_forced_url, get_global_unique_hash
+from ..common.settings import settings
+from .types import SandboxConfiguration
+
+
+class ResponseError(Exception):
+    def __init__(self, response: httpx.Response, data: Any = None, error: Any = None):
+        data_error = {}
+        if isinstance(data, dict) and "error" in data:
+            data_error = data
+        if isinstance(error, dict) and "error" in error:
+            data_error["error"] = error["error"]
+        if response.status_code:
+            data_error["status"] = response.status_code
+        if response.reason_phrase:
+            data_error["statusText"] = response.reason_phrase
+
+        super().__init__(str(data_error))
+        self.response = response
+        self.data = data
+        self.error = error
+
+
+class SandboxAction:
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        self.sandbox_config = sandbox_config
+
+    @property
+    def name(self) -> str:
+        return self.sandbox_config.metadata.name if self.sandbox_config.metadata else ""
+
+    @property
+    def external_url(self) -> str:
+        return f"{settings.run_url}/{settings.workspace}/sandboxes/{self.name}"
+
+    @property
+    def internal_url(self) -> str:
+        hash_value = get_global_unique_hash(settings.workspace, "sandbox", self.name)
+        return f"{settings.run_internal_protocol}://bl-{settings.env}-{hash_value}.{settings.run_internal_hostname}"
+
+    @property
+    def forced_url(self) -> Optional[str]:
+        if self.sandbox_config.force_url:
+            return self.sandbox_config.force_url
+        return get_forced_url("sandbox", self.name)
+
+    @property
+    def url(self) -> str:
+        if self.forced_url:
+            return self.forced_url
+        # Uncomment when mk3 is fully available
+        # if settings.run_internal_hostname:
+        #     return self.internal_url
+        return self.external_url
+
+    @property
+    def fallback_url(self) -> Optional[str]:
+        if self.external_url != self.url:
+            return self.external_url
+        return None
+
+    def get_client(self) -> httpx.AsyncClient:
+        if self.sandbox_config.force_url:
+            return httpx.AsyncClient(
+                base_url=self.sandbox_config.force_url, headers=self.sandbox_config.headers
+            )
+        # Create a new client instance each time to avoid "Cannot open a client instance more than once" error
+        return httpx.AsyncClient(
+            base_url=self.url,
+            headers={**settings.headers, **self.sandbox_config.headers},
+        )
+
+    def handle_response_error(self, response: httpx.Response, data: Any, error: Any):
+        if not response.is_success or not data:
+            raise ResponseError(response, data, error)

blaxel/{sandbox → core/sandbox}/client/api/filesystem/put_filesystem_path.py

@@ -23,7 +23,7 @@ def _get_kwargs(
         "url": f"/filesystem/{path}",
     }
 
-    if type(body) == dict:
+    if type(body) is dict:
         _body = body
     else:
         _body = body.to_dict()

blaxel/{sandbox → core/sandbox}/client/api/network/post_network_process_pid_monitor.py

@@ -25,7 +25,7 @@ def _get_kwargs(
         "url": f"/network/process/{pid}/monitor",
     }
 
-    if type(body) == dict:
+    if type(body) is dict:
         _body = body
     else:
         _body = body.to_dict()

blaxel/{sandbox → core/sandbox}/client/api/process/post_process.py

@@ -22,7 +22,7 @@ def _get_kwargs(
         "url": "/process",
     }
 
-    if type(body) == dict:
+    if type(body) is dict:
         _body = body
     else:
         _body = body.to_dict()

blaxel/{sandbox → core/sandbox}/client/models/directory.py

@@ -34,7 +34,7 @@ class Directory:
         if not isinstance(self.files, Unset):
             files = []
             for files_item_data in self.files:
-                if type(files_item_data) == dict:
+                if type(files_item_data) is dict:
                     files_item = files_item_data
                 else:
                     files_item = files_item_data.to_dict()
@@ -48,7 +48,7 @@ class Directory:
         if not isinstance(self.subdirectories, Unset):
             subdirectories = []
             for subdirectories_item_data in self.subdirectories:
-                if type(subdirectories_item_data) == dict:
+                if type(subdirectories_item_data) is dict:
                     subdirectories_item = subdirectories_item_data
                 else:
                     subdirectories_item = subdirectories_item_data.to_dict()

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