blaxel 0.1.9rc36__py3-none-any.whl → 0.1.10__py3-none-any.whl

blaxel/sandbox/client/models/subdirectory.py

@@ -0,0 +1,60 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="Subdirectory")
+
+
+@_attrs_define
+class Subdirectory:
+    """
+    Attributes:
+        path (Union[Unset, str]):
+    """
+
+    path: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        path = self.path
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if path is not UNSET:
+            field_dict["path"] = path
+
+        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()
+        path = d.pop("path", UNSET)
+
+        subdirectory = cls(
+            path=path,
+        )
+
+        subdirectory.additional_properties = d
+        return subdirectory
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> 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

blaxel/sandbox/client/models/success_response.py

@@ -0,0 +1,69 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="SuccessResponse")
+
+
+@_attrs_define
+class SuccessResponse:
+    """
+    Attributes:
+        message (Union[Unset, str]):  Example: File created successfully.
+        path (Union[Unset, str]):  Example: /path/to/file.
+    """
+
+    message: Union[Unset, str] = UNSET
+    path: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        message = self.message
+
+        path = self.path
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if message is not UNSET:
+            field_dict["message"] = message
+        if path is not UNSET:
+            field_dict["path"] = path
+
+        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()
+        message = d.pop("message", UNSET)
+
+        path = d.pop("path", UNSET)
+
+        success_response = cls(
+            message=message,
+            path=path,
+        )
+
+        success_response.additional_properties = d
+        return success_response
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> 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

blaxel/sandbox/client/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

blaxel/sandbox/client/types.py

@@ -0,0 +1,46 @@
+"""Contains some shared types for properties"""
+
+from collections.abc import MutableMapping
+from http import HTTPStatus
+from typing import BinaryIO, Generic, Literal, Optional, TypeVar
+
+from attrs import define
+
+
+class Unset:
+    def __bool__(self) -> Literal[False]:
+        return False
+
+
+UNSET: Unset = Unset()
+
+FileJsonType = tuple[Optional[str], BinaryIO, Optional[str]]
+
+
+@define
+class File:
+    """Contains information for file uploads"""
+
+    payload: BinaryIO
+    file_name: Optional[str] = None
+    mime_type: Optional[str] = None
+
+    def to_tuple(self) -> FileJsonType:
+        """Return a tuple representation that httpx will accept for multipart/form-data"""
+        return self.file_name, self.payload, self.mime_type
+
+
+T = TypeVar("T")
+
+
+@define
+class Response(Generic[T]):
+    """A response from an endpoint"""
+
+    status_code: HTTPStatus
+    content: bytes
+    headers: MutableMapping[str, str]
+    parsed: Optional[T]
+
+
+__all__ = ["UNSET", "File", "FileJsonType", "Response", "Unset"]

blaxel/sandbox/filesystem.py

