blaxel 0.2.37__tar.gz → 0.2.38rc122__tar.gz

{blaxel-0.2.37 → blaxel-0.2.38rc122}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.2.37
+Version: 0.2.38rc122
 Summary: Blaxel - AI development platform SDK
 Project-URL: Homepage, https://blaxel.ai
 Project-URL: Documentation, https://docs.blaxel.ai

{blaxel-0.2.37 → blaxel-0.2.38rc122}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blaxel"
-version = "0.2.37"
+version = "0.2.38.pre-122"
 description = "Blaxel - AI development platform SDK"
 readme = "README.md"
 authors = [{ name = "cploujoux", email = "cploujoux@blaxel.ai" }]

{blaxel-0.2.37 → blaxel-0.2.38rc122}/src/blaxel/__init__.py

@@ -4,8 +4,8 @@ from .core.common.autoload import autoload
 from .core.common.env import env
 from .core.common.settings import settings
 
-__version__ = "0.2.37"
-__commit__ = "a09d7ea6c786074ceaa4639ce6415958524a7a5f"
+__version__ = "0.2.38.pre-122"
+__commit__ = "e7e19d2c82b2b0c2a44914245734597adcc00f49"
 __sentry_dsn__ = "https://9711de13cd02b285ca4378c01de8dc30@o4508714045276160.ingest.us.sentry.io/4510461121462272"
 __all__ = ["autoload", "settings", "env"]
 

blaxel-0.2.38rc122/src/blaxel/llamaindex/model.py

@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+import os
+from logging import getLogger
+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
+
+# 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,
+    )
+    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.
+
+    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):
+        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_data
+        model_type = config["type"]
+        model = config["model"]
+        url = config["url"]
+        kwargs = config.get("kwargs", {})
+
+        if model_type == "anthropic":
+            from llama_index.llms.anthropic import Anthropic
+
+            return Anthropic(
+                model=model,
+                api_key=settings.auth.token,
+                base_url=url,
+                default_headers=settings.auth.get_headers(),
+                **kwargs,
+            )
+        elif model_type == "xai":
+            from llama_index.llms.groq import Groq
+
+            return Groq(
+                model=model,
+                api_key=settings.auth.token,
+                api_base=f"{url}/v1",
+                **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,
+                api_base=f"{url}/v1",
+                http_options=HttpOptions(
+                    base_url=url,
+                    headers=settings.auth.get_headers(),
+                ),
+                **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,
+                api_base=f"{url}/v1",
+                **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,
+                api_base=f"{url}/v1",
+                **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"
+                )
+
+            return OpenAI(
+                model=model,
+                api_key=settings.auth.token,
+                api_base=f"{url}/v1",
+                **kwargs,
+            )
+
+    def _refresh_token(self):
+        """Refresh the token and recreate the model if needed."""
+        current_token = settings.auth.token
+
+        new_token = settings.auth.token
+
+        if current_token != new_token:
+            self._wrapped = self._create_model()
+
+    # --- Core LLM methods with token refresh ---
+
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        self._refresh_token()
+        return self._wrapped.chat(messages, **kwargs)
+
+    async def achat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        self._refresh_token()
+        return await self._wrapped.achat(messages, **kwargs)
+
+    def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
+        self._refresh_token()
+        return self._wrapped.complete(prompt, formatted=formatted, **kwargs)
+
+    async def acomplete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
+        self._refresh_token()
+        return await self._wrapped.acomplete(prompt, formatted=formatted, **kwargs)
+
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        self._refresh_token()
+        return self._wrapped.stream_chat(messages, **kwargs)
+
+    async def astream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseAsyncGen:
+        self._refresh_token()
+        result = self._wrapped.astream_chat(messages, **kwargs)
+        # Handle both coroutine and async generator patterns
+        if hasattr(result, "__aiter__"):
+            return result
+        return await result
+
+    def stream_complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponseGen:
+        self._refresh_token()
+        return self._wrapped.stream_complete(prompt, formatted=formatted, **kwargs)
+
+    async def astream_complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponseAsyncGen:
+        self._refresh_token()
+        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
+
+    # --- FunctionCallingLLM methods (delegate to wrapped model) ---
+
+    def _prepare_chat_with_tools(
+        self,
+        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,
+    ) -> 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):
+    url, type, model = await bl_model_core(name).get_parameters()
+
+    # Store model configuration for recreation
+    model_config = {"type": type, "model": model, "url": url, "kwargs": kwargs}
+
+    # Create and return the wrapper
+    return TokenRefreshingLLM(model_config)

blaxel-0.2.37/src/blaxel/llamaindex/model.py

@@ -1,223 +0,0 @@
-from __future__ import annotations
-
-import os
-from logging import getLogger
-from typing import TYPE_CHECKING, Any, Sequence
-
-from blaxel.core import bl_model as bl_model_core
-from blaxel.core import settings
-
-# 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__)
-
-
-class TokenRefreshingWrapper:
-    """Base wrapper class that refreshes token before each call."""
-
-    def __init__(self, model_config: dict):
-        self.model_config = model_config
-        self.wrapped_model = self._create_model()
-
-    def _create_model(self):
-        """Create the model instance with current token."""
-        config = self.model_config
-        model_type = config["type"]
-        model = config["model"]
-        url = config["url"]
-        kwargs = config.get("kwargs", {})
-
-        if model_type == "anthropic":
-            from llama_index.llms.anthropic import Anthropic
-
-            return Anthropic(
-                model=model,
-                api_key=settings.auth.token,
-                base_url=url,
-                default_headers=settings.auth.get_headers(),
-                **kwargs,
-            )
-        elif model_type == "xai":
-            from llama_index.llms.groq import Groq
-
-            return Groq(
-                model=model,
-                api_key=settings.auth.token,
-                api_base=f"{url}/v1",
-                **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,
-                api_base=f"{url}/v1",
-                http_options=HttpOptions(
-                    base_url=url,
-                    headers=settings.auth.get_headers(),
-                ),
-                **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,
-                api_base=f"{url}/v1",
-                **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,
-                api_base=f"{url}/v1",
-                **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"
-                )
-
-            return OpenAI(
-                model=model,
-                api_key=settings.auth.token,
-                api_base=f"{url}/v1",
-                **kwargs,
-            )
-
-    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()
-
-    def __getattr__(self, name):
-        """Delegate attribute access to wrapped model."""
-        return getattr(self.wrapped_model, name)
-
-
-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."""
-        self._refresh_token()
-        return await self.wrapped_model.achat(messages, **kwargs)
-
-    def chat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
-    ) -> ChatResponse:
-        """Sync chat with token refresh."""
-        self._refresh_token()
-        return self.wrapped_model.chat(messages, **kwargs)
-
-    async def astream_chat(
-        self,
-        messages: Sequence[ChatMessage],
-        **kwargs: Any,
-    ) -> ChatResponseAsyncGen:
-        """Async stream chat with token refresh."""
-        self._refresh_token()
-        async for chunk in self.wrapped_model.astream_chat(messages, **kwargs):
-            yield chunk
-
-    def stream_chat(
-        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
-
-    async def acomplete(
-        self,
-        prompt: str,
-        **kwargs: Any,
-    ) -> CompletionResponse:
-        """Async complete with token refresh."""
-        self._refresh_token()
-        return await self.wrapped_model.acomplete(prompt, **kwargs)
-
-    def complete(
-        self,
-        prompt: str,
-        **kwargs: Any,
-    ) -> CompletionResponse:
-        """Sync complete with token refresh."""
-        self._refresh_token()
-        return self.wrapped_model.complete(prompt, **kwargs)
-
-    async def astream_complete(
-        self,
-        prompt: str,
-        **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
-
-    def stream_complete(
-        self,
-        prompt: str,
-        **kwargs: Any,
-    ) -> CompletionResponseGen:
-        """Sync stream complete with token refresh."""
-        self._refresh_token()
-        for chunk in self.wrapped_model.stream_complete(prompt, **kwargs):
-            yield chunk
-
-
-async def bl_model(name, **kwargs):
-    url, type, model = await bl_model_core(name).get_parameters()
-
-    # Store model configuration for recreation
-    model_config = {"type": type, "model": model, "url": url, "kwargs": kwargs}
-
-    # Create and return the wrapper
-    return TokenRefreshingLLM(model_config)