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

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

@@ -8,17 +8,63 @@ 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
+
+try:
+    from opentelemetry.trace import Span, StatusCode
+
+    HAS_OPENTELEMETRY = True
+except ImportError:
+    HAS_OPENTELEMETRY = False
+
+    # Create dummy classes for when opentelemetry is not available
+    class Span:
+        def set_attributes(self, *args, **kwargs):
+            pass
+
+        def set_status(self, *args, **kwargs):
+            pass
+
+        def record_exception(self, *args, **kwargs):
+            pass
+
+        def end(self, *args, **kwargs):
+            pass
+
+        def add_event(self, *args, **kwargs):
+            pass
+
+        def get_span_context(self, *args, **kwargs):
+            return None
+
+        def is_recording(self, *args, **kwargs):
+            return False
+
+        def set_attribute(self, *args, **kwargs):
+            pass
+
+        def update_name(self, *args, **kwargs):
+            pass
+
+    class StatusCode:
+        ERROR = "ERROR"
+
 
 from ..common.env import env
-from ..instrumentation.span import SpanManager
 
 logger = logging.getLogger(__name__)
 
 
+class DummySpanManager:
+    """Dummy span manager for when opentelemetry is not available."""
+
+    def create_span(self, name: str, attributes: dict = None):
+        return Span()
+
+
 class BlaxelMcpServerTransport:
     """WebSocket server transport for MCP."""
+
     spans: Dict[str, Span] = {}
 
     def __init__(self, port: int = 8080):
@@ -33,7 +79,13 @@ class BlaxelMcpServerTransport:
             self.port = port
         self.clients = {}
         self.server = None
-        self.span_manager = SpanManager("blaxel-tracer")
+
+        # Initialize span manager
+        if HAS_OPENTELEMETRY:
+            # TODO: Implement proper OpenTelemetry span manager when telemetry is available
+            self.span_manager = DummySpanManager()
+        else:
+            self.span_manager = DummySpanManager()
 
     @asynccontextmanager
     async def websocket_server(self):
@@ -47,7 +99,7 @@ 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}")
@@ -61,14 +113,18 @@ class BlaxelMcpServerTransport:
                         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
+                            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)
@@ -102,13 +158,15 @@ 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)
+                        span = self.spans.get(client_id + ":" + msg_id)
                         try:
                             await websocket.send(data)
                             if span:
-                                span.set_attributes({
-                                    "mcp.message.response_sent": True,
-                                })
+                                span.set_attributes(
+                                    {
+                                        "mcp.message.response_sent": True,
+                                    }
+                                )
                         except Exception as e:
                             if span:
                                 span.set_status(StatusCode.ERROR)
@@ -140,6 +198,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