@@ -0,0 +1,104 @@
+import asyncio
+from typing import Dict
+
+from ..common.settings import settings
+from .base import SandboxHandleBase
+from .client.api.filesystem.delete_filesystem_path import (
+    asyncio_detailed as delete_filesystem_by_path,
+)
+from .client.api.filesystem.get_filesystem_path import asyncio_detailed as get_filesystem_by_path
+from .client.api.filesystem.put_filesystem_path import asyncio_detailed as put_filesystem_by_path
+from .client.client import client
+from .client.models import Directory, FileRequest, SuccessResponse
+
+
+class SandboxFileSystem(SandboxHandleBase):
+    def __init__(self, sandbox):
+        super().__init__(sandbox)
+        self.client = client.with_base_url(self.url).with_headers(settings.headers)
+
+    async def mkdir(self, path: str, permissions: str = "0755") -> SuccessResponse:
+        path = self.format_path(path)
+        body = FileRequest(is_directory=True, permissions=permissions)
+        response = await put_filesystem_by_path(path=path, client=self.client, body=body)
+        self.handle_response(response)
+        return response.parsed
+
+    async def write(self, path: str, content: str) -> SuccessResponse:
+        path = self.format_path(path)
+        body = FileRequest(content=content)
+        response = await put_filesystem_by_path(path=path, client=self.client, body=body)
+        self.handle_response(response)
+        return response.parsed
+
+    async def read(self, path: str) -> str:
+        path = self.format_path(path)
+        response = await get_filesystem_by_path(path=path, client=self.client)
+        self.handle_response(response)
+        if "content" not in response.parsed.additional_properties:
+            raise Exception('{"error": "File not found"}')
+        return response.parsed.additional_properties["content"]
+
+    async def rm(self, path: str, recursive: bool = False) -> SuccessResponse:
+        path = self.format_path(path)
+        response = await delete_filesystem_by_path(path=path, client=self.client, recursive=recursive)
+        self.handle_response(response)
+        return response.parsed
+
+    async def ls(self, path: str) -> Directory:
+        path = self.format_path(path)
+        response = await get_filesystem_by_path(path=path, client=self.client)
+        self.handle_response(response)
+        if not hasattr(response.parsed, "files") and not hasattr(response.parsed, "subdirectories"):
+            raise Exception('{"error": "Directory not found"}')
+        return response.parsed
+
+    async def cp(self, source: str, destination: str) -> Dict[str, str]:
+        source = self.format_path(source)
+        destination = self.format_path(destination)
+        response = await get_filesystem_by_path(path=source, client=self.client)
+        self.handle_response(response)
+        data = response.parsed
+        if "content" in data.additional_properties:
+            await self.write(destination, data.additional_properties["content"])
+            return {
+                "message": "File copied successfully",
+                "source": source,
+                "destination": destination,
+            }
+        elif hasattr(data, "subdirectories") or hasattr(data, "files"):
+            # Create destination directory
+            await self.mkdir(destination)
+            # Process subdirectories in batches of 5
+            subdirectories = getattr(data, "subdirectories", []) or []
+            for i in range(0, len(subdirectories), 5):
+                batch = subdirectories[i:i+5]
+                await asyncio.gather(*[
+                    self.cp(
+                        getattr(subdir, "path", f"{source}/{getattr(subdir, 'path', '')}"),
+                        f"{destination}/{getattr(subdir, 'path', '')}"
+                    ) for subdir in batch
+                ])
+            # Process files in batches of 10
+            files = getattr(data, "files", []) or []
+            for i in range(0, len(files), 10):
+                batch = files[i:i+10]
+                await asyncio.gather(*[
+                    self.write(
+                        f"{destination}/{getattr(file, 'path', '')}",
+                        await self.read(getattr(file, "path", f"{source}/{getattr(file, 'path', '')}"))
+                    ) for file in batch
+                ])
+            return {
+                "message": "Directory copied successfully",
+                "source": source,
+                "destination": destination,
+            }
+        raise Exception("Unsupported file type")
+
+    def format_path(self, path: str) -> str:
+        if path == "/":
+            return "%2F"
+        if path.startswith("/"):
+            path = path[1:]
+        return path

blaxel/sandbox/process.py

@@ -0,0 +1,57 @@
+from blaxel.sandbox.client.models.process_request import ProcessRequest
+
+from .base import SandboxHandleBase
+from .client.api.process.delete_process_identifier import (
+    asyncio_detailed as delete_process_by_identifier,
+)
+from .client.api.process.delete_process_identifier_kill import (
+    asyncio_detailed as delete_process_by_identifier_kill,
+)
+from .client.api.process.get_process import asyncio_detailed as get_process
+from .client.api.process.get_process_identifier import asyncio_detailed as get_process_by_identifier
+from .client.api.process.get_process_identifier_logs import (
+    asyncio_detailed as get_process_by_identifier_logs,
+)
+from .client.api.process.post_process import asyncio_detailed as post_process
+from .client.models import (
+    GetProcessIdentifierLogsResponse200,
+    ProcessKillRequest,
+    ProcessResponse,
+    SuccessResponse,
+)
+
+
+class SandboxProcess(SandboxHandleBase):
+    async def exec(self, process: ProcessRequest) -> ProcessResponse:
+        response = await post_process(client=self.client, body=process)
+        self.handle_response(response)
+        return response.parsed
+
+    async def get(self, identifier: str) -> ProcessResponse:
+        response = await get_process_by_identifier(identifier=identifier, client=self.client)
+        self.handle_response(response)
+        return response.parsed
+
+    async def list(self) -> list[ProcessResponse]:
+        response = await get_process(client=self.client)
+        self.handle_response(response)
+        return response.parsed
+
+    async def stop(self, identifier: str) -> SuccessResponse:
+        response = await delete_process_by_identifier(identifier=identifier, client=self.client)
+        self.handle_response(response)
+        return response.parsed
+
+    async def kill(self, identifier: str, signal: str = "SIGKILL") -> SuccessResponse:
+        kill_request = ProcessKillRequest(signal=signal)
+        response = await delete_process_by_identifier_kill(identifier=identifier, client=self.client, body=kill_request)
+        self.handle_response(response)
+        return response.parsed
+
+    async def logs(self, identifier: str, type_: str = "stdout") -> str:
+        response = await get_process_by_identifier_logs(identifier=identifier, client=self.client)
+        self.handle_response(response)
+        data: GetProcessIdentifierLogsResponse200 = response.parsed
+        if type_ in data.additional_properties:
+            return data.additional_properties[type_]
+        raise Exception("Unsupported log type")

blaxel/sandbox/sandbox.py

