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

blaxel/{sandbox → core/sandbox}/sandbox.py

@@ -1,26 +1,33 @@
 import asyncio
 import logging
 import time
-from typing import List
+from typing import Any, Dict, List, Union
 
 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 ..client.models import Metadata, Sandbox
 from .filesystem import SandboxFileSystem
+from .network import SandboxNetwork
 from .preview import SandboxPreviews
 from .process import SandboxProcess
+from .session import SandboxSessions
+from .types import SandboxConfiguration, SessionWithToken
 
 logger = logging.getLogger(__name__)
 
+
 class SandboxInstance:
     def __init__(self, sandbox: Sandbox):
         self.sandbox = sandbox
-        self.fs = SandboxFileSystem(sandbox)
-        self.process = SandboxProcess(sandbox)
+        self.config = SandboxConfiguration(sandbox)
+        self.fs = SandboxFileSystem(self.config)
+        self.process = SandboxProcess(self.config)
         self.previews = SandboxPreviews(sandbox)
+        self.sessions = SandboxSessions(self.config)
+        self.network = SandboxNetwork(self.config)
 
     @property
     def metadata(self):
@@ -49,6 +56,7 @@ class SandboxInstance:
                 )
                 logger.info(f"Waiting for sandbox to be deployed, status: {response.status}")
                 self.sandbox = response
+                self.config = SandboxConfiguration(self.sandbox)
             except Exception as e:
                 logger.error("Could not retrieve sandbox", exc_info=e)
 
@@ -59,7 +67,10 @@ class SandboxInstance:
                 raise Exception("Sandbox did not deploy in time")
 
     @classmethod
-    async def create(cls, sandbox: Sandbox) -> "SandboxInstance":
+    async def create(cls, sandbox: Union[Sandbox, Dict[str, Any]]) -> "SandboxInstance":
+        if isinstance(sandbox, dict):
+            sandbox = Sandbox.from_dict(sandbox)
+
         if not sandbox.spec:
             raise Exception("Sandbox spec is required")
         if not sandbox.spec.runtime:
@@ -92,3 +103,49 @@ class SandboxInstance:
             client=client,
         )
         return response
+
+    @classmethod
+    async def create_if_not_exists(
+        cls, sandbox: Union[Sandbox, Dict[str, Any]]
+    ) -> "SandboxInstance":
+        """Create a sandbox if it doesn't exist, otherwise return existing."""
+        if isinstance(sandbox, dict):
+            sandbox = Sandbox.from_dict(sandbox)
+
+        try:
+            sandbox_instance = await cls.get(sandbox.metadata.name)
+            return sandbox_instance
+        except Exception as e:
+            # Check if it's a 404 error (sandbox not found)
+            if hasattr(e, "status_code") and e.status_code == 404:
+                return await cls.create(sandbox)
+            raise e
+
+    @classmethod
+    async def from_session(
+        cls, session: Union[SessionWithToken, Dict[str, Any]]
+    ) -> "SandboxInstance":
+        """Create a sandbox instance from a session with token."""
+        if isinstance(session, dict):
+            session = SessionWithToken.from_dict(session)
+
+        # Create a minimal sandbox configuration for session-based access
+        sandbox_name = session.name.split("-")[0] if "-" in session.name else session.name
+        sandbox = Sandbox(metadata=Metadata(name=sandbox_name))
+        config = SandboxConfiguration(
+            sandbox=sandbox,
+            force_url=session.url,
+            headers={"X-Blaxel-Preview-Token": session.token},
+            params={"bl_preview_token": session.token},
+        )
+
+        instance = cls.__new__(cls)
+        instance.sandbox = sandbox
+        instance.config = config
+        instance.fs = SandboxFileSystem(config)
+        instance.process = SandboxProcess(config)
+        instance.previews = SandboxPreviews(sandbox)
+        instance.sessions = SandboxSessions(config)
+        instance.network = SandboxNetwork(config)
+
+        return instance

blaxel/core/sandbox/session.py

