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

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):