blaxel 0.2.36__py3-none-any.whl → 0.2.38rc122__py3-none-any.whl

blaxel/core/sandbox/types.py

@@ -4,7 +4,14 @@ from typing import Any, Callable, Dict, List, TypeVar, Union
 import httpx
 from attrs import define as _attrs_define
 
-from ..client.models import Port, Sandbox, SandboxLifecycle, VolumeAttachment
+from ..client.models import (
+    Env,
+    Port,
+    PortProtocol,
+    Sandbox,
+    SandboxLifecycle,
+    VolumeAttachment,
+)
 from ..client.types import UNSET
 from .client.models.process_request import ProcessRequest
 from .client.models.process_response import ProcessResponse
@@ -205,7 +212,7 @@ class SandboxCreateConfiguration:
             if isinstance(port, Port):
                 # If it's already a Port object, ensure protocol defaults to HTTP
                 if port.protocol is UNSET or not port.protocol:
-                    port.protocol = "HTTP"
+                    port.protocol = PortProtocol.HTTP
                 port_objects.append(port)
             elif isinstance(port, dict):
                 # Convert dict to Port object with HTTP as default protocol
@@ -218,20 +225,22 @@ class SandboxCreateConfiguration:
 
         return port_objects
 
-    def _normalize_envs(self) -> List[Dict[str, str]] | None:
+    def _normalize_envs(self) -> List[Env] | None:
         """Convert envs to list of dicts with name and value keys."""
         if not self.envs:
             return None
 
         env_objects = []
         for env in self.envs:
-            if isinstance(env, dict):
+            if isinstance(env, Env):
+                env_objects.append(env)
+            elif isinstance(env, dict):
                 # Validate that the dict has the required keys
                 if "name" not in env or "value" not in env:
                     raise ValueError(
                         f"Environment variable dict must have 'name' and 'value' keys: {env}"
                     )
-                env_objects.append({"name": env["name"], "value": env["value"]})
+                env_objects.append(Env(name=env["name"], value=env["value"]))
             else:
                 raise ValueError(
                     f"Invalid env type: {type(env)}. Expected dict with 'name' and 'value' keys."
@@ -385,3 +394,201 @@ class Context:
     @classmethod
     def from_json(cls, data: Dict[str, Any]) -> "Context":
         return cls(id=str(data.get("id") or data.get("context_id") or ""))
+
+
+class StreamHandle:
+    """Handle for managing a streaming operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "StreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncStreamHandle:
+    """Handle for managing a streaming operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncStreamHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncStreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class WatchHandle:
+    """Handle for managing a file system watch operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "WatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncWatchHandle:
+    """Handle for managing a file system watch operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncWatchHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncWatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)

blaxel/core/volume/volume.py

@@ -1,3 +1,5 @@
+import asyncio
+import time
 import uuid
 from typing import Callable, Dict, List, Union
 
