beamlit 0.0.57rc111__py3-none-any.whl → 0.0.57rc112__py3-none-any.whl

beamlit/agents/decorator.py

@@ -4,25 +4,123 @@ Defines decorators for agent functionalities.
 """
 
 # Import necessary modules
+import asyncio
 import functools
 import inspect
 from logging import getLogger
 from typing import Callable
 
-from langgraph.checkpoint.memory import MemorySaver
-from langgraph.prebuilt import create_react_agent
-
 from beamlit.api.models import get_model, list_models
 from beamlit.authentication import new_client
-from beamlit.common.settings import init
+from beamlit.common.settings import Settings, init
 from beamlit.errors import UnexpectedStatus
 from beamlit.functions import get_functions
 from beamlit.models import Agent, AgentSpec, Metadata
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import create_react_agent
 
 from .chat import get_chat_model_full
 from .voice.openai import OpenAIVoiceReactAgent
 
 
+async def initialize_agent(
+    settings: Settings,
+    agent: Agent | dict = None,
+    override_model=None,
+    override_agent=None,
+    override_functions=None,
+    remote_functions=None,
+    local_functions=None,
+):
+    logger = getLogger(__name__)
+    client = new_client()
+    chat_model = override_model or None
+
+    if agent is not None:
+        metadata = Metadata(**agent.get("metadata", {}))
+        spec = AgentSpec(**agent.get("spec", {}))
+        agent = Agent(metadata=metadata, spec=spec)
+        if agent.spec.model and chat_model is None:
+            try:
+                response = get_model.sync_detailed(
+                    agent.spec.model, client=client
+                )
+                settings.agent.model = response.parsed
+            except UnexpectedStatus as e:
+                if e.status_code == 404:
+                    if e.status_code == 404:
+                        raise ValueError(f"Model {agent.spec.model} not found")
+                raise e
+            except Exception as e:
+                raise e
+
+            if settings.agent.model:
+                chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
+                settings.agent.chat_model = chat_model
+                logger.info(f"Chat model configured, using: {provider}:{model}")
+
+    if override_functions is not None:
+        functions = override_functions
+    else:
+        functions = await get_functions(
+            client=client,
+            dir=settings.agent.functions_directory,
+            remote_functions=remote_functions,
+            chain=agent.spec.agent_chain,
+            local_functions=local_functions,
+            remote_functions_empty=not remote_functions,
+            warning=chat_model is not None,
+        )
+    settings.agent.functions = functions
+
+    if override_agent is None:
+        if chat_model is None:
+            models_select = ""
+            try:
+                models = list_models.sync_detailed(
+                    client=client
+                )
+                models = ", ".join([model.metadata.name for model in models.parsed])
+                models_select = f"You can select one from your models: {models}"
+            except Exception:
+                pass
+
+            raise ValueError(
+                f"You must provide a model.\n"
+                f"{models_select}\n"
+                f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
+                "Add it to your agent spec\n"
+                "agent={\n"
+                '    "metadata": {\n'
+                f'        "name": "{agent.metadata.name}",\n'
+                "    },\n"
+                '    "spec": {\n'
+                '        "model": "MODEL_NAME",\n'
+                f'        "description": "{agent.spec.description}",\n'
+                f'        "prompt": "{agent.spec.prompt}",\n'
+                "    },\n"
+                "}")
+        if isinstance(chat_model, OpenAIVoiceReactAgent):
+            _agent = chat_model
+        else:
+            memory = MemorySaver()
+            if len(functions) == 0:
+                raise ValueError("You can define this function in directory "
+                    f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
+                    "from beamlit.functions import function\n\n"
+                    "@function()\n"
+                    "def hello_world(query: str):\n"
+                    "    return 'Hello, world!'\n")
+            try:
+                _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
+            except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
+                logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
+                _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
+
+        settings.agent.agent = _agent
+    else:
+        settings.agent.agent = override_agent
+
 def agent(
     agent: Agent | dict = None,
     override_model=None,
@@ -58,17 +156,17 @@ def agent(
         Re-raises exceptions encountered during model retrieval and agent setup.
     """
     logger = getLogger(__name__)
+    settings = init()
+    _is_initialized = False
     try:
         if agent is not None and not isinstance(agent, dict):
             raise Exception(
                 'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
             )
 
-        client = new_client()
-        chat_model = override_model or None
-        settings = init()
 
         def wrapper(func):
