blaxel 0.2.35__py3-none-any.whl → 0.2.37__py3-none-any.whl

blaxel/core/volume/volume.py

@@ -1,3 +1,5 @@
+import asyncio
+import time
 import uuid
 from typing import Callable, Dict, List, Union
 
@@ -9,10 +11,13 @@ from ..client.api.volumes.get_volume import asyncio as get_volume
 from ..client.api.volumes.get_volume import sync as get_volume_sync
 from ..client.api.volumes.list_volumes import asyncio as list_volumes
 from ..client.api.volumes.list_volumes import sync as list_volumes_sync
+from ..client.api.volumes.update_volume import asyncio as update_volume
+from ..client.api.volumes.update_volume import sync as update_volume_sync
 from ..client.client import client
 from ..client.models import Metadata, Volume, VolumeSpec
 from ..client.models.error import Error
 from ..client.types import UNSET
+from ..common.settings import settings
 
 
 class VolumeAPIError(Exception):
@@ -60,6 +65,46 @@ class _SyncDeleteDescriptor:
             return instance_delete
 
 
+class _AsyncUpdateDescriptor:
+    """Descriptor that provides both class-level and instance-level update functionality."""
+
+    def __init__(self, update_func: Callable):
+        self._update_func = update_func
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            # Called on the class: VolumeInstance.update("name", updates)
+            return self._update_func
+        else:
+            # Called on an instance: instance.update(updates)
+            async def instance_update(
+                updates: Union["VolumeCreateConfiguration", Volume, Dict[str, any]],
+            ) -> "VolumeInstance":
+                return await self._update_func(instance.metadata.name or "", updates)
+
+            return instance_update
+
+
+class _SyncUpdateDescriptor:
+    """Descriptor that provides both class-level and instance-level update functionality (sync)."""
+
+    def __init__(self, update_func: Callable):
+        self._update_func = update_func
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            # Called on the class: SyncVolumeInstance.update("name", updates)
+            return self._update_func
+        else:
+            # Called on an instance: instance.update(updates)
+            def instance_update(
+                updates: Union["VolumeCreateConfiguration", Volume, Dict[str, any]],
+            ) -> "SyncVolumeInstance":
+                return self._update_func(instance.metadata.name or "", updates)
+
+            return instance_update
+
+
 class VolumeCreateConfiguration:
     """Simplified configuration for creating volumes with default values."""
 
@@ -147,7 +192,7 @@ class VolumeInstance:
                 ),
                 spec=VolumeSpec(
                     size=config.size or default_size,
-                    region=config.region or UNSET,
+                    region=config.region or settings.region or UNSET,
                     template=config.template or UNSET,
                 ),
             )
@@ -161,7 +206,7 @@ class VolumeInstance:
                 ),
                 spec=VolumeSpec(
                     size=volume_config.size or default_size,
-                    region=volume_config.region or UNSET,
+                    region=volume_config.region or settings.region or UNSET,
                     template=volume_config.template or UNSET,
                 ),
             )
@@ -288,7 +333,7 @@ class SyncVolumeInstance:
                 ),
                 spec=VolumeSpec(
                     size=config.size or default_size,
-                    region=config.region or UNSET,
+                    region=config.region or settings.region or UNSET,
                     template=config.template or UNSET,
                 ),
             )
@@ -302,7 +347,7 @@ class SyncVolumeInstance:
                 ),
                 spec=VolumeSpec(
                     size=volume_config.size or default_size,
-                    region=volume_config.region or UNSET,
+                    region=volume_config.region or settings.region or UNSET,
                     template=volume_config.template or UNSET,
                 ),
             )
@@ -384,6 +429,166 @@ def _delete_volume_by_name_sync(volume_name: str) -> Volume:
     return response
 
 