@@ -500,6 +502,8 @@ async def _update_volume_by_name(
         status_code = int(response.code) if response.code is not UNSET else None
         message = response.message if response.message is not UNSET else response.error
         raise VolumeAPIError(message, status_code=status_code, code=response.error)
+    # This is for safe update
+    await asyncio.sleep(0.5)
     return VolumeInstance(response)
 
 
@@ -576,6 +580,8 @@ def _update_volume_by_name_sync(
         status_code = int(response.code) if response.code is not UNSET else None
         message = response.message if response.message is not UNSET else response.error
         raise VolumeAPIError(message, status_code=status_code, code=response.error)
+    # This is for safe update
+    time.sleep(0.5)
     return SyncVolumeInstance(response)
 
 

blaxel/langgraph/tools.py

@@ -5,7 +5,6 @@ from blaxel.core.tools.types import Tool, ToolException
 
 if TYPE_CHECKING:
     from langchain_core.tools import StructuredTool
-    from mcp.types import EmbeddedResource, ImageContent
 
 
 def _clean_schema_for_openai(schema: Dict[str, Any]) -> Dict[str, Any]:

blaxel/llamaindex/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import os
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any, Dict, List, Sequence, Union
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -20,20 +20,61 @@ if TYPE_CHECKING:
         CompletionResponseAsyncGen,
         CompletionResponseGen,
     )
+    from llama_index.core.llms.llm import ToolSelection
+    from llama_index.core.tools.types import BaseTool
+
+# Runtime imports needed for class inheritance and construction
+from llama_index.core.base.llms.types import LLMMetadata  # noqa: E402
+from llama_index.core.llms.function_calling import FunctionCallingLLM  # noqa: E402
+from pydantic import PrivateAttr  # noqa: E402
 
 logger = getLogger(__name__)
 
+DEFAULT_CONTEXT_WINDOW = 128000
+DEFAULT_NUM_OUTPUT = 4096
+
+
+class TokenRefreshingLLM(FunctionCallingLLM):
+    """Wrapper for LlamaIndex LLMs that refreshes token before each call.
 
-class TokenRefreshingWrapper:
-    """Base wrapper class that refreshes token before each call."""
+    Inherits from FunctionCallingLLM to maintain type compatibility with
+    LlamaIndex's agents and components that validate isinstance(model, LLM).
+    """
+
+    _model_config_data: dict = PrivateAttr(default_factory=dict)
+    _wrapped: Any = PrivateAttr(default=None)
 
     def __init__(self, model_config: dict):
-        self.model_config = model_config
-        self.wrapped_model = self._create_model()
+        super().__init__()
+        self._model_config_data = model_config
+        self._wrapped = self._create_model()
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "TokenRefreshingLLM"
+
+    @property
+    def wrapped_model(self) -> Any:
+        """Access the underlying wrapped LLM model."""
+        return self._wrapped
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        """Get LLM metadata, with fallback for unknown model names."""
+        try:
+            return self._wrapped.metadata
+        except (ValueError, KeyError) as e:
+            logger.warning(f"Could not get metadata from wrapped model: {e}. Using defaults.")
+            return LLMMetadata(
+                context_window=DEFAULT_CONTEXT_WINDOW,
+                num_output=DEFAULT_NUM_OUTPUT,
+                is_chat_model=True,
+                model_name=self._model_config_data.get("model", "unknown"),
+            )
 
     def _create_model(self):
         """Create the model instance with current token."""
-        config = self.model_config
+        config = self._model_config_data
         model_type = config["type"]
         model = config["model"]
         url = config["url"]
@@ -115,102 +156,106 @@ class TokenRefreshingWrapper:
 
     def _refresh_token(self):
         """Refresh the token and recreate the model if needed."""
-        # Only refresh if using ClientCredentials (which has get_token method)
         current_token = settings.auth.token
 
-        if hasattr(settings.auth, "get_token"):
-            # This will trigger token refresh if needed
-            settings.auth.get_token()
-
         new_token = settings.auth.token
 
-        # If token changed, recreate the model
         if current_token != new_token:
-            self.wrapped_model = self._create_model()
+            self._wrapped = self._create_model()
 
-    def __getattr__(self, name):
-        """Delegate attribute access to wrapped model."""
-        return getattr(self.wrapped_model, name)
+    # --- Core LLM methods with token refresh ---
 
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        self._refresh_token()
+        return self._wrapped.chat(messages, **kwargs)
 
-class TokenRefreshingLLM(TokenRefreshingWrapper):
-    """Wrapper for LlamaIndex LLMs that refreshes token before each call."""
-
-    async def achat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
-    ) -> ChatResponse:
-        """Async chat with token refresh."""
+    async def achat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         self._refresh_token()
-        return await self.wrapped_model.achat(messages, **kwargs)
+        return await self._wrapped.achat(messages, **kwargs)
 
-    def chat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
-    ) -> ChatResponse:
-        """Sync chat with token refresh."""
+    def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
         self._refresh_token()
-        return self.wrapped_model.chat(messages, **kwargs)
+        return self._wrapped.complete(prompt, formatted=formatted, **kwargs)
 