+
             agent_kwargs = any(
                 param.name == "agent"
                 for param in inspect.signature(func).parameters.values()
@@ -83,99 +181,23 @@ def agent(
             )
 
             @functools.wraps(func)
-            def wrapped(*args, **kwargs):
+            async def wrapped(*args, **kwargs):
+                nonlocal _is_initialized
+                if not _is_initialized:
+                    async with asyncio.Lock():
+                        if not _is_initialized:
+                            await initialize_agent(settings, agent, override_model, override_agent, override_functions, remote_functions, local_functions)
+                            _is_initialized = True
                 if agent_kwargs:
                     kwargs["agent"] = settings.agent.agent
                 if model_kwargs:
                     kwargs["model"] = settings.agent.chat_model
                 if functions_kwargs:
                     kwargs["functions"] = settings.agent.functions
-                return func(*args, **kwargs)
+                return await func(*args, **kwargs)
 
             return wrapped
 
-        if agent is not None:
-            metadata = Metadata(**agent.get("metadata", {}))
-            spec = AgentSpec(**agent.get("spec", {}))
-            agent = Agent(metadata=metadata, spec=spec)
-            if agent.spec.model and chat_model is None:
-                try:
-                    response = get_model.sync_detailed(
-                        agent.spec.model, client=client
-                    )
-                    settings.agent.model = response.parsed
-                except UnexpectedStatus as e:
-                    raise e
-                except Exception as e:
-                    raise e
-
-                if settings.agent.model:
-                    chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
-                    settings.agent.chat_model = chat_model
-                    logger.info(f"Chat model configured, using: {provider}:{model}")
-
-        if override_functions is not None:
-            functions = override_functions
-        else:
-            functions = get_functions(
-                client=client,
-                dir=settings.agent.functions_directory,
-                remote_functions=remote_functions,
-                chain=agent.spec.agent_chain,
-                local_functions=local_functions,
-                remote_functions_empty=not remote_functions,
-                warning=chat_model is not None,
-            )
-
-        settings.agent.functions = functions
-
-        if override_agent is None:
-            if chat_model is None:
-                models_select = ""
-                try:
-                    models = list_models.sync_detailed(
-                        client=client
-                    )
-                    models = ", ".join([model.metadata.name for model in models.parsed])
-                    models_select = f"You can select one from your models: {models}"
-                except Exception:
-                    pass
-
-                raise ValueError(
-                    f"You must provide a model.\n"
-                    f"{models_select}\n"
-                    f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
-                    "Add it to your agent spec\n"
-                    "agent={\n"
-                    '    "metadata": {\n'
-                    f'        "name": "{agent.metadata.name}",\n'
-                    "    },\n"
-                    '    "spec": {\n'
-                    '        "model": "MODEL_NAME",\n'
-                    f'        "description": "{agent.spec.description}",\n'
-                    f'        "prompt": "{agent.spec.prompt}",\n'
-                    "    },\n"
-                    "}")
-            if isinstance(chat_model, OpenAIVoiceReactAgent):
-                _agent = chat_model
-            else:
-                memory = MemorySaver()
-                if len(functions) == 0:
-                    raise ValueError("You can define this function in directory "
-                        f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
-                        "from beamlit.functions import function\n\n"
-                        "@function()\n"
-                        "def hello_world(query: str):\n"
-                        "    return 'Hello, world!'\n")
-                try:
-                    _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
-                except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
-                    logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
-                    _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
-
-            settings.agent.agent = _agent
-        else:
-            settings.agent.agent = override_agent
         return wrapper
     except Exception as e:
         logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")

beamlit/functions/common.py

@@ -21,9 +21,6 @@ import traceback
 from logging import getLogger
 from typing import Union
 
-from langchain_core.tools import StructuredTool
-from langchain_core.tools.base import create_schema_from_function
-
 from beamlit.authentication import new_client
 from beamlit.client import AuthenticatedClient
 from beamlit.common import slugify
@@ -31,10 +28,12 @@ from beamlit.common.settings import get_settings
 from beamlit.functions.local.local import LocalToolKit
 from beamlit.functions.remote.remote import RemoteToolkit
 from beamlit.models import AgentChain
+from langchain_core.tools import StructuredTool
+from langchain_core.tools.base import create_schema_from_function
 
 logger = getLogger(__name__)
 
-def get_functions(
+async def get_functions(
     remote_functions: Union[list[str], None] = None,
     local_functions: Union[list[dict], None] = None,
     client: Union[AuthenticatedClient, None] = None,
@@ -139,7 +138,7 @@ def get_functions(
                                             ):
                                                 is_kit = keyword.value.value
                                     if is_kit and not settings.remote:
-                                        kit_functions = get_functions(
+                                        kit_functions = await get_functions(
                                             client=client,
                                             dir=os.path.join(root),
                                             remote_functions_empty=remote_functions_empty,
@@ -183,20 +182,21 @@ def get_functions(
         for function in remote_functions:
             try:
                 toolkit = RemoteToolkit(client, function)
-                toolkit.initialize()
-                functions.extend(toolkit.get_tools())
+                await toolkit.initialize()
+                functions.extend(await toolkit.get_tools())
             except Exception as e:
-                logger.debug(
-                    f"Failed to initialize remote function {function}: {e!s}\n"
-                    f"Traceback:\n{traceback.format_exc()}"
-                )
-                logger.warn(f"Failed to initialize remote function {function}: {e!s}")
+                if not isinstance(e, RuntimeError):
+                    logger.debug(
+                        f"Failed to initialize remote function {function}: {e!s}\n"
+                        f"Traceback:\n{traceback.format_exc()}"
+                    )
+                    logger.warn(f"Failed to initialize remote function {function}: {e!s}")
     if local_functions:
         for function in local_functions:
             try:
                 toolkit = LocalToolKit(client, function)
-                toolkit.initialize()
-                functions.extend(toolkit.get_tools())
+                await toolkit.initialize()
+                functions.extend(await toolkit.get_tools())
             except Exception as e:
                 logger.debug(
                     f"Failed to initialize local function {function}: {e!s}\n"
@@ -206,8 +206,7 @@ def get_functions(
 
     if chain:
         toolkit = ChainToolkit(client, chain)
-        toolkit.initialize()
-        functions.extend(toolkit.get_tools())
-
+        await toolkit.initialize()
+        functions.extend(await toolkit.get_tools())
     return functions
 

beamlit/functions/decorator.py

@@ -4,9 +4,8 @@ import functools
 from collections.abc import Callable
 from logging import getLogger
 
-from fastapi import Request
-
 from beamlit.models import Function, FunctionKit
+from fastapi import Request
 
 logger = getLogger(__name__)
 

beamlit/functions/local/local.py

@@ -1,11 +1,10 @@
 from dataclasses import dataclass
 
 import pydantic
-from langchain_core.tools.base import BaseTool
-
 from beamlit.authentication.authentication import AuthenticatedClient
 from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
 from beamlit.models import Function
+from langchain_core.tools.base import BaseTool
 
 
 @dataclass
@@ -23,7 +22,7 @@ class LocalToolKit:
     _function: Function | None = None
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         """Initialize the session and retrieve the local function details."""
         if self._function is None:
             try:
@@ -34,7 +33,6 @@ class LocalToolKit:
                     spec={
                         "configurations": {
                             "url": self.local_function['url'],
-                            "sse": self.local_function['sse'],
                         },
                         "description": self.local_function['description'] or "",
                     }
@@ -42,8 +40,8 @@ class LocalToolKit:
             except Exception as e:
                 raise RuntimeError(f"Failed to initialize local function: {e}")
 
-    def get_tools(self) -> list[BaseTool]:
+    async def get_tools(self) -> list[BaseTool]:
         mcp_client = MCPClient(self.client, self._function.spec["configurations"]["url"], sse=self._function.spec["configurations"]["sse"])
         mcp_toolkit = MCPToolkit(client=mcp_client)
-        mcp_toolkit.initialize()
-        return mcp_toolkit.get_tools()
+        await mcp_toolkit.initialize()
+        return await mcp_toolkit.get_tools()

beamlit/functions/mcp/client.py

@@ -0,0 +1,98 @@
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+from urllib.parse import urljoin, urlparse
+
+import anyio
+import httpx
+from anyio.abc import TaskStatus
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from websockets.client import connect as ws_connect
+from websockets.client import WebSocketClientProtocol
+
+import mcp.types as types
+
+logger = logging.getLogger(__name__)
+
+
+def remove_request_params(url: str) -> str:
+    return urljoin(url, urlparse(url).path)
+
+
+@asynccontextmanager
+async def websocket_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5,
+):
+    """
+    Client transport for WebSocket.
+
+    The `timeout` parameter controls connection timeout.
+    """
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    async with anyio.create_task_group() as tg:
+        try:
+            # Convert http(s):// to ws(s)://
+            ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
+            logger.debug(f"Connecting to WebSocket endpoint: {remove_request_params(ws_url)}")
+            
+            async with ws_connect(ws_url, extra_headers=headers, open_timeout=timeout) as websocket:
+                logger.debug("WebSocket connection established")
+
+                async def ws_reader(
+                    websocket: WebSocketClientProtocol,
+                    task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+                ):
+                    try:
+                        task_status.started()
+                        async for message in websocket:
+                            logger.debug(f"Received WebSocket message: {message}")
+                            try:
+                                parsed_message = types.JSONRPCMessage.model_validate_json(message)
+                                logger.debug(f"Received server message: {parsed_message}")
+                                await read_stream_writer.send(parsed_message)
+                            except Exception as exc:
+                                logger.error(f"Error parsing server message: {exc}")
+                                await read_stream_writer.send(exc)
+                    except Exception as exc:
+                        logger.error(f"Error in ws_reader: {exc}")
+                        await read_stream_writer.send(exc)
+                    finally:
+                        await read_stream_writer.aclose()
+
+                async def ws_writer(websocket: WebSocketClientProtocol):
+                    try:
+                        async with write_stream_reader:
+                            async for message in write_stream_reader:
+                                logger.debug(f"Sending client message: {message}")
+                                await websocket.send(
+                                    message.model_dump_json(
+                                        by_alias=True,
+                                        exclude_none=True,
+                                    )
+                                )
+                                logger.debug("Client message sent successfully")
+                    except Exception as exc:
+                        logger.error(f"Error in ws_writer: {exc}")
+                    finally:
+                        await write_stream.aclose()
+
+                await tg.start(ws_reader, websocket)
+                tg.start_soon(ws_writer, websocket)
+
+                try:
+                    yield read_stream, write_stream
+                finally:
+                    tg.cancel_scope.cancel()
+        finally:
+            await read_stream_writer.aclose()
+            await write_stream.aclose()

beamlit/functions/mcp/mcp.py

@@ -12,14 +12,13 @@ import pydantic
 import pydantic_core
 import requests
 import typing_extensions as t
-from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-from mcp.types import CallToolResult, ListToolsResult
-
 from beamlit.authentication import get_authentication_headers
 from beamlit.authentication.authentication import AuthenticatedClient
 from beamlit.common.settings import get_settings
+from beamlit.functions.mcp.client import websocket_client
+from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
+from mcp import ClientSession
+from mcp.types import CallToolResult, ListToolsResult
 
 from .utils import create_schema_model
 
@@ -29,63 +28,65 @@ logger = logging.getLogger(__name__)
 
 
 class MCPClient:
-    def __init__(self, client: AuthenticatedClient, url: str, sse: bool = False):
+    def __init__(self, client: AuthenticatedClient, url: str, fallback_url: str | None = None):
         self.client = client
         self.url = url
-        self._sse = False
+        self.fallback_url = fallback_url
 
-    async def list_sse_tools(self) -> ListToolsResult:
-        # Create a new context for each SSE connection
+    async def list_ws_tools(self, is_fallback: bool = False) -> ListToolsResult:
+        if is_fallback:
+            url = self.fallback_url
+        else:
+            url = self.url
         try:
-            async with sse_client(f"{self.url}/sse", headers=get_authentication_headers(settings)) as (read_stream, write_stream):
-                async with ClientSession(read_stream, write_stream) as session:
-                    await session.initialize()
-                    response = await session.list_tools()
+            async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
+                logger.debug("WebSocket connection established")
+                async with ClientSession(read_stream, write_stream) as client:
+                    await client.initialize()
+                    response = await client.list_tools()
+                    logger.debug(f"WebSocket tools: {response}")
                     return response
-        except Exception:
-            self._sse = False
-            logger.info("SSE not available, trying HTTP")
+        except Exception as e:
+            logger.error(f"Error listing SSE tools: {e}")
+            logger.debug("WebSocket not available, trying HTTP")
             return None  # Signal to list_tools() to try HTTP instead
 
-    def list_tools(self) -> ListToolsResult:
+    async def list_tools(self) -> ListToolsResult:
+        logger.debug(f"Listing tools for {self.url}")
         try:
-            loop = asyncio.get_event_loop()
-            result = loop.run_until_complete(self.list_sse_tools())
-            if result is None:  # SSE failed, try HTTP
-                raise Exception("SSE failed")
-            self._sse = True
+            result = await self.list_ws_tools(is_fallback=False)
             return result
-        except Exception:  # Fallback to HTTP
-            client = self.client.get_httpx_client()
-            response = client.request("GET", f"{self.url}/tools/list")
-            response.raise_for_status()
-            return ListToolsResult(**response.json())
+        except Exception as e: # Fallback to Public endpoint
+            if self.fallback_url:
+                try:
+                    result = await self.list_ws_tools(is_fallback=True)
+                    return result
+                except Exception as e:
+                    raise e
+            else:
+                raise e
+
 
     async def call_tool(
         self,
         tool_name: str,
         arguments: dict[str, Any] = None,
+        is_fallback: bool = False,
     ) -> requests.Response | AsyncIterator[CallToolResult]:
-        if self._sse:
-            async with sse_client(f"{self.url}/sse", headers=get_authentication_headers(settings)) as (read_stream, write_stream):
+        if is_fallback:
+            url = self.fallback_url
+        else:
+            url = self.url
+        try:
+            async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
                 async with ClientSession(read_stream, write_stream) as session:
                     await session.initialize()
                     response = await session.call_tool(tool_name, arguments or {})
                     content = pydantic_core.to_json(response).decode()
                     return content
-        else: # Fallback to HTTP
-            client = self.client.get_httpx_client()
-            response = client.request(
-                "POST",
-                f"{self.url}/tools/call",
-                json={"name": tool_name, "arguments": arguments},
-            )
-            response.raise_for_status()
-            result = CallToolResult(response.json())
-            if result.isError:
-                raise ToolException(result.content)
-            content = pydantic_core.to_json(result.content).decode()
-            return content
+        except Exception as e:
+            raise e
+
 
 class MCPTool(BaseTool):
     """
@@ -94,7 +95,6 @@ class MCPTool(BaseTool):
     Attributes:
         client (MCPClient): The MCP client instance.
         handle_tool_error (bool | str | Callable[[ToolException], str] | None): Error handling strategy.
-        sse (bool): Whether to use SSE streaming for responses.
     """
 
     client: MCPClient
@@ -110,7 +110,16 @@ class MCPTool(BaseTool):
 
     @t.override
     async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
-        return await self.client.call_tool(self.name, arguments=kwargs)
+        try:
+            return await self.client.call_tool(self.name, arguments=kwargs)
+        except Exception as e:
+            if self.client.fallback_url:
+                try:
+                    return await self.client.call_tool(self.name, arguments=kwargs, is_fallback=True) # Fallback to Public endpoint
+                except Exception as e:
+                    raise e
+            else:
+                raise e
 
     @t.override
     @property
@@ -133,14 +142,14 @@ class MCPToolkit(BaseToolkit):
     _tools: ListToolsResult | None = None
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         """Initialize the session and retrieve tools list"""
         if self._tools is None:
-            response = self.client.list_tools()
+            response = await self.client.list_tools()
             self._tools = response
 
     @t.override
-    def get_tools(self) -> list[BaseTool]:
+    async def get_tools(self) -> list[BaseTool]:
         if self._tools is None:
             raise RuntimeError("Must initialize the toolkit first")
 
@@ -150,7 +159,6 @@ class MCPToolkit(BaseToolkit):
                 name=tool.name,
                 description=tool.description or "",
                 args_schema=create_schema_model(tool.name, tool.inputSchema),
-                sse=self.sse,
             )
             # list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
             for tool in self._tools.tools

beamlit/functions/remote/remote.py

@@ -11,8 +11,6 @@ from typing import Callable
 
 import pydantic
 import typing_extensions as t
-from langchain_core.tools.base import BaseTool, ToolException
-
 from beamlit.api.functions import get_function, list_functions
 from beamlit.authentication.authentication import AuthenticatedClient
 from beamlit.common.settings import get_settings
@@ -20,6 +18,7 @@ from beamlit.errors import UnexpectedStatus
 from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
 from beamlit.models import Function, StoreFunctionParameter
 from beamlit.run import RunClient
+from langchain_core.tools.base import BaseTool, ToolException
 
 
 def create_dynamic_schema(name: str, parameters: list[StoreFunctionParameter]) -> type[pydantic.BaseModel]:
@@ -115,7 +114,7 @@ class RemoteToolkit:
     _service_name: str | None = None
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         """Initialize the session and retrieve the remote function details."""
         if self._function is None:
             try:
@@ -136,19 +135,21 @@ class RemoteToolkit:
                     f"error: {e.status_code}. Available functions: {', '.join(names)}"
                 )
 
-    def get_tools(self) -> list[BaseTool]:
+    async def get_tools(self) -> list[BaseTool]:
         settings = get_settings()
         if self._function is None:
             raise RuntimeError("Must initialize the toolkit first")
 
         if self._function.spec.integration_connections:
+            fallback_url = None
             url = f"{settings.run_url}/{settings.workspace}/functions/{self._function.metadata.name}"
             if self._service_name:
+                fallback_url = f"https://{self._service_name}.{settings.run_internal_hostname}"
                 url = f"https://{self._service_name}.{settings.run_internal_hostname}"
-            mcp_client = MCPClient(self.client, url)
-            mcp_toolkit = MCPToolkit(client=mcp_client)
-            mcp_toolkit.initialize()
-            return mcp_toolkit.get_tools()
+            mcp_client = MCPClient(self.client, url, fallback_url)
+            mcp_toolkit = MCPToolkit(client=mcp_client, url=url)
+            await mcp_toolkit.initialize()
+            return await mcp_toolkit.get_tools()
 
         if self._function.spec.kit:
             return [

{beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beamlit
-Version: 0.0.57rc111
+Version: 0.0.57rc112
 Summary: Add your description here
 Author-email: cploujoux <ch.ploujoux@gmail.com>
 License-File: LICENSE
@@ -43,6 +43,7 @@ Requires-Dist: pyjwt>=2.10.1
 Requires-Dist: python-dateutil>=2.8.0
 Requires-Dist: pyyaml<6.1.0,>=6.0.2
 Requires-Dist: requests<2.33.0,>=2.32.3
+Requires-Dist: websocket>=0.2.1
 Requires-Dist: websockets>=14.1
 Description-Content-Type: text/markdown
 

{beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/RECORD

@@ -7,7 +7,7 @@ beamlit/types.py,sha256=E1hhDh_zXfsSQ0NCt9-uw90_Mr5iIlsdfnfvxv5HarU,1005
 beamlit/agents/__init__.py,sha256=bWsFaXUbAps3IsL3Prti89m1s714vICXodbQi77h3vY,206
 beamlit/agents/chain.py,sha256=JsinjAYBr3oaM4heouZaiaV2jMmi779LHAMtD_4P59s,4867
 beamlit/agents/chat.py,sha256=ufuydptucLNe_Jyr7lQO1WfQ5pe0I5YKh-y0smwWABM,8301
-beamlit/agents/decorator.py,sha256=oXhaLXeyuw3gOoiDFoszoilJgJLo9hX9WHL30DSw5go,7754
+beamlit/agents/decorator.py,sha256=pKA9NIaNoToCsjMkvBpOvfQUsb0B20XwdskTonz-isI,8310
 beamlit/agents/thread.py,sha256=PVP9Gey8fMZHA-Cs8WbfmdSlD7g-n4HKuk1sTVf8yOQ,1087
 beamlit/agents/voice/openai.py,sha256=-RDBwl16i4TbUhFo5-77Ci3zmI3Y8U2yf2MmvXR2haQ,9479
 beamlit/agents/voice/utils.py,sha256=tQidyM40Ewuy12wKqpvJLvfJgneQ0sZf50dqnerPGHg,836
@@ -131,12 +131,13 @@ beamlit/deploy/deploy.py,sha256=qpNXHvLj712Uo34MPT8z0AkPNsG03lEPe08p1XpGQhc,1116
 beamlit/deploy/format.py,sha256=W3ESUHyFv-iZDjVnHOf9YFDDXZSXYIFFbwCoL1GInE0,1162
 beamlit/deploy/parser.py,sha256=gjRUhOVtfKnc1UNc_FhXsEfj9zrMNuq8W93pNsJBpo0,7586
 beamlit/functions/__init__.py,sha256=Mnoqpa1dm7TXwjodBbF_40JyD78aXsOYWmqjDSnA1lU,317
-beamlit/functions/common.py,sha256=4SK1N0VoQiydUT4frGT9OWMRxgHgMHUcsBhicMkcMmc,9996
-beamlit/functions/decorator.py,sha256=iQbLwUo0K83DFJ3ub8O5jKtkbSINnku6GZcKJ9h7-5E,2292
-beamlit/functions/local/local.py,sha256=KjkHWBxGlG9fliXnOOwZQEop1g2o10IaLiPo4Zu_XAk,1929
-beamlit/functions/mcp/mcp.py,sha256=itd7i5Ej0WSBmw8Z6pbdUL1-NJwDoIQ1edeFyg1o4Co,5814
+beamlit/functions/common.py,sha256=uPmFx8odqHVas8kSfc41B9g_raWtAAzQBNyuRjtDDb0,10114
+beamlit/functions/decorator.py,sha256=P17iKgDPe5VBEgwLPep1yVTD67C8YtRXuoe_kNEfBFg,2291
+beamlit/functions/local/local.py,sha256=GQlp-dlEd465qMwJ-oSkVtLdDGtIYEE4PsVoAzbJJyg,1890
+beamlit/functions/mcp/client.py,sha256=SwGINxSx9a3_LlDAiYuf8NNw8KZf8B5PIOls_wjJlkU,4086
+beamlit/functions/mcp/mcp.py,sha256=T9dUc1PvYO1FXADLu0rZY4FIjUFFQAOuOkiZghLd8eo,5949
 beamlit/functions/mcp/utils.py,sha256=V7bah6cymdtjJ_LJUrNcHDeApDHA6uXvaGVeFJGKj2U,1850
-beamlit/functions/remote/remote.py,sha256=iB8DmJBUj-NchQoPlAPTNLASywIWYlkSlqvHDl-ACYE,6553
+beamlit/functions/remote/remote.py,sha256=O7zq315jHkcvNbrd4AZfsxn4PVdqGvrecKxBR-EwV9U,6727
 beamlit/models/__init__.py,sha256=042wT7sG_YUmJJAb9edrui-DesMxLjOGVvnEtqc92CY,8916
 beamlit/models/acl.py,sha256=tH67gsl_BMaviSbTaaIkO1g9cWZgJ6VgAnYVjQSzGZY,3952
 beamlit/models/agent.py,sha256=DPZZOrDyOAXlh7uvPmhK1jgz8L4UWmaF_SH-OnAdjlw,4162
@@ -253,8 +254,8 @@ beamlit/serve/app.py,sha256=5XZci-R95Zjl97wMtQd1BRtonnkJJ2AeoTVFPKGAOfA,4283
 beamlit/serve/middlewares/__init__.py,sha256=O7fyfE1DIYmajFY9WWdzxCgeAQWZzJfeUjzHGbpWaAk,309
 beamlit/serve/middlewares/accesslog.py,sha256=lcu33j4epFSHRBaeTpyt8deNb3kaM3K91-andw4fp80,1112
 beamlit/serve/middlewares/processtime.py,sha256=3x5w1yQexB0xFNKK6fgLbINxT-eLLunfZ6UDV0bIIF4,944
-beamlit-0.0.57rc111.dist-info/METADATA,sha256=5HLIcsxG56ZK1gCw97tpXUpjnlJ2iSlQAk-gs8Na2wg,3515
-beamlit-0.0.57rc111.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-beamlit-0.0.57rc111.dist-info/entry_points.txt,sha256=zxhgdn7SP-Otk4rEv7LMPAAa9w4TUCLbu9TJi9-K3xg,115
-beamlit-0.0.57rc111.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
-beamlit-0.0.57rc111.dist-info/RECORD,,
+beamlit-0.0.57rc112.dist-info/METADATA,sha256=KHt8c7Zlqs1EnvWmSVAEYgO3hasTEWMbR8AnnSmUIw0,3547
+beamlit-0.0.57rc112.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+beamlit-0.0.57rc112.dist-info/entry_points.txt,sha256=zxhgdn7SP-Otk4rEv7LMPAAa9w4TUCLbu9TJi9-K3xg,115
+beamlit-0.0.57rc112.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
+beamlit-0.0.57rc112.dist-info/RECORD,,