+async def _update_volume_by_name(
+    volume_name: str, updates: Union[VolumeCreateConfiguration, Volume, Dict[str, any]]
+) -> "VolumeInstance":
+    """Update a volume by name (async)."""
+    # Get the current volume
+    volume_instance = await VolumeInstance.get(volume_name)
+    current_volume = volume_instance.volume
+
+    # Build the update body
+    if isinstance(updates, Volume):
+        new_metadata = updates.metadata
+        new_spec = updates.spec
+    elif isinstance(updates, VolumeCreateConfiguration):
+        new_metadata = Metadata(
+            name=current_volume.metadata.name if current_volume.metadata else volume_name,
+            display_name=updates.display_name,
+            labels=updates.labels,
+        )
+        new_spec = VolumeSpec(
+            size=updates.size,
+            region=updates.region,
+            template=updates.template,
+        )
+    elif isinstance(updates, dict):
+        config = VolumeCreateConfiguration.from_dict(updates)
+        new_metadata = Metadata(
+            name=current_volume.metadata.name if current_volume.metadata else volume_name,
+            display_name=config.display_name,
+            labels=config.labels,
+        )
+        new_spec = VolumeSpec(
+            size=config.size,
+            region=config.region,
+            template=config.template,
+        )
+    else:
+        raise ValueError(
+            f"Invalid updates type: {type(updates)}. Expected VolumeCreateConfiguration, Volume, or dict."
+        )
+
+    # Merge current values with updates
+    merged_metadata = Metadata(
+        name=current_volume.metadata.name if current_volume.metadata else volume_name,
+        display_name=new_metadata.display_name
+        if new_metadata and new_metadata.display_name
+        else (current_volume.metadata.display_name if current_volume.metadata else None),
+        labels=new_metadata.labels
+        if new_metadata and new_metadata.labels
+        else (current_volume.metadata.labels if current_volume.metadata else None),
+    )
+
+    merged_spec = VolumeSpec(
+        size=new_spec.size
+        if new_spec and new_spec.size
+        else (current_volume.spec.size if current_volume.spec else None),
+        region=new_spec.region
+        if new_spec and new_spec.region
+        else (current_volume.spec.region if current_volume.spec else None),
+        template=new_spec.template
+        if new_spec and new_spec.template
+        else (current_volume.spec.template if current_volume.spec else None),
+    )
+
+    body = Volume(
+        metadata=merged_metadata,
+        spec=merged_spec,
+    )
+
+    response = await update_volume(volume_name=volume_name, client=client, body=body)
+    if isinstance(response, Error):
+        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)
+
+
+def _update_volume_by_name_sync(
+    volume_name: str, updates: Union[VolumeCreateConfiguration, Volume, Dict[str, any]]
+) -> "SyncVolumeInstance":
+    """Update a volume by name (sync)."""
+    # Get the current volume
+    volume_instance = SyncVolumeInstance.get(volume_name)
+    current_volume = volume_instance.volume
+
+    # Build the update body
+    if isinstance(updates, Volume):
+        new_metadata = updates.metadata
+        new_spec = updates.spec
+    elif isinstance(updates, VolumeCreateConfiguration):
+        new_metadata = Metadata(
+            name=current_volume.metadata.name if current_volume.metadata else volume_name,
+            display_name=updates.display_name,
+            labels=updates.labels,
+        )
+        new_spec = VolumeSpec(
+            size=updates.size,
+            region=updates.region,
+            template=updates.template,
+        )
+    elif isinstance(updates, dict):
+        config = VolumeCreateConfiguration.from_dict(updates)
+        new_metadata = Metadata(
+            name=current_volume.metadata.name if current_volume.metadata else volume_name,
+            display_name=config.display_name,
+            labels=config.labels,
+        )
+        new_spec = VolumeSpec(
+            size=config.size,
+            region=config.region,
+            template=config.template,
+        )
+    else:
+        raise ValueError(
+            f"Invalid updates type: {type(updates)}. Expected VolumeCreateConfiguration, Volume, or dict."
+        )
+
+    # Merge current values with updates
+    merged_metadata = Metadata(
+        name=current_volume.metadata.name if current_volume.metadata else volume_name,
+        display_name=new_metadata.display_name
+        if new_metadata and new_metadata.display_name
+        else (current_volume.metadata.display_name if current_volume.metadata else None),
+        labels=new_metadata.labels
+        if new_metadata and new_metadata.labels
+        else (current_volume.metadata.labels if current_volume.metadata else None),
+    )
+
+    merged_spec = VolumeSpec(
+        size=new_spec.size
+        if new_spec and new_spec.size
+        else (current_volume.spec.size if current_volume.spec else None),
+        region=new_spec.region
+        if new_spec and new_spec.region
+        else (current_volume.spec.region if current_volume.spec else None),
+        template=new_spec.template
+        if new_spec and new_spec.template
+        else (current_volume.spec.template if current_volume.spec else None),
+    )
+
+    body = Volume(
+        metadata=merged_metadata,
+        spec=merged_spec,
+    )
+
+    response = update_volume_sync(volume_name=volume_name, client=client, body=body)
+    if isinstance(response, Error):
+        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)
+
+
 # Assign the delete descriptors to support both class-level and instance-level calls
 VolumeInstance.delete = _AsyncDeleteDescriptor(_delete_volume_by_name)
 SyncVolumeInstance.delete = _SyncDeleteDescriptor(_delete_volume_by_name_sync)