-    async def astream_chat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
-    ) -> ChatResponseAsyncGen:
-        """Async stream chat with token refresh."""
+    async def acomplete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
         self._refresh_token()
-        async for chunk in self.wrapped_model.astream_chat(messages, **kwargs):
-            yield chunk
+        return await self._wrapped.acomplete(prompt, formatted=formatted, **kwargs)
 
     def stream_chat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
+        self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponseGen:
-        """Sync stream chat with token refresh."""
         self._refresh_token()
-        for chunk in self.wrapped_model.stream_chat(messages, **kwargs):
-            yield chunk
+        return self._wrapped.stream_chat(messages, **kwargs)
 
-    async def acomplete(
-        self,
-        prompt: str,
-        **kwargs: Any,
-    ) -> CompletionResponse:
-        """Async complete with token refresh."""
+    async def astream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseAsyncGen:
         self._refresh_token()
-        return await self.wrapped_model.acomplete(prompt, **kwargs)
+        result = self._wrapped.astream_chat(messages, **kwargs)
+        # Handle both coroutine and async generator patterns
+        if hasattr(result, "__aiter__"):
+            return result
+        return await result
 
-    def complete(
-        self,
-        prompt: str,
-        **kwargs: Any,
-    ) -> CompletionResponse:
-        """Sync complete with token refresh."""
+    def stream_complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponseGen:
         self._refresh_token()
-        return self.wrapped_model.complete(prompt, **kwargs)
+        return self._wrapped.stream_complete(prompt, formatted=formatted, **kwargs)
 
     async def astream_complete(
-        self,
-        prompt: str,
-        **kwargs: Any,
+        self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponseAsyncGen:
-        """Async stream complete with token refresh."""
         self._refresh_token()
-        async for chunk in self.wrapped_model.astream_complete(prompt, **kwargs):
-            yield chunk
+        result = self._wrapped.astream_complete(prompt, formatted=formatted, **kwargs)
+        # Handle both coroutine and async generator patterns
+        if hasattr(result, "__aiter__"):
+            return result
+        return await result
 
-    def stream_complete(
+    # --- FunctionCallingLLM methods (delegate to wrapped model) ---
+
+    def _prepare_chat_with_tools(
         self,
-        prompt: str,
+        tools: Sequence[BaseTool],
+        user_msg: Union[str, ChatMessage, None] = None,
+        chat_history: List[ChatMessage] | None = None,
+        verbose: bool = False,
+        allow_parallel_tool_calls: bool = False,
+        tool_required: Any = None,
         **kwargs: Any,
-    ) -> CompletionResponseGen:
-        """Sync stream complete with token refresh."""
-        self._refresh_token()
-        for chunk in self.wrapped_model.stream_complete(prompt, **kwargs):
-            yield chunk
+    ) -> Dict[str, Any]:
+        if hasattr(self._wrapped, "_prepare_chat_with_tools"):
+            return self._wrapped._prepare_chat_with_tools(
+                tools,
+                user_msg=user_msg,
+                chat_history=chat_history,
+                verbose=verbose,
+                allow_parallel_tool_calls=allow_parallel_tool_calls,
+                tool_required=tool_required,
+                **kwargs,
+            )
+        raise NotImplementedError(
+            f"The wrapped model ({type(self._wrapped).__name__}) does not support function calling"
+        )
+
+    def get_tool_calls_from_response(
+        self,
+        response: ChatResponse,
+        error_on_no_tool_call: bool = True,
+        **kwargs: Any,
+    ) -> List[ToolSelection]:
+        if hasattr(self._wrapped, "get_tool_calls_from_response"):
+            return self._wrapped.get_tool_calls_from_response(
+                response,
+                error_on_no_tool_call=error_on_no_tool_call,
+                **kwargs,
+            )
+        raise NotImplementedError(
+            f"The wrapped model ({type(self._wrapped).__name__}) does not support function calling"
+        )
 
 
 async def bl_model(name, **kwargs):