@@ -0,0 +1,92 @@
+import asyncio
+import logging
+import time
+from typing import List
+
+from ..client.api.compute.create_sandbox import asyncio as create_sandbox
+from ..client.api.compute.delete_sandbox import asyncio as delete_sandbox
+from ..client.api.compute.get_sandbox import asyncio as get_sandbox
+from ..client.api.compute.list_sandboxes import asyncio as list_sandboxes
+from ..client.client import client
+from ..client.models import Sandbox
+from .filesystem import SandboxFileSystem
+from .process import SandboxProcess
+
+logger = logging.getLogger(__name__)
+
+class SandboxInstance:
+    def __init__(self, sandbox: Sandbox):
+        self.sandbox = sandbox
+        self.fs = SandboxFileSystem(sandbox)
+        self.process = SandboxProcess(sandbox)
+
+    @property
+    def metadata(self):
+        return self.sandbox.metadata
+
+    @property
+    def status(self):
+        return self.sandbox.status
+
+    @property
+    def events(self):
+        return self.sandbox.events
+
+    @property
+    def spec(self):
+        return self.sandbox.spec
+
+    async def wait(self, max_wait: int = 60000, interval: int = 1000) -> None:
+        start_time = time.time() * 1000  # Convert to milliseconds
+        while self.sandbox.status != "DEPLOYED":
+            await asyncio.sleep(interval / 1000)  # Convert to seconds
+            try:
+                response = await get_sandbox(
+                    self.sandbox.metadata.name,
+                    client=client,
+                )
+                logger.info(f"Waiting for sandbox to be deployed, status: {response.status}")
+                self.sandbox = response
+            except Exception as e:
+                logger.error("Could not retrieve sandbox", exc_info=e)
+
+            if self.sandbox.status == "FAILED":
+                raise Exception("Sandbox failed to deploy")
+
+            if (time.time() * 1000) - start_time > max_wait:
+                raise Exception("Sandbox did not deploy in time")
+
+    @classmethod
+    async def create(cls, sandbox: Sandbox) -> "SandboxInstance":
+        if not sandbox.spec:
+            raise Exception("Sandbox spec is required")
+        if not sandbox.spec.runtime:
+            raise Exception("Sandbox runtime is required")
+        sandbox.spec.runtime.generation = sandbox.spec.runtime.generation or "mk3"
+
+        response = await create_sandbox(
+            client=client,
+            body=sandbox,
+        )
+        return cls(response)
+
+    @classmethod
+    async def get(cls, sandbox_name: str) -> "SandboxInstance":
+        response = await get_sandbox(
+            sandbox_name,
+            client=client,
+        )
+        return cls(response)
+
+    @classmethod
+    async def list(cls) -> List["SandboxInstance"]:
+        response = await list_sandboxes()
+        return [cls(sandbox) for sandbox in response]
+
+    @classmethod
+    async def delete(cls, sandbox_name: str) -> Sandbox:
+        response = await delete_sandbox(
+            sandbox_name,
+            client=client,
+        )
+        return response

blaxel/tools/__init__.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import traceback
 from contextlib import AsyncExitStack
 from logging import getLogger
 from typing import Any, cast
@@ -10,6 +11,7 @@ from mcp.types import CallToolResult
 from mcp.types import Tool as MCPTool
 
 from ..common.env import env
+from ..common.internal import get_global_unique_hash
 from ..common.settings import settings
 from ..instrumentation.span import SpanManager
 from ..mcp.client import websocket_client
@@ -22,15 +24,19 @@ if os.getenv("BL_SERVER_PORT"):
     DEFAULT_TIMEOUT = 5
 
 class PersistentWebSocket:
-    def __init__(self, url: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
+    def __init__(self, url: str, name: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
         self.url = url
+        self.name = name
         self.timeout = timeout
-        self.timeout_enabled = timeout_enabled
         self.session_exit_stack = AsyncExitStack()
         self.client_exit_stack = AsyncExitStack()
         self.session: ClientSession = None
         self.timer_task = None
         self.tools_cache = []
+        if settings.bl_cloud:
+            self.timeout_enabled = False
+        else:
+            self.timeout_enabled = timeout_enabled
 
     def with_metas(self, metas: dict[str, Any]):
         self.metas = metas
@@ -46,19 +52,32 @@ class PersistentWebSocket:
         logger.debug(f"Tool {tool_name} returned {call_tool_result}")
         if self.timeout_enabled:
             self._reset_timer()
+        else:
+            await self._close()
         return call_tool_result
 
     async def list_tools(self):
-        await self._initialize()
-        if self.timeout_enabled:
-            self._remove_timer()
-        logger.debug("Listing tools")
-        list_tools_result = await self.session.list_tools()
-        self.tools_cache = list_tools_result.tools
-        logger.debug(f"Tools listed: {list_tools_result}")
-        if self.timeout_enabled:
-            self._reset_timer()
-        return list_tools_result
+        logger.debug(f"Listing tools for {self.name}")
+        span_attributes = {
+            "tool.server": self.url,
+            "tool.server_name": self.name,
+            "span.type": "tool.list",
+        }
+        with SpanManager("blaxel-tracer").create_active_span(self.name, span_attributes) as span:
+            await self._initialize()
+            logger.debug(f"Initialized websocket for {self.name}")
+            if self.timeout_enabled:
+                self._remove_timer()
+            logger.debug("Listing tools")
+            list_tools_result = await self.session.list_tools()
+            self.tools_cache = list_tools_result.tools
+            logger.debug(f"Tools listed: {list_tools_result}")
+            if self.timeout_enabled:
+                self._reset_timer()
+            else:
+                await self._close()
+            span.set_attribute("tool.list.result", list_tools_result.model_dump_json())
+            return list_tools_result
 
     def get_tools(self):
         return self.tools_cache
@@ -125,6 +144,7 @@ def convert_mcp_tool_to_blaxel_tool(
             "tool.args": json.dumps(arguments),
             "tool.server": url,
             "tool.server_name": name,
+            "span.type": "tool.call",
         }
         with SpanManager("blaxel-tracer").create_active_span("blaxel-tool-call", span_attributes):
             logger.debug(f"Calling tool {tool.name} with arguments {arguments}")
@@ -163,22 +183,37 @@ class BlTools:
         self.timeout = timeout
         self.timeout_enabled = timeout_enabled
 
-    def _external_url(self, name: str) -> str:
-        return f"{settings.run_url}/{settings.auth.workspace_name}/functions/{name}"
+    def _internal_url(self, name: str):
+        """Get the internal URL for the agent using a hash of workspace and agent name."""
+        hash = get_global_unique_hash(settings.workspace, "function", name)
+        return f"{settings.run_internal_protocol}://bl-{settings.env}-{hash}.{settings.run_internal_hostname}"
 
-    def _url(self, name: str) -> str:
+    def _forced_url(self, name: str):
+        """Get the forced URL from environment variables if set."""
         env_var = name.replace("-", "_").upper()
         if env[f"BL_FUNCTION_{env_var}_URL"]:
             return env[f"BL_FUNCTION_{env_var}_URL"]
-        elif env[f"BL_FUNCTION_{env_var}_SERVICE_NAME"]:
-            return f"https://{env[f'BL_FUNCTION_{env_var}_SERVICE_NAME']}.{settings.run_internal_hostname}"
-        return self._external_url(name)
+        return None
+
+    def _external_url(self, name: str):
+        return f"{settings.run_url}/{settings.workspace}/functions/{name}"
 
-    def _fallback_url(self, name: str) -> str | None:
+    def _fallback_url(self, name: str):
         if self._external_url(name) != self._url(name):
             return self._external_url(name)
         return None
 
+    def _url(self, name: str):
+        logger.debug(f"Getting URL for {name}")
+        if self._forced_url(name):
+            logger.debug(f"Forced URL found for {name}: {self._forced_url(name)}")
+            return self._forced_url(name)
+        if settings.run_internal_hostname:
+            logger.debug(f"Internal hostname found for {name}: {self._internal_url(name)}")
+            return self._internal_url(name)
+        logger.debug(f"No URL found for {name}, using external URL")
+        return self._external_url(name)
+
     def get_tools(self) -> list[Tool]:
         """Get a list of all tools from all connected servers."""
         all_tools: list[Tool] = []
@@ -240,8 +275,13 @@ class BlTools:
         except Exception as e:
             if not self._fallback_url(name):
                 raise e
+            logger.warning(f"Error connecting to {name}: {e}\n{traceback.format_exc()}")
             url = self._fallback_url(name)
-            await self.connect_with_url(name, url)
+            try:
+                await self.connect_with_url(name, url)
+            except Exception as e:
+                logger.error(f"Error connecting to {name} with fallback URL: {e}\n{traceback.format_exc()}")
+                raise e
 
     async def connect_with_url(
         self, name: str, url: str
@@ -255,7 +295,8 @@ class BlTools:
         logger.debug(f"Initializing session and loading tools from {url}")
 
         if not toolPersistances.get(name):
-            toolPersistances[name] = PersistentWebSocket(url, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
+            logger.debug(f"Creating new persistent websocket for {name}")
+            toolPersistances[name] = PersistentWebSocket(url, name, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
             await toolPersistances[name].list_tools()
         logger.debug(f"Loaded {len(toolPersistances[name].get_tools())} tools from {url}")
         return toolPersistances[name].with_metas(self.metas)

{blaxel-0.1.9rc36.dist-info → blaxel-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.1.9rc36
+Version: 0.1.10
 Summary: Add your description here
 Author-email: cploujoux <cploujoux@blaxel.ai>
 License-File: LICENSE