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

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

@@ -7,11 +7,13 @@ 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 mcp.shared.message import SessionMessage
+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)
 
@@ -27,11 +29,11 @@ async def websocket_client(
 
     The `timeout` parameter controls connection timeout.
     """
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: MemoryObjectSendStream[SessionMessage]
+    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -44,11 +46,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,9 +61,13 @@ 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)
+                                    await read_stream_writer.send(
+                                        SessionMessage(message=parsed_message)
+                                    )
                                 except Exception as exc:
                                     logger.error(f"Error parsing server message: {exc}")
                                     await read_stream_writer.send(exc)
@@ -68,13 +77,15 @@ 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:
-                                    logger.debug(f"Sending client message: {message}")
+                                async for session_message in write_stream_reader:
+                                    logger.debug(
+                                        f"Sending client message: {session_message.message}"
+                                    )
                                     await websocket.send(
-                                        message.model_dump_json(
+                                        session_message.message.model_dump_json(
                                             by_alias=True,
                                             exclude_none=True,
                                         )
@@ -107,4 +118,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,22 @@ 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 mcp.shared.message import SessionMessage
+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,47 +31,34 @@ class BlaxelMcpServerTransport:
             self.port = port
         self.clients = {}
         self.server = None
-        self.span_manager = SpanManager("blaxel-tracer")
 
     @asynccontextmanager
     async def websocket_server(self):
         """Create and run a WebSocket server for MCP communication."""
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-        read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+        read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-        write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+        write_stream: MemoryObjectSendStream[SessionMessage]
+        write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
         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)
+                        await read_stream_writer.send(SessionMessage(message=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:
@@ -85,11 +70,11 @@ class BlaxelMcpServerTransport:
 
         async def message_sender():
             async with write_stream_reader:
-                async for message in write_stream_reader:
+                async for session_message in write_stream_reader:
                     # Extract client ID from message ID
                     client_id = None
                     msg_id = None
-
+                    message = session_message.message
                     if hasattr(message, "id") and message.id is not None:
                         parts = str(message.id).split(":", 1)
                         if len(parts) == 2:
@@ -102,23 +87,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 +114,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()