+
+# Assign the update descriptors to support both class-level and instance-level calls
+VolumeInstance.update = _AsyncUpdateDescriptor(_update_volume_by_name)
+SyncVolumeInstance.update = _SyncUpdateDescriptor(_update_volume_by_name_sync)

blaxel/langgraph/model.py

@@ -1,22 +1,17 @@
+from __future__ import annotations
+
 from logging import getLogger
-from typing import Any, AsyncIterator, Iterator, List
-
-from langchain_anthropic import ChatAnthropic
-from langchain_cerebras import ChatCerebras
-from langchain_cohere import ChatCohere
-from langchain_core.callbacks import Callbacks
-from langchain_core.language_models import LanguageModelInput
-from langchain_core.messages import BaseMessage
-from langchain_core.outputs import LLMResult
-from langchain_core.runnables import RunnableConfig
-from langchain_deepseek import ChatDeepSeek
-from langchain_openai import ChatOpenAI
-from langchain_xai import ChatXAI
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
 
-from .custom.gemini import ChatGoogleGenerativeAI
+if TYPE_CHECKING:
+    from langchain_core.callbacks import Callbacks
+    from langchain_core.language_models import LanguageModelInput
+    from langchain_core.messages import BaseMessage
+    from langchain_core.outputs import LLMResult
+    from langchain_core.runnables import RunnableConfig
 
 logger = getLogger(__name__)
 
@@ -37,6 +32,8 @@ class TokenRefreshingWrapper:
         kwargs = config.get("kwargs", {})
 
         if model_type == "mistral":
+            from langchain_openai import ChatOpenAI
+
             return ChatOpenAI(
                 api_key=settings.auth.token,
                 model=model,
@@ -44,6 +41,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "cohere":
+            from langchain_cohere import ChatCohere
+
             return ChatCohere(
                 cohere_api_key=settings.auth.token,
                 model=model,
@@ -51,6 +50,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "xai":
+            from langchain_xai import ChatXAI
+
             return ChatXAI(
                 model=model,
                 api_key=settings.auth.token,
@@ -58,6 +59,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "deepseek":
+            from langchain_deepseek import ChatDeepSeek
+
             return ChatDeepSeek(
                 api_key=settings.auth.token,
                 model=model,
@@ -65,6 +68,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "anthropic":
+            from langchain_anthropic import ChatAnthropic
+
             return ChatAnthropic(
                 api_key=settings.auth.token,
                 anthropic_api_url=url,
@@ -73,6 +78,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "gemini":
+            from .custom.gemini import ChatGoogleGenerativeAI
+
             return ChatGoogleGenerativeAI(
                 model=model,
                 client_options={"api_endpoint": url},
@@ -81,6 +88,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "cerebras":
+            from langchain_cerebras import ChatCerebras
+
             return ChatCerebras(
                 api_key=settings.auth.token,
                 model=model,
@@ -88,6 +97,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         else:
+            from langchain_openai import ChatOpenAI
+
             if model_type != "openai":
                 logger.warning(f"Model {model} is not supported by Langchain, defaulting to OpenAI")
             return ChatOpenAI(

blaxel/langgraph/tools.py

@@ -1,17 +1,10 @@
-from typing import Any, Dict
-
-from langchain_core.tools import StructuredTool
-from mcp.types import (
-    CallToolResult,
-    EmbeddedResource,
-    ImageContent,
-    TextContent,
-)
+from typing import TYPE_CHECKING, Any, Dict
 
 from blaxel.core.tools import bl_tools as bl_tools_core
 from blaxel.core.tools.types import Tool, ToolException
 
-NonTextContent = ImageContent | EmbeddedResource
+if TYPE_CHECKING:
+    from langchain_core.tools import StructuredTool
 
 
 def _clean_schema_for_openai(schema: Dict[str, Any]) -> Dict[str, Any]:
@@ -43,7 +36,17 @@ def _clean_schema_for_openai(schema: Dict[str, Any]) -> Dict[str, Any]:
     return cleaned
 
 
-def get_langchain_tool(tool: Tool) -> StructuredTool:
+def get_langchain_tool(tool: Tool) -> "StructuredTool":
+    from langchain_core.tools import StructuredTool
+    from mcp.types import (
+        CallToolResult,
+        EmbeddedResource,
+        ImageContent,
+        TextContent,
+    )
+
+    NonTextContent = ImageContent | EmbeddedResource
+
     async def langchain_coroutine(
         **arguments: dict[str, Any],
     ) -> tuple[str | list[str], list[NonTextContent] | None]:
@@ -77,7 +80,7 @@ def get_langchain_tool(tool: Tool) -> StructuredTool:
     )
 
 
-async def bl_tools(tools_names: list[str], **kwargs) -> list[StructuredTool]:
+async def bl_tools(tools_names: list[str], **kwargs) -> list["StructuredTool"]:
     tools = bl_tools_core(tools_names, **kwargs)
     await tools.initialize()
     return [get_langchain_tool(tool) for tool in tools.get_tools()]

blaxel/llamaindex/model.py

@@ -1,33 +1,25 @@
-import os
-
-# Transformers is a dependency of DeepSeek, and it logs a lot of warnings that are not useful
-os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
+from __future__ import annotations
 
+import os
 from logging import getLogger
-from typing import Any, Sequence
-
-from google.genai.types import HttpOptions
-from llama_index.core.base.llms.types import (
-    ChatMessage,
-    ChatResponse,
-    ChatResponseAsyncGen,
-    ChatResponseGen,
-    CompletionResponse,
-    CompletionResponseAsyncGen,
-    CompletionResponseGen,
-)
-from llama_index.llms.anthropic import Anthropic
-from llama_index.llms.cerebras import Cerebras
-from llama_index.llms.deepseek import DeepSeek
-from llama_index.llms.google_genai import GoogleGenAI
-from llama_index.llms.groq import Groq
-from llama_index.llms.mistralai import MistralAI
-from llama_index.llms.openai import OpenAI
+from typing import TYPE_CHECKING, Any, Sequence
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
 
-from .custom.cohere import Cohere
+# Transformers is a dependency of DeepSeek, and it logs a lot of warnings that are not useful
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
+
+if TYPE_CHECKING:
+    from llama_index.core.base.llms.types import (
+        ChatMessage,
+        ChatResponse,
+        ChatResponseAsyncGen,
+        ChatResponseGen,
+        CompletionResponse,
+        CompletionResponseAsyncGen,
+        CompletionResponseGen,
+    )
 
 logger = getLogger(__name__)
 
@@ -48,6 +40,8 @@ class TokenRefreshingWrapper:
         kwargs = config.get("kwargs", {})
 
         if model_type == "anthropic":
+            from llama_index.llms.anthropic import Anthropic
+
             return Anthropic(
                 model=model,
                 api_key=settings.auth.token,
@@ -56,6 +50,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "xai":
+            from llama_index.llms.groq import Groq
+
             return Groq(
                 model=model,
                 api_key=settings.auth.token,
@@ -63,6 +59,9 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "gemini":
+            from google.genai.types import HttpOptions
+            from llama_index.llms.google_genai import GoogleGenAI
+
             return GoogleGenAI(
                 api_key=settings.auth.token,
                 model=model,
@@ -74,8 +73,12 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "cohere":
+            from .custom.cohere import Cohere
+
             return Cohere(model=model, api_key=settings.auth.token, api_base=url, **kwargs)
         elif model_type == "deepseek":
+            from llama_index.llms.deepseek import DeepSeek
+
             return DeepSeek(
                 model=model,
                 api_key=settings.auth.token,
@@ -83,8 +86,12 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "mistral":
+            from llama_index.llms.mistralai import MistralAI
+
             return MistralAI(model=model, api_key=settings.auth.token, endpoint=url, **kwargs)
         elif model_type == "cerebras":
+            from llama_index.llms.cerebras import Cerebras
+
             return Cerebras(
                 model=model,
                 api_key=settings.auth.token,
@@ -92,6 +99,8 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         else:
+            from llama_index.llms.openai import OpenAI
+
             if model_type != "openai":
                 logger.warning(
                     f"Model {model} is not supported by LlamaIndex, defaulting to OpenAI"

blaxel/llamaindex/tools.py

@@ -1,12 +1,17 @@
-from llama_index.core.tools import FunctionTool
-from llama_index.core.tools.types import ToolMetadata
+from typing import TYPE_CHECKING
 
 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
 
+if TYPE_CHECKING:
+    from llama_index.core.tools import FunctionTool
+
+
+def get_llamaindex_tool(tool: Tool) -> "FunctionTool":
+    from llama_index.core.tools import FunctionTool
+    from llama_index.core.tools.types import ToolMetadata
 
-def get_llamaindex_tool(tool: Tool) -> FunctionTool:
     model_schema = create_model_from_json_schema(
         tool.input_schema, model_name=f"{tool.name}_Schema"
     )
@@ -21,7 +26,7 @@ def get_llamaindex_tool(tool: Tool) -> FunctionTool:
     )
 
 
-async def bl_tools(tools_names: list[str], **kwargs) -> list[FunctionTool]:
+async def bl_tools(tools_names: list[str], **kwargs) -> list["FunctionTool"]:
     tools = bl_tools_core(tools_names, **kwargs)
     await tools.initialize()
     return [get_llamaindex_tool(tool) for tool in tools.get_tools()]

blaxel/pydantic/model.py

@@ -1,19 +1,7 @@
 import logging
 from typing import Any
 
-from anthropic import AsyncAnthropic
-from cohere import AsyncClientV2
-from mistralai.sdk import Mistral
 from pydantic_ai.models import Model
-from pydantic_ai.models.anthropic import AnthropicModel
-from pydantic_ai.models.cohere import CohereModel
-from pydantic_ai.models.gemini import GeminiModel
-from pydantic_ai.models.mistral import MistralModel
-from pydantic_ai.models.openai import OpenAIModel
-from pydantic_ai.providers.anthropic import AnthropicProvider
-from pydantic_ai.providers.cohere import CohereProvider
-from pydantic_ai.providers.mistral import MistralProvider
-from pydantic_ai.providers.openai import OpenAIProvider
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -41,6 +29,10 @@ class TokenRefreshingModel(Model):
         kwargs = config.get("kwargs", {})
 
         if type == "mistral":
+            from mistralai.sdk import Mistral
+            from pydantic_ai.models.mistral import MistralModel
+            from pydantic_ai.providers.mistral import MistralProvider
+
             return MistralModel(
                 model_name=model,
                 provider=MistralProvider(
@@ -52,6 +44,10 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "cohere":
+            from cohere import AsyncClientV2
+            from pydantic_ai.models.cohere import CohereModel
+            from pydantic_ai.providers.cohere import CohereProvider
+
             return CohereModel(
                 model_name=model,
                 provider=CohereProvider(
@@ -62,6 +58,9 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "xai":
+            from pydantic_ai.models.openai import OpenAIModel
+            from pydantic_ai.providers.openai import OpenAIProvider
+
             return OpenAIModel(
                 model_name=model,
                 provider=OpenAIProvider(
@@ -69,6 +68,9 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "deepseek":
+            from pydantic_ai.models.openai import OpenAIModel
+            from pydantic_ai.providers.openai import OpenAIProvider
+
             return OpenAIModel(
                 model_name=model,
                 provider=OpenAIProvider(
@@ -76,6 +78,9 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "cerebras":
+            from pydantic_ai.models.openai import OpenAIModel
+            from pydantic_ai.providers.openai import OpenAIProvider
+
             return OpenAIModel(
                 model_name=model,
                 provider=OpenAIProvider(
@@ -83,6 +88,10 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "anthropic":
+            from anthropic import AsyncAnthropic
+            from pydantic_ai.models.anthropic import AnthropicModel
+            from pydantic_ai.providers.anthropic import AnthropicProvider
+
             return AnthropicModel(
                 model_name=model,
                 provider=AnthropicProvider(
@@ -95,6 +104,8 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "gemini":
+            from pydantic_ai.models.gemini import GeminiModel
+
             return GeminiModel(
                 model_name=model,
                 provider=GoogleGLAProvider(
@@ -105,6 +116,9 @@ class TokenRefreshingModel(Model):
                 ),
             )
         else:
+            from pydantic_ai.models.openai import OpenAIModel
+            from pydantic_ai.providers.openai import OpenAIProvider
+
             if type != "openai":
                 logger.warning(f"Model {model} is not supported by Pydantic, defaulting to OpenAI")
             return OpenAIModel(