@@ -0,0 +1,124 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+from ..client.api.compute import (
+    create_sandbox_preview,
+    delete_sandbox_preview,
+    get_sandbox_preview,
+    list_sandbox_preview_tokens,
+    list_sandbox_previews,
+)
+from ..client.client import client
+from ..client.models import Metadata, Preview, PreviewSpec
+from .preview import SandboxPreview
+from .types import SandboxConfiguration, SessionCreateOptions, SessionWithToken
+
+
+class SandboxSessions:
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        self.sandbox_config = sandbox_config
+
+    @property
+    def sandbox_name(self) -> str:
+        return self.sandbox_config.metadata.name if self.sandbox_config.metadata else ""
+
+    async def create(
+        self, options: Optional[Union[SessionCreateOptions, Dict[str, Any]]] = None
+    ) -> SessionWithToken:
+        if options is None:
+            options = SessionCreateOptions()
+        elif isinstance(options, dict):
+            options = SessionCreateOptions.from_dict(options)
+
+        expires_at = options.expires_at or datetime.now() + timedelta(days=1)
+
+        preview_body = Preview(
+            metadata=Metadata(name=f"session-{int(datetime.now().timestamp() * 1000)}"),
+            spec=PreviewSpec(
+                port=443,
+                public=False,
+                request_headers=options.request_headers,
+                response_headers=options.response_headers,
+            ),
+        )
+
+        preview_response = await create_sandbox_preview.asyncio(
+            self.sandbox_name, client=client, body=preview_body
+        )
+
+        preview = SandboxPreview(preview_response, self.sandbox_name)
+        token_obj = await preview.tokens.create(expires_at)
+
+        return SessionWithToken(
+            name=preview_body.metadata.name,
+            url=preview.spec.url or "",
+            token=token_obj.value,
+            expires_at=token_obj.expires_at,
+        )
+
+    async def create_if_expired(
+        self,
+        options: Optional[Union[SessionCreateOptions, Dict[str, Any]]] = None,
+        delta_seconds: int = 3600,  # 1 hour
+    ) -> SessionWithToken:
+        if options is None:
+            options = SessionCreateOptions()
+        elif isinstance(options, dict):
+            options = SessionCreateOptions.from_dict(options)
+
+        all_sessions = await self.list()
+        now = datetime.now()
+        threshold = now + timedelta(seconds=delta_seconds)
+
+        if all_sessions:
+            session_data = all_sessions[0]
+            if session_data.expires_at < threshold:
+                await self.delete(session_data.name)
+                session_data = await self.create(options)
+        else:
+            session_data = await self.create(options)
+
+        return session_data
+
+    async def list(self) -> List[SessionWithToken]:
+        previews_response = await list_sandbox_previews.asyncio(self.sandbox_name, client=client)
+
+        sessions = []
+        for preview in previews_response:
+            if preview.metadata and preview.metadata.name and "session-" in preview.metadata.name:
+                token = await self.get_token(preview.metadata.name)
+                if token:
+                    sessions.append(
+                        SessionWithToken(
+                            name=preview.metadata.name,
+                            url=preview.spec.url or "",
+                            token=token.spec.token or "",
+                            expires_at=token.spec.expires_at or datetime.now(),
+                        )
+                    )
+
+        return sessions
+
+    async def get(self, name: str) -> SessionWithToken:
+        preview_response = await get_sandbox_preview.asyncio(self.sandbox_name, name, client=client)
+
+        token = await self.get_token(name)
+
+        return SessionWithToken(
+            name=name,
+            url=preview_response.spec.url or "",
+            token=token.spec.token or "" if token else "",
+            expires_at=token.spec.expires_at or datetime.now() if token else datetime.now(),
+        )
+
+    async def delete(self, name: str):
+        return await delete_sandbox_preview.asyncio(self.sandbox_name, name, client=client)
+
+    async def get_token(self, preview_name: str):
+        tokens_response = await list_sandbox_preview_tokens.asyncio(
+            self.sandbox_name, preview_name, client=client
+        )
+
+        if not tokens_response:
+            return None
+        return tokens_response[0]

blaxel/core/sandbox/types.py

@@ -0,0 +1,103 @@
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from ..client.models import Sandbox
+
+
+class SessionCreateOptions:
+    def __init__(
+        self,
+        expires_at: Optional[datetime] = None,
+        response_headers: Optional[Dict[str, str]] = None,
+        request_headers: Optional[Dict[str, str]] = None,
+    ):
+        self.expires_at = expires_at
+        self.response_headers = response_headers or {}
+        self.request_headers = request_headers or {}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionCreateOptions":
+        expires_at = None
+        if "expires_at" in data and data["expires_at"]:
+            if isinstance(data["expires_at"], str):
+                expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))
+            elif isinstance(data["expires_at"], datetime):
+                expires_at = data["expires_at"]
+
+        return cls(
+            expires_at=expires_at,
+            response_headers=data.get("response_headers"),
+            request_headers=data.get("request_headers"),
+        )
+
+
+class SessionWithToken:
+    def __init__(self, name: str, url: str, token: str, expires_at: datetime):
+        self.name = name
+        self.url = url
+        self.token = token
+        self.expires_at = expires_at
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionWithToken":
+        expires_at = data["expires_at"]
+        if isinstance(expires_at, str):
+            expires_at = datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
+
+        return cls(
+            name=data["name"],
+            url=data["url"],
+            token=data["token"],
+            expires_at=expires_at,
+        )
+
+
+class SandboxConfiguration:
+    def __init__(
+        self,
+        sandbox: Sandbox,
+        force_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, str]] = None,
+    ):
+        self.sandbox = sandbox
+        self.force_url = force_url
+        self.headers = headers or {}
+        self.params = params or {}
+
+    @property
+    def metadata(self):
+        return self.sandbox.metadata
+
+    @property
+    def status(self):
+        return self.sandbox.status
+
+    @property
+    def spec(self):
+        return self.sandbox.spec
+
+
+class WatchEvent:
+    def __init__(self, op: str, path: str, name: str, content: Optional[str] = None):
+        self.op = op
+        self.path = path
+        self.name = name
+        self.content = content
+
+
+class SandboxFilesystemFile:
+    def __init__(self, path: str, content: str):
+        self.path = path
+        self.content = content
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SandboxFilesystemFile":
+        return cls(data["path"], data["content"])
+
+
+class CopyResponse:
+    def __init__(self, message: str, source: str, destination: str):
+        self.message = message
+        self.source = source
+        self.destination = destination

