blaxel 0.1.22rc70__py3-none-any.whl → 0.2.0rc2__py3-none-any.whl

blaxel/{tools → core/tools}/__init__.py

@@ -1,5 +1,4 @@
 import asyncio
-import json
 import os
 import traceback
 from contextlib import AsyncExitStack
@@ -12,7 +11,6 @@ from mcp.types import Tool as MCPTool
 
 from ..common.internal import get_forced_url, get_global_unique_hash
 from ..common.settings import settings
-from ..instrumentation.span import SpanManager
 from ..mcp.client import websocket_client
 from .types import Tool
 
@@ -22,8 +20,11 @@ DEFAULT_TIMEOUT = 1
 if os.getenv("BL_SERVER_PORT"):
     DEFAULT_TIMEOUT = 5
 
+
 class PersistentWebSocket:
-    def __init__(self, url: str, name: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
+    def __init__(
+        self, url: str, name: str, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True
+    ):
         self.url = url
         self.name = name
         self.timeout = timeout
@@ -58,32 +59,30 @@ class PersistentWebSocket:
         except Exception as e:
             logger.error(f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}")
             return CallToolResult(
-                content=[{"type": "text", "text": f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}"}],
+                content=[
+                    {
+                        "type": "text",
+                        "text": f"Error calling tool {tool_name}: {e}\n{traceback.format_exc()}",
+                    }
+                ],
                 isError=True,
             )
 
     async def list_tools(self):
         logger.debug(f"Listing tools for {self.name}")
-        span_attributes = {
-            "tool.server": self.url,
-            "tool.server_name": self.name,
-            "span.type": "tool.list",
-        }
-        with SpanManager("blaxel-tracer").create_active_span(self.name, span_attributes) as span:
-            await self._initialize()
-            logger.debug(f"Initialized websocket for {self.name}")
-            if self.timeout_enabled:
-                self._remove_timer()
-            logger.debug("Listing tools")
-            list_tools_result = await self.session.list_tools()
-            self.tools_cache = list_tools_result.tools
-            logger.debug(f"Tools listed: {list_tools_result}")
-            if self.timeout_enabled:
-                self._reset_timer()
-            else:
-                await self._close()
-            span.set_attribute("tool.list.result", list_tools_result.model_dump_json())
-            return list_tools_result
+        await self._initialize()
+        logger.debug(f"Initialized websocket for {self.name}")
+        if self.timeout_enabled:
+            self._remove_timer()
+        logger.debug("Listing tools")
+        list_tools_result = await self.session.list_tools()
+        self.tools_cache = list_tools_result.tools
+        logger.debug(f"Tools listed: {list_tools_result}")
+        if self.timeout_enabled:
+            self._reset_timer()
+        else:
+            await self._close()
+        return list_tools_result
 
     def get_tools(self):
         return self.tools_cache
@@ -91,8 +90,13 @@ class PersistentWebSocket:
     async def _initialize(self):
         if not self.session:
             logger.debug(f"Initializing websocket client for {self.url}")
-            read, write = await self.client_exit_stack.enter_async_context(websocket_client(self.url, settings.headers))
-            self.session = cast(ClientSession, await self.session_exit_stack.enter_async_context(ClientSession(read, write)))
+            read, write = await self.client_exit_stack.enter_async_context(
+                websocket_client(self.url, settings.headers)
+            )
+            self.session = cast(
+                ClientSession,
+                await self.session_exit_stack.enter_async_context(ClientSession(read, write)),
+            )
             await self.session.initialize()
 
     def _reset_timer(self):
@@ -108,7 +112,6 @@ class PersistentWebSocket:
         await self._close()
         self.session = None
 
-
     async def _close(self):
         logger.debug(f"Closing websocket client {self.url}")
         if self.session:
@@ -123,6 +126,7 @@ class PersistentWebSocket:
                 logger.debug(f"Error closing client exit stack: {e}")
             logger.debug("WebSocket connection closed due to inactivity.")
 
+
 def convert_mcp_tool_to_blaxel_tool(
     websocket_client: PersistentWebSocket,
     name: str,
@@ -145,18 +149,10 @@ def convert_mcp_tool_to_blaxel_tool(
         *args: Any,
         **arguments: dict[str, Any],
     ) -> CallToolResult:
-        span_attributes = {
-            "tool.name": tool.name,
-            "tool.args": json.dumps(arguments),
-            "tool.server": url,
-            "tool.server_name": name,
-            "span.type": "tool.call",
-        }
-        with SpanManager("blaxel-tracer").create_active_span("blaxel-tool-call", span_attributes):
-            logger.debug(f"Calling tool {tool.name} with arguments {arguments}")
-            call_tool_result = await websocket_client.call_tool(tool.name, arguments)
-            logger.debug(f"Tool {tool.name} returned {call_tool_result}")
-            return call_tool_result
+        logger.debug(f"Calling tool {tool.name} with arguments {arguments}")
+        call_tool_result = await websocket_client.call_tool(tool.name, arguments)
+        logger.debug(f"Tool {tool.name} returned {call_tool_result}")
+        return call_tool_result
 
     async def call_tool(
         *args: Any,
@@ -180,10 +176,18 @@ def convert_mcp_tool_to_blaxel_tool(
         response_format="content_and_artifact",
     )
 
+
 toolPersistances: dict[str, PersistentWebSocket] = {}
 
+
 class BlTools:
-    def __init__(self, functions: list[str], metas: dict[str, Any] = {}, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True):
+    def __init__(
+        self,
+        functions: list[str],
+        metas: dict[str, Any] = {},
+        timeout: int = DEFAULT_TIMEOUT,
+        timeout_enabled: bool = True,
+    ):
         self.functions = functions
         self.metas = metas
         self.timeout = timeout
@@ -224,52 +228,13 @@ class BlTools:
             toolPersistances.get(name).with_metas(self.metas)
             websocket = toolPersistances.get(name)
             tools = websocket.get_tools()
-            converted_tools = [convert_mcp_tool_to_blaxel_tool(websocket, name, self._url(name), tool) for tool in tools]
+            converted_tools = [
+                convert_mcp_tool_to_blaxel_tool(websocket, name, self._url(name), tool)
+                for tool in tools
+            ]
             all_tools.extend(converted_tools)
         return all_tools
 
-    async def to_langchain(self):
-        from .langchain import get_langchain_tools
-
-        await self.initialize()
-        return get_langchain_tools(self.get_tools())
-
-    async def to_llamaindex(self):
-        from .llamaindex import get_llamaindex_tools
-
-        await self.initialize()
-        return get_llamaindex_tools(self.get_tools())
-
-    async def to_crewai(self):
-        from .crewai import get_crewai_tools
-
-        await self.initialize()
-        return get_crewai_tools(self.get_tools())
-
-    async def to_openai(self):
-        from .openai import get_openai_tools
-
-        await self.initialize()
-        return get_openai_tools(self.get_tools())
-
-    async def to_pydantic(self):
-        from .pydantic import get_pydantic_tools
-
-        await self.initialize()
-        return get_pydantic_tools(self.get_tools())
-
-    async def to_google_adk(self):
-        from .googleadk import get_google_adk_tools
-
-        await self.initialize()
-        return get_google_adk_tools(self.get_tools())
-
-    async def to_livekit(self):
-        from .livekit import get_livekit_tools
-
-        await self.initialize()
-        return get_livekit_tools(self.get_tools())
-
     async def connect(self, name: str):
         # Create and store the connection
         try:
@@ -283,12 +248,12 @@ class BlTools:
             try:
                 await self.connect_with_url(name, url)
             except Exception as e:
-                logger.error(f"Error connecting to {name} with fallback URL: {e}\n{traceback.format_exc()}")
+                logger.error(
+                    f"Error connecting to {name} with fallback URL: {e}\n{traceback.format_exc()}"
+                )
                 raise e
 
-    async def connect_with_url(
-        self, name: str, url: str
-    ) -> None:
+    async def connect_with_url(self, name: str, url: str) -> None:
         """Initialize a session and load tools from it.
 
         Args:
@@ -299,17 +264,24 @@ class BlTools:
 
         if not toolPersistances.get(name):
             logger.debug(f"Creating new persistent websocket for {name}")
-            toolPersistances[name] = PersistentWebSocket(url, name, timeout=self.timeout, timeout_enabled=self.timeout_enabled)
+            toolPersistances[name] = PersistentWebSocket(
+                url, name, timeout=self.timeout, timeout_enabled=self.timeout_enabled
+            )
             await toolPersistances[name].list_tools()
         logger.debug(f"Loaded {len(toolPersistances[name].get_tools())} tools from {url}")
         return toolPersistances[name].with_metas(self.metas)
 
     async def initialize(self) -> "BlTools":
         for i in range(0, len(self.functions), 10):
-            batch = self.functions[i:i+10]
+            batch = self.functions[i : i + 10]
             await asyncio.gather(*(self.connect(name) for name in batch))
         return self
 
 
-def bl_tools(functions: list[str], metas: dict[str, Any] = {}, timeout: int = DEFAULT_TIMEOUT, timeout_enabled: bool = True) -> BlTools:
+def bl_tools(
+    functions: list[str],
+    metas: dict[str, Any] = {},
+    timeout: int = DEFAULT_TIMEOUT,
+    timeout_enabled: bool = True,
+) -> BlTools:
     return BlTools(functions, metas=metas, timeout=timeout, timeout_enabled=timeout_enabled)

blaxel/crewai/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel CrewAI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/crewai.py → crewai/model.py}

@@ -2,11 +2,13 @@ from logging import getLogger
 
 from crewai import LLM
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
-async def get_crewai_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return LLM(
             model=f"mistral/{model}",

blaxel/crewai/tools.py

@@ -0,0 +1,26 @@
+from crewai.tools import BaseTool
+
+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
+
+
+class CrewAITool(BaseTool):
+    _tool: Tool
+
+    def __init__(self, tool: Tool):
+        super().__init__(
+            name=tool.name,
+            description=tool.description,
+            args_schema=create_model_from_json_schema(tool.input_schema),
+        )
+        self._tool = tool
+
+    def _run(self, *args, **kwargs):
+        return self._tool.sync_coroutine(**kwargs)
+
+async def bl_tools(tools_names: list[str], **kwargs) -> list[BaseTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [CrewAITool(tool) for tool in tools.get_tools()]
+    return [CrewAITool(tool) for tool in tools.get_tools()]

blaxel/googleadk/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel Google ADK integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/googleadk.py → googleadk/model.py}

@@ -2,7 +2,8 @@ from logging import getLogger
 
 from google.adk.models.lite_llm import LiteLlm
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
@@ -64,4 +65,9 @@ async def get_google_adk_model(url: str, type: str, model: str, **kwargs):
             api_key=settings.auth.token,
             api_base=f"{url}/v1",
             **kwargs
-        )
+        )
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
+    return await get_google_adk_model(url, type, model, **kwargs)
+    return await get_google_adk_model(url, type, model, **kwargs)

blaxel/googleadk/tools.py

@@ -0,0 +1,72 @@
+import inspect
+from typing import Any, Optional, override
+
+from google.adk.tools import BaseTool, ToolContext
+from google.genai import types
+
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
+
+
+class GoogleADKTool(BaseTool):
+    _tool: Tool
+
+    def __init__(self, tool: Tool):
+        super().__init__(
+            name=tool.name,
+            description=tool.description,
+        )
+        self._tool = tool
+
+    def _clean_schema(self, schema: dict) -> dict:
+        if not isinstance(schema, dict):
+            return schema
+
+        # Create a copy of the schema
+        cleaned_schema = schema.copy()
+
+        # Remove $schema and additionalProperties at current level
+        if "$schema" in cleaned_schema:
+            del cleaned_schema["$schema"]
+        if "additionalProperties" in cleaned_schema:
+            del cleaned_schema["additionalProperties"]
+
+        # Recursively clean properties if they exist
+        if "properties" in cleaned_schema:
+            cleaned_schema["properties"] = {
+                k: self._clean_schema(v) for k, v in cleaned_schema["properties"].items()
+            }
+
+        return cleaned_schema
+
+    @override
+    def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
+        # Clean the schema recursively
+        schema = self._clean_schema(self._tool.input_schema)
+
+        function_decl = types.FunctionDeclaration.model_validate(
+            types.FunctionDeclaration(
+                name=self._tool.name,
+                description=self._tool.description,
+                parameters=schema,
+            )
+        )
+
+        return function_decl
+
+    @override
+    async def run_async(
+        self, *, args: dict[str, Any], tool_context: ToolContext
+    ) -> Any:
+        args_to_call = args.copy()
+        signature = inspect.signature(self._tool.coroutine)
+        if 'tool_context' in signature.parameters:
+            args_to_call['tool_context'] = tool_context
+        return await self._tool.coroutine(**args_to_call) or {}
+
+
+async def bl_tools(tools_names: list[str], **kwargs: Any) -> list[BaseTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [GoogleADKTool(tool) for tool in tools.get_tools()]
+    return [GoogleADKTool(tool) for tool in tools.get_tools()]

blaxel/langgraph/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LangGraph integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/langchain.py → langgraph/model.py}

@@ -7,12 +7,16 @@ from langchain_deepseek import ChatDeepSeek
 from langchain_openai import ChatOpenAI
 from langchain_xai import ChatXAI
 
-from ..common.settings import settings
-from .custom.langchain.gemini import ChatGoogleGenerativeAI
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+
+from .custom.gemini import ChatGoogleGenerativeAI
 
 logger = getLogger(__name__)
 
-async def get_langchain_model(url: str,type: str, model: str, **kwargs):
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return ChatOpenAI(
             api_key=settings.auth.token,
@@ -72,4 +76,4 @@ async def get_langchain_model(url: str,type: str, model: str, **kwargs):
             model=model,
             base_url=f"{url}/v1",
             **kwargs
-        )
+        )

blaxel/{tools/langchain.py → langgraph/tools.py}

@@ -3,7 +3,8 @@ from typing import Any
 from langchain_core.tools import StructuredTool
 from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
 
-from .types import Tool, ToolException
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool, ToolException
 
 NonTextContent = ImageContent | EmbeddedResource
 
@@ -37,5 +38,8 @@ def get_langchain_tool(tool: Tool) -> StructuredTool:
         sync_coroutine=tool.sync_coroutine,
     )
 
-def get_langchain_tools(tools: list[Tool]) -> list[StructuredTool]:
-    return [get_langchain_tool(tool) for tool in tools]
+
+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/livekit/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LiveKit integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/livekit.py → livekit/model.py}

@@ -2,7 +2,8 @@ from logging import getLogger
 
 from livekit.plugins import groq, openai
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 logger = getLogger(__name__)
 
@@ -23,3 +24,8 @@ async def get_livekit_model(url: str, type: str, model: str, **kwargs):
             base_url=f"{url}/v1",
             **kwargs
         )
+
+async def bl_model(name: str, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
+    return await get_livekit_model(url, type, model, **kwargs)
+    return await get_livekit_model(url, type, model, **kwargs)

blaxel/{tools/livekit.py → livekit/tools.py}

@@ -1,6 +1,7 @@
 from livekit.agents import function_tool, llm
 
-from .types import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
 
 
 def livekit_coroutine(tool: Tool):
@@ -21,3 +22,9 @@ def get_livekit_tools(tools: list[Tool]) -> list[llm.FunctionTool]:
             },
         ))
     return livekit_tools
+
+async def bl_tools(tools_names: list[str], **kwargs) -> list[llm.FunctionTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return get_livekit_tools(tools.get_tools())
+    return get_livekit_tools(tools.get_tools())

blaxel/llamaindex/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel LlamaIndex integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/llamaindex.py → llamaindex/model.py}

@@ -14,12 +14,15 @@ from llama_index.llms.groq import Groq
 from llama_index.llms.mistralai import MistralAI
 from llama_index.llms.openai import OpenAI
 
-from ..common.settings import settings
-from .custom.llamaindex.cohere import Cohere
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+
+from .custom.cohere import Cohere
 
 logger = getLogger(__name__)
 
-async def get_llamaindex_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'anthropic':
         return Anthropic(
             model=model,

blaxel/{tools/llamaindex.py → llamaindex/tools.py}

@@ -3,8 +3,9 @@
 from llama_index.core.tools import FunctionTool
 from llama_index.core.tools.types import ToolMetadata
 
-from .common import create_model_from_json_schema
-from .types import Tool
+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
 
 
 def get_llamaindex_tool(tool: Tool) -> FunctionTool:
@@ -21,5 +22,7 @@ def get_llamaindex_tool(tool: Tool) -> FunctionTool:
         ),
     )
 
-def get_llamaindex_tools(tools: list[Tool]) -> list[FunctionTool]:
-    return [get_llamaindex_tool(tool) for tool in tools]
+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/openai/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel OpenAI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/openai.py → openai/model.py}

@@ -1,9 +1,11 @@
 from agents import AsyncOpenAI, OpenAIChatCompletionsModel
 
-from ..common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
 
 
-async def get_openai_model(url: str, type: str, model: str, **kwargs):
+async def bl_model(name, **kwargs):
+    url, type, model = await bl_model_core(name).get_parameters()
     if type != "openai":
         raise ValueError(f"Invalid model type: {type}")
     external_client = AsyncOpenAI(

blaxel/{tools/openai.py → openai/tools.py}

@@ -3,7 +3,8 @@ from typing import Any
 
 from agents import FunctionTool, RunContextWrapper
 
-from .types import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
+from blaxel.core.tools.types import Tool
 
 
 def get_openai_tool(tool: Tool) -> FunctionTool:
@@ -21,5 +22,8 @@ def get_openai_tool(tool: Tool) -> FunctionTool:
         on_invoke_tool=openai_coroutine,
     )
 
-def get_openai_tools(tools: list[Tool]) -> list[FunctionTool]:
-    return [get_openai_tool(tool) for tool in tools]
+async def bl_tools(tools_names: list[str], **kwargs) -> list[FunctionTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_openai_tool(tool) for tool in tools.get_tools()]
+    return [get_openai_tool(tool) for tool in tools.get_tools()]

blaxel/pydantic/__init__.py

@@ -0,0 +1,4 @@
+"""Blaxel Pydantic AI integration module."""
+
+from .model import *  # noqa: F403, F401
+from .tools import *  # noqa: F403, F401

blaxel/{models/custom/pydantic → pydantic/custom}/gemini.py

@@ -1,4 +1,3 @@
-
 import httpx
 from httpx import AsyncClient
 from pydantic_ai.providers import Provider

blaxel/{models/pydantic.py → pydantic/model.py}

@@ -14,13 +14,15 @@ from pydantic_ai.providers.cohere import CohereProvider
 from pydantic_ai.providers.mistral import MistralProvider
 from pydantic_ai.providers.openai import OpenAIProvider
 
-from blaxel.common.settings import settings
+from blaxel.core import bl_model as bl_model_core
+from blaxel.core import settings
+from blaxel.core.client import client
 
-from ..client import client
-from .custom.pydantic.gemini import GoogleGLAProvider
+from .custom.gemini import GoogleGLAProvider
 
 logger = logging.getLogger(__name__)
-async def get_pydantic_model(url: str,type: str, model: str, **kwargs) -> Model:
+async def bl_model(name: str, **kwargs) -> Model:
+    url, type, model = await bl_model_core(name).get_parameters()
     if type == 'mistral':
         return MistralModel(
             model_name=model,

blaxel/{tools/pydantic.py → pydantic/tools.py}

@@ -2,7 +2,8 @@ from pydantic_ai import RunContext
 from pydantic_ai.tools import Tool as PydanticTool
 from pydantic_ai.tools import ToolDefinition
 
-from .types import Tool
+from blaxel.core.tools import Tool
+from blaxel.core.tools import bl_tools as bl_tools_core
 
 
 def get_pydantic_tool(tool: Tool) -> PydanticTool:
@@ -44,5 +45,7 @@ def get_pydantic_tool(tool: Tool) -> PydanticTool:
     )
 
 
-def get_pydantic_tools(tools: list[Tool]) -> list[PydanticTool]:
-    return [get_pydantic_tool(tool) for tool in tools]
+async def bl_tools(tools_names: list[str], **kwargs) -> list[PydanticTool]:
+    tools = bl_tools_core(tools_names, **kwargs)
+    await tools.initialize()
+    return [get_pydantic_tool(tool) for tool in tools.get_tools()]