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

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
@@ -11,7 +11,7 @@ from blaxel.core import settings
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 
 if TYPE_CHECKING:
-    from llama_index.core.base.llms.types import (
+    from llama_index.core.base.llms.types import (  # type: ignore[import-not-found]
         ChatMessage,
         ChatResponse,
         ChatResponseAsyncGen,
@@ -20,27 +20,76 @@ if TYPE_CHECKING:
         CompletionResponseAsyncGen,
         CompletionResponseGen,
     )
+    from llama_index.core.llms.llm import (  # type: ignore[import-not-found]
+        ToolSelection,
+    )
+    from llama_index.core.tools.types import BaseTool  # type: ignore[import-not-found]
+
+# Runtime imports needed for class inheritance and construction
+from llama_index.core.base.llms.types import (  # type: ignore[import-not-found]
+    LLMMetadata,
+)
+from llama_index.core.llms.function_calling import (  # type: ignore[import-not-found]
+    FunctionCallingLLM,
+)
+from pydantic import PrivateAttr  # type: ignore[import-not-found]
 
 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"]
         kwargs = config.get("kwargs", {})
 
         if model_type == "anthropic":
-            from llama_index.llms.anthropic import Anthropic
+            from llama_index.llms.anthropic import (  # type: ignore[import-not-found]
+                Anthropic,
+            )
 
             return Anthropic(
                 model=model,
@@ -50,7 +99,7 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "xai":
-            from llama_index.llms.groq import Groq
+            from llama_index.llms.groq import Groq  # type: ignore[import-not-found]
 
             return Groq(
                 model=model,
@@ -60,7 +109,9 @@ class TokenRefreshingWrapper:
             )
         elif model_type == "gemini":
             from google.genai.types import HttpOptions
-            from llama_index.llms.google_genai import GoogleGenAI
+            from llama_index.llms.google_genai import (  # type: ignore[import-not-found]
+                GoogleGenAI,
+            )
 
             return GoogleGenAI(
                 api_key=settings.auth.token,
@@ -73,11 +124,13 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "cohere":
-            from .custom.cohere import Cohere
+            from .custom.cohere import Cohere  # type: ignore[import-not-found]
 
             return Cohere(model=model, api_key=settings.auth.token, api_base=url, **kwargs)
         elif model_type == "deepseek":
-            from llama_index.llms.deepseek import DeepSeek
+            from llama_index.llms.deepseek import (  # type: ignore[import-not-found]
+                DeepSeek,
+            )
 
             return DeepSeek(
                 model=model,
@@ -86,11 +139,15 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         elif model_type == "mistral":
-            from llama_index.llms.mistralai import MistralAI
+            from llama_index.llms.mistralai import (  # type: ignore[import-not-found]
+                MistralAI,
+            )
 
             return MistralAI(model=model, api_key=settings.auth.token, endpoint=url, **kwargs)
         elif model_type == "cerebras":
-            from llama_index.llms.cerebras import Cerebras
+            from llama_index.llms.cerebras import (  # type: ignore[import-not-found]
+                Cerebras,
+            )
 
             return Cerebras(
                 model=model,
@@ -99,7 +156,7 @@ class TokenRefreshingWrapper:
                 **kwargs,
             )
         else:
-            from llama_index.llms.openai import OpenAI
+            from llama_index.llms.openai import OpenAI  # type: ignore[import-not-found]
 
             if model_type != "openai":
                 logger.warning(
@@ -115,102 +172,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()
-
-    def __getattr__(self, name):
-        """Delegate attribute access to wrapped model."""
-        return getattr(self.wrapped_model, name)
+            self._wrapped = self._create_model()
 
+    # --- Core LLM methods with token refresh ---
 
-class TokenRefreshingLLM(TokenRefreshingWrapper):
-    """Wrapper for LlamaIndex LLMs that refreshes token before each call."""
+    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:
-        """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):
@@ -220,4 +281,4 @@ async def bl_model(name, **kwargs):
     model_config = {"type": type, "model": model, "url": url, "kwargs": kwargs}
 
     # Create and return the wrapper
-    return TokenRefreshingLLM(model_config)
+    return TokenRefreshingLLM(model_config)

blaxel/llamaindex/tools.py

@@ -5,12 +5,14 @@ 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
+    from llama_index.core.tools import FunctionTool  # type: ignore[import-not-found]
 
 
 def get_llamaindex_tool(tool: Tool) -> "FunctionTool":
-    from llama_index.core.tools import FunctionTool
-    from llama_index.core.tools.types import ToolMetadata
+    from llama_index.core.tools import FunctionTool  # type: ignore[import-not-found]
+    from llama_index.core.tools.types import (  # type: ignore[import-not-found]
+        ToolMetadata,
+    )
 
     model_schema = create_model_from_json_schema(
         tool.input_schema, model_name=f"{tool.name}_Schema"
@@ -29,4 +31,4 @@ def get_llamaindex_tool(tool: Tool) -> "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()]
+    return [get_llamaindex_tool(tool) for tool in tools.get_tools()]

blaxel/openai/model.py

@@ -1,5 +1,6 @@
 import httpx
-from agents import AsyncOpenAI, OpenAIChatCompletionsModel
+from agents import OpenAIChatCompletionsModel
+from openai import AsyncOpenAI
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -14,6 +15,11 @@ class DynamicHeadersHTTPClient(httpx.AsyncClient):
     async def send(self, request, *args, **kwargs):
         # Update headers with the latest auth headers before each request
         auth_headers = settings.auth.get_headers()
+        # Remove the SDK's default "Authorization: Bearer replaced" header
+        # when our auth uses a different header (e.g. X-Blaxel-Authorization with API keys)
+        if "Authorization" not in auth_headers:
+            request.headers.pop("Authorization", None)
+            request.headers.pop("authorization", None)
         for key, value in auth_headers.items():
             request.headers[key] = value
         return await super().send(request, *args, **kwargs)

blaxel/openai/tools.py

@@ -1,7 +1,8 @@
 import json
 from typing import Any
 
-from agents import FunctionTool, RunContextWrapper
+from agents import FunctionTool  # type: ignore[import-not-found]
+from agents.tool_context import ToolContext  # type: ignore[import-not-found]
 
 from blaxel.core.tools import bl_tools as bl_tools_core
 from blaxel.core.tools.types import Tool
@@ -24,6 +25,13 @@ def _clean_schema_for_openai(schema: dict) -> dict:
     if "additionalProperties" in cleaned_schema:
         del cleaned_schema["additionalProperties"]
 
+    # Ensure object type schemas have properties
+    if cleaned_schema.get("type") == "object":
+        if "properties" not in cleaned_schema:
+            cleaned_schema["properties"] = {}
+        if "required" not in cleaned_schema:
+            cleaned_schema["required"] = []
+
     # Recursively clean properties if they exist
     if "properties" in cleaned_schema:
         cleaned_schema["properties"] = {
@@ -39,9 +47,11 @@ def _clean_schema_for_openai(schema: dict) -> dict:
 
 def get_openai_tool(tool: Tool) -> FunctionTool:
     async def openai_coroutine(
-        _: RunContextWrapper,
-        arguments: dict[str, Any],
+        _: ToolContext[Any],
+        arguments: str,
     ) -> Any:
+        if not tool.coroutine:
+            raise ValueError(f"Tool {tool.name} does not have a coroutine defined")
         result = await tool.coroutine(**json.loads(arguments))
         return result
 

blaxel/pydantic/model.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any
 
-from pydantic_ai.models import Model
+from pydantic_ai.models import Model  # type: ignore[import-not-found]
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -30,8 +30,12 @@ class TokenRefreshingModel(Model):
 
         if type == "mistral":
             from mistralai.sdk import Mistral
-            from pydantic_ai.models.mistral import MistralModel
-            from pydantic_ai.providers.mistral import MistralProvider
+            from pydantic_ai.models.mistral import (  # type: ignore[import-not-found]
+                MistralModel,
+            )
+            from pydantic_ai.providers.mistral import (  # type: ignore[import-not-found]
+                MistralProvider,
+            )
 
             return MistralModel(
                 model_name=model,
@@ -45,8 +49,12 @@ class TokenRefreshingModel(Model):
             )
         elif type == "cohere":
             from cohere import AsyncClientV2
-            from pydantic_ai.models.cohere import CohereModel
-            from pydantic_ai.providers.cohere import CohereProvider
+            from pydantic_ai.models.cohere import (  # type: ignore[import-not-found]
+                CohereModel,
+            )
+            from pydantic_ai.providers.cohere import (  # type: ignore[import-not-found]
+                CohereProvider,
+            )
 
             return CohereModel(
                 model_name=model,
@@ -58,30 +66,42 @@ class TokenRefreshingModel(Model):
                 ),
             )
         elif type == "xai":
-            from pydantic_ai.models.openai import OpenAIModel
-            from pydantic_ai.providers.openai import OpenAIProvider
+            from pydantic_ai.models.openai import (  # type: ignore[import-not-found]
+                OpenAIChatModel,
+            )
+            from pydantic_ai.providers.openai import (  # type: ignore[import-not-found]
+                OpenAIProvider,
+            )
 
-            return OpenAIModel(
+            return OpenAIChatModel(
                 model_name=model,
                 provider=OpenAIProvider(
                     base_url=f"{url}/v1", api_key=settings.auth.token, **kwargs
                 ),
             )
         elif type == "deepseek":
-            from pydantic_ai.models.openai import OpenAIModel
-            from pydantic_ai.providers.openai import OpenAIProvider
+            from pydantic_ai.models.openai import (  # type: ignore[import-not-found]
+                OpenAIChatModel,
+            )
+            from pydantic_ai.providers.openai import (  # type: ignore[import-not-found]
+                OpenAIProvider,
+            )
 
-            return OpenAIModel(
+            return OpenAIChatModel(
                 model_name=model,
                 provider=OpenAIProvider(
                     base_url=f"{url}/v1", api_key=settings.auth.token, **kwargs
                 ),
             )
         elif type == "cerebras":
-            from pydantic_ai.models.openai import OpenAIModel
-            from pydantic_ai.providers.openai import OpenAIProvider
+            from pydantic_ai.models.openai import (  # type: ignore[import-not-found]
+                OpenAIChatModel,
+            )
+            from pydantic_ai.providers.openai import (  # type: ignore[import-not-found]
+                OpenAIProvider,
+            )
 
-            return OpenAIModel(
+            return OpenAIChatModel(
                 model_name=model,
                 provider=OpenAIProvider(
                     base_url=f"{url}/v1", api_key=settings.auth.token, **kwargs
@@ -116,12 +136,12 @@ class TokenRefreshingModel(Model):
                 ),
             )
         else:
-            from pydantic_ai.models.openai import OpenAIModel
+            from pydantic_ai.models.openai import OpenAIChatModel
             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(
+            return OpenAIChatModel(
                 model_name=model,
                 provider=OpenAIProvider(
                     base_url=f"{url}/v1", api_key=settings.auth.token, **kwargs
@@ -130,12 +150,6 @@ class TokenRefreshingModel(Model):
 
     def _get_fresh_model(self) -> Model:
         """Get or create a model with fresh token if needed."""
-        # Only refresh if using ClientCredentials (which has get_token method)
-        if hasattr(settings.auth, "get_token"):
-            # This will trigger token refresh if needed
-            logger.debug(f"Calling get_token for {self.model_config['type']} model")
-            settings.auth.get_token()
-
         new_token = settings.auth.token
 
         # If token changed or no cached model, create new one
@@ -152,10 +166,10 @@ class TokenRefreshingModel(Model):
         return model.model_name
 
     @property
-    def system(self) -> Any | None:
+    def system(self) -> str:
         """Return the system property from the wrapped model."""
         model = self._get_fresh_model()
-        return model.system if hasattr(model, "system") else None
+        return model.system if hasattr(model, "system") else ""
 
     async def request(self, *args, **kwargs):
         """Make a request to the model with token refresh."""

blaxel/pydantic/tools.py

@@ -1,11 +1,44 @@
-from pydantic_ai import RunContext
-from pydantic_ai.tools import Tool as PydanticTool
-from pydantic_ai.tools import ToolDefinition
+from typing import Any
+
+from pydantic_ai import RunContext  # type: ignore[import-not-found]
+from pydantic_ai.tools import Tool as PydanticTool  # type: ignore[import-not-found]
+from pydantic_ai.tools import ToolDefinition  # type: ignore[import-not-found]
 
 from blaxel.core.tools import Tool
 from blaxel.core.tools import bl_tools as bl_tools_core
 
 
+def _clean_schema_for_openai(schema: dict[str, Any]) -> dict[str, Any]:
+    """Clean JSON schema to be compatible with OpenAI function calling.
+
+    OpenAI requires object schemas to have a 'properties' field, even if empty.
+    """
+    if not isinstance(schema, dict):
+        return schema
+
+    cleaned = schema.copy()
+
+    if cleaned.get("type") == "object":
+        if "properties" not in cleaned:
+            cleaned["properties"] = {}
+        if "required" not in cleaned:
+            cleaned["required"] = []
+
+    if "additionalProperties" in cleaned:
+        del cleaned["additionalProperties"]
+    if "$schema" in cleaned:
+        del cleaned["$schema"]
+
+    if "properties" in cleaned:
+        cleaned["properties"] = {
+            k: _clean_schema_for_openai(v) for k, v in cleaned["properties"].items()
+        }
+    if "items" in cleaned and isinstance(cleaned["items"], dict):
+        cleaned["items"] = _clean_schema_for_openai(cleaned["items"])
+
+    return cleaned
+
+
 def get_pydantic_tool(tool: Tool) -> PydanticTool:
     """
     Converts a custom Tool object into a Pydantic AI Tool object.
@@ -27,7 +60,7 @@ def get_pydantic_tool(tool: Tool) -> PydanticTool:
         """Dynamically prepares the ToolDefinition using the custom Tool's attributes."""
         tool_def.name = tool.name  # Override inferred name
         tool_def.description = tool.description  # Override inferred description
-        tool_def.parameters_json_schema = tool.input_schema
+        tool_def.parameters_json_schema = _clean_schema_for_openai(tool.input_schema)
         return tool_def
 
     async def pydantic_function(**kwargs):