blaxel/{tools → core/tools}/__init__.py

@@ -1,5 +1,4 @@
 import asyncio
-import json
 import os
 import traceback
 from contextlib import AsyncExitStack
@@ -12,7 +11,6 @@ from mcp.types import Tool as MCPTool
 
 from ..common.internal import get_forced_url, get_global_unique_hash
 from ..common.settings import settings
-from ..instrumentation.span import SpanManager
 from ..mcp.client import websocket_client
 from .types import Tool
 
@@ -22,8 +20,11 @@ DEFAULT_TIMEOUT = 1
 if os.getenv("BL_SERVER_PORT"):
     DEFAULT_TIMEOUT = 5
 
+
 class PersistentWebSocket:
-    def __init__(self, url: str, name: 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
@@ -58,32 +59,30 @@ class PersistentWebSocket:
         except Exception as e:
             logger.error(f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}")
             return CallToolResult(
-                content=[{"type": "text", "text": f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}"}],
+                content=[
+                    {
+                        "type": "text",
+                        "text": f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}",
+                    }
+                ],
                 isError=True,
             )
 
     async def list_tools(self):
         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
+        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()
+        return list_tools_result
 
     def get_tools(self):
         return self.tools_cache
@@ -91,8 +90,13 @@ class PersistentWebSocket:
     async def _initialize(self):
         if not self.session:
             logger.debug(f"Initializing websocket client for {self.url}")
-            read, write = await self.client_exit_stack.enter_async_context(websocket_client(self.url, settings.headers))
-            self.session = cast(ClientSession, await self.session_exit_stack.enter_async_context(ClientSession(read, write)))
+            read, write = await self.client_exit_stack.enter_async_context(
+                websocket_client(self.url, settings.headers)
+            )
+            self.session = cast(
+                ClientSession,
+                await self.session_exit_stack.enter_async_context(ClientSession(read, write)),
+            )
             await self.session.initialize()
 
     def _reset_timer(self):
@@ -108,7 +112,6 @@ class PersistentWebSocket:
         await self._close()
         self.session = None
 
-
     async def _close(self):
         logger.debug(f"Closing websocket client {self.url}")
         if self.session:
@@ -123,6 +126,7 @@ class PersistentWebSocket:
                 logger.debug(f"Error closing client exit stack: {e}")
             logger.debug("WebSocket connection closed due to inactivity.")
 
+
 def convert_mcp_tool_to_blaxel_tool(
     websocket_client: PersistentWebSocket,
     name: str,
@@ -145,18 +149,10 @@ def convert_mcp_tool_to_blaxel_tool(
         *args: Any,
         **arguments: dict[str, Any],
     ) -> CallToolResult:
-        span_attributes = {
-            "tool.name": tool.name,
-            "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}")
-            call_tool_result = await websocket_client.call_tool(tool.name, arguments)
-            logger.debug(f"Tool {tool.name} returned {call_tool_result}")
-            return call_tool_result
+        logger.debug(f"Calling tool {tool.name} with arguments {arguments}")
+        call_tool_result = await websocket_client.call_tool(tool.name, arguments)
+        logger.debug(f"Tool {tool.name} returned {call_tool_result}")
+        return call_tool_result
 
     async def call_tool(
         *args: Any,
@@ -180,10 +176,18 @@ def convert_mcp_tool_to_blaxel_tool(
         response_format="content_and_artifact",
     )
 
+
 toolPersistances: dict[str, PersistentWebSocket] = {}
 
+
 class BlTools:
-    def __init__(self, functions: list[str], metas: dict[str, Any] = {}, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
+    def __init__(
+        self,
+        functions: list[str],
+        metas: dict[str, Any] = {},
+        timeout: int = DEFAULT_TIMEOUT,
+        timeout_enabled: bool = True,
+    ):
         self.functions = functions
         self.metas = metas
         self.timeout = timeout
@@ -224,52 +228,13 @@ class BlTools:
             toolPersistances.get(name).with_metas(self.metas)
             websocket = toolPersistances.get(name)
             tools = websocket.get_tools()
-            converted_tools = [convert_mcp_tool_to_blaxel_tool(websocket, name, self._url(name), tool) for tool in tools]
+            converted_tools = [
+                convert_mcp_tool_to_blaxel_tool(websocket, name, self._url(name), tool)
+                for tool in tools
+            ]
             all_tools.extend(converted_tools)
         return all_tools
 
-    async def to_langchain(self):
-        from .langchain import get_langchain_tools
-
-        await self.initialize()
-        return get_langchain_tools(self.get_tools())
-
-    async def to_llamaindex(self):
-        from .llamaindex import get_llamaindex_tools
-
-        await self.initialize()
-        return get_llamaindex_tools(self.get_tools())
-
-    async def to_crewai(self):
-        from .crewai import get_crewai_tools
-
-        await self.initialize()
-        return get_crewai_tools(self.get_tools())
-
-    async def to_openai(self):
-        from .openai import get_openai_tools
-
-        await self.initialize()
-        return get_openai_tools(self.get_tools())
-
-    async def to_pydantic(self):
-        from .pydantic import get_pydantic_tools
-
-        await self.initialize()
-        return get_pydantic_tools(self.get_tools())
-
-    async def to_google_adk(self):
-        from .googleadk import get_google_adk_tools
-
-        await self.initialize()
-        return get_google_adk_tools(self.get_tools())
-
-    async def to_livekit(self):
-        from .livekit import get_livekit_tools
-
-        await self.initialize()
-        return get_livekit_tools(self.get_tools())
-
     async def connect(self, name: str):
         # Create and store the connection
         try:
@@ -283,12 +248,12 @@ class BlTools:
             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()}")
+                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
-    ) -> None:
+    async def connect_with_url(self, name: str, url: str) -> None:
         """Initialize a session and load tools from it.
 
         Args:
@@ -299,17 +264,24 @@ class BlTools:
 
         if not toolPersistances.get(name):
             logger.debug(f"Creating new persistent websocket for {name}")
-            toolPersistances[name] = PersistentWebSocket(url, name, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
+            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)
 
     async def initialize(self) -> "BlTools":
         for i in range(0, len(self.functions), 10):
-            batch = self.functions[i:i+10]
+            batch = self.functions[i : i + 10]
             await asyncio.gather(*(self.connect(name) for name in batch))
         return self
 
 
-def bl_tools(functions: list[str], metas: dict[str, Any] = {}, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True) -> BlTools:
+def bl_tools(
+    functions: list[str],
+    metas: dict[str, Any] = {},
+    timeout: int = DEFAULT_TIMEOUT,
+    timeout_enabled: bool = True,
+) -> BlTools:
     return BlTools(functions, metas=metas, timeout=timeout, timeout_enabled=timeout_enabled)

blaxel/crewai/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel CrewAI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/crewai.py → crewai/model.py}

@@ -2,11 +2,13 @@ from logging import getLogger
 
 from crewai import LLM
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
-async def get_crewai_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return LLM(
             model=f"mistral/{model}",

blaxel/crewai/tools.py

@@ -0,0 +1,26 @@
+from crewai.tools import BaseTool
+
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.common import create_model_from_json_schema
+from blaxel.core.tools.types import Tool
+
+
+class CrewAITool(BaseTool):
+    _tool: Tool
+
+    def __init__(self, tool: Tool):
+        super().__init__(
+            name=tool.name,
+            description=tool.description,
+            args_schema=create_model_from_json_schema(tool.input_schema),
+        )
+        self._tool = tool
+
+    def _run(self, *args, **kwargs):
+        return self._tool.sync_coroutine(**kwargs)
+
+async def bl_tools(tools_names: list[str], **kwargs) -> list[BaseTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [CrewAITool(tool) for tool in tools.get_tools()]
+    return [CrewAITool(tool) for tool in tools.get_tools()]

blaxel/googleadk/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel Google ADK integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/googleadk.py → googleadk/model.py}

@@ -2,7 +2,8 @@ from logging import getLogger
 
 from google.adk.models.lite_llm import LiteLlm
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
@@ -64,4 +65,9 @@ async def get_google_adk_model(url: str, type: str, model: str, **kwargs):
             api_key=settings.auth.token,
             api_base=f"{url}/v1",
             **kwargs
-        )
+        )
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
+    return await get_google_adk_model(url, type, model, **kwargs)
+    return await get_google_adk_model(url, type, model, **kwargs)