blaxel 0.1.21rc69__py3-none-any.whl → 0.2.0__py3-none-any.whl

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/delete_process_identifier_kill.py

@@ -6,32 +6,18 @@ import httpx
 from ... import errors
 from ...client import Client
 from ...models.error_response import ErrorResponse
-from ...models.process_kill_request import ProcessKillRequest
 from ...models.success_response import SuccessResponse
 from ...types import Response
 
 
 def _get_kwargs(
     identifier: str,
-    *,
-    body: ProcessKillRequest,
 ) -> dict[str, Any]:
-    headers: dict[str, Any] = {}
-
     _kwargs: dict[str, Any] = {
         "method": "delete",
         "url": f"/process/{identifier}/kill",
     }
 
-    if type(body) == dict:
-        _body = body
-    else:
-        _body = body.to_dict()
-
-    _kwargs["json"] = _body
-    headers["Content-Type"] = "application/json"
-
-    _kwargs["headers"] = headers
     return _kwargs
 
 
@@ -71,7 +57,6 @@ def sync_detailed(
     identifier: str,
     *,
     client: Union[Client],
-    body: ProcessKillRequest,
 ) -> Response[Union[ErrorResponse, SuccessResponse]]:
     """Kill a process
 
@@ -79,7 +64,6 @@ def sync_detailed(
 
     Args:
         identifier (str):
-        body (ProcessKillRequest):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -91,7 +75,6 @@ def sync_detailed(
 
     kwargs = _get_kwargs(
         identifier=identifier,
-        body=body,
     )
 
     response = client.get_httpx_client().request(
@@ -105,7 +88,6 @@ def sync(
     identifier: str,
     *,
     client: Union[Client],
-    body: ProcessKillRequest,
 ) -> Optional[Union[ErrorResponse, SuccessResponse]]:
     """Kill a process
 
@@ -113,7 +95,6 @@ def sync(
 
     Args:
         identifier (str):
-        body (ProcessKillRequest):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -126,7 +107,6 @@ def sync(
     return sync_detailed(
         identifier=identifier,
         client=client,
-        body=body,
     ).parsed
 
 
@@ -134,7 +114,6 @@ async def asyncio_detailed(
     identifier: str,
     *,
     client: Union[Client],
-    body: ProcessKillRequest,
 ) -> Response[Union[ErrorResponse, SuccessResponse]]:
     """Kill a process
 
@@ -142,7 +121,6 @@ async def asyncio_detailed(
 
     Args:
         identifier (str):
-        body (ProcessKillRequest):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -154,7 +132,6 @@ async def asyncio_detailed(
 
     kwargs = _get_kwargs(
         identifier=identifier,
-        body=body,
     )
 
     response = await client.get_async_httpx_client().request(**kwargs)
@@ -166,7 +143,6 @@ async def asyncio(
     identifier: str,
     *,
     client: Union[Client],
-    body: ProcessKillRequest,
 ) -> Optional[Union[ErrorResponse, SuccessResponse]]:
     """Kill a process
 
@@ -174,7 +150,6 @@ async def asyncio(
 
     Args:
         identifier (str):
-        body (ProcessKillRequest):
 
     Raises:
         errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
@@ -188,6 +163,5 @@ async def asyncio(
         await asyncio_detailed(
             identifier=identifier,
             client=client,
-            body=body,
         )
     ).parsed

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/__init__.py

@@ -11,9 +11,9 @@ from .file_with_content import FileWithContent
 from .get_network_process_pid_ports_response_200 import GetNetworkProcessPidPortsResponse200
 from .port_monitor_request import PortMonitorRequest
 from .post_network_process_pid_monitor_response_200 import PostNetworkProcessPidMonitorResponse200
-from .process_kill_request import ProcessKillRequest
 from .process_logs import ProcessLogs
 from .process_request import ProcessRequest
+from .process_request_env import ProcessRequestEnv
 from .process_response import ProcessResponse
 from .process_response_status import ProcessResponseStatus
 from .subdirectory import Subdirectory
@@ -29,9 +29,9 @@ __all__ = (
     "GetNetworkProcessPidPortsResponse200",
     "PortMonitorRequest",
     "PostNetworkProcessPidMonitorResponse200",
-    "ProcessKillRequest",
     "ProcessLogs",
     "ProcessRequest",
+    "ProcessRequestEnv",
     "ProcessResponse",
     "ProcessResponseStatus",
     "Subdirectory",

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/{sandbox → core/sandbox}/client/models/process_request.py

@@ -1,10 +1,14 @@
-from typing import Any, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, TypeVar, Union, cast
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
 from ..types import UNSET, Unset
 
+if TYPE_CHECKING:
+    from ..models.process_request_env import ProcessRequestEnv
+
+
 T = TypeVar("T", bound="ProcessRequest")
 
 
@@ -13,6 +17,7 @@ class ProcessRequest:
     """
     Attributes:
         command (str):  Example: ls -la.
+        env (Union[Unset, ProcessRequestEnv]):  Example: {'{"PORT"': ' "3000"}'}.
         name (Union[Unset, str]):  Example: my-process.
         timeout (Union[Unset, int]):  Example: 30.
         wait_for_completion (Union[Unset, bool]):
@@ -21,6 +26,7 @@ class ProcessRequest:
     """
 
     command: str
+    env: Union[Unset, "ProcessRequestEnv"] = UNSET
     name: Union[Unset, str] = UNSET
     timeout: Union[Unset, int] = UNSET
     wait_for_completion: Union[Unset, bool] = UNSET
@@ -31,6 +37,12 @@ class ProcessRequest:
     def to_dict(self) -> dict[str, Any]:
         command = self.command
 
+        env: Union[Unset, dict[str, Any]] = UNSET
+        if self.env and not isinstance(self.env, Unset) and not isinstance(self.env, dict):
+            env = self.env.to_dict()
+        elif self.env and isinstance(self.env, dict):
+            env = self.env
+
         name = self.name
 
         timeout = self.timeout
@@ -50,6 +62,8 @@ class ProcessRequest:
                 "command": command,
             }
         )
+        if env is not UNSET:
+            field_dict["env"] = env
         if name is not UNSET:
             field_dict["name"] = name
         if timeout is not UNSET:
@@ -65,11 +79,20 @@ class ProcessRequest:
 
     @classmethod
     def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        from ..models.process_request_env import ProcessRequestEnv
+
         if not src_dict:
             return None
         d = src_dict.copy()
         command = d.pop("command")
 
+        _env = d.pop("env", UNSET)
+        env: Union[Unset, ProcessRequestEnv]
+        if isinstance(_env, Unset):
+            env = UNSET
+        else:
+            env = ProcessRequestEnv.from_dict(_env)
+
         name = d.pop("name", UNSET)
 
         timeout = d.pop("timeout", UNSET)
@@ -82,6 +105,7 @@ class ProcessRequest:
 
         process_request = cls(
             command=command,
+            env=env,
             name=name,
             timeout=timeout,
             wait_for_completion=wait_for_completion,

blaxel/core/sandbox/client/models/process_request_env.py

@@ -0,0 +1,49 @@
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="ProcessRequestEnv")
+
+
+@_attrs_define
+class ProcessRequestEnv:
+    """
+    Example:
+        {'{"PORT"': ' "3000"}'}
+
+    """
+
+    additional_properties: dict[str, str] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        process_request_env = cls()
+
+        process_request_env.additional_properties = d
+        return process_request_env
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> str:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: str) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties