beamlit 0.0.57rc111__tar.gz → 0.0.57rc113__tar.gz

{beamlit-0.0.57rc111 → beamlit-0.0.57rc113}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beamlit
-Version: 0.0.57rc111
+Version: 0.0.57rc113
 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 → beamlit-0.0.57rc113}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "beamlit"
-version = "0.0.57.pre-111"
+version = "0.0.57.pre-113"
 description = "Add your description here"
 readme = "README.md"
 authors = [
@@ -51,6 +51,7 @@ dependencies = [
     "opentelemetry-instrumentation-weaviate>=0.35.0",
     "mcp>=1.2.1",
     "websockets>=14.1",
+    "websocket>=0.2.1",
 ]
 [dependency-groups]
 dev = [

beamlit-0.0.57rc113/src/beamlit/agents/decorator.py

@@ -0,0 +1,205 @@
+"""Module: decorator
+
+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 Settings, init
+from beamlit.errors import UnexpectedStatus
+from beamlit.functions import get_functions
+from beamlit.models import Agent, AgentSpec, Metadata
+
+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,
+    override_agent=None,
+    override_functions=None,
+    remote_functions=None,
+    local_functions=None,
+) -> Callable:
+    """
+    A decorator factory that configures and wraps functions to integrate with Beamlit agents.
+    Handles model initialization, function retrieval, and agent setup.
+
+    Parameters:
+        agent (Agent | dict, optional): An `Agent` instance or a dictionary containing agent metadata and specifications.
+        override_model (Any, optional): An optional model to override the default agent model.
+        override_agent (Any, optional): An optional agent instance to override the default agent.
+        mcp_hub (Any, optional): An optional MCP hub configuration.
+        remote_functions (Any, optional): An optional list of remote functions to be integrated.
+        local_functions (Any, optional): An optional list of local functions to be integrated.
+
+    Returns:
+        Callable: A decorator that wraps the target function, injecting agent-related configurations and dependencies.
+
+    Behavior:
+        - Validates and initializes the agent configuration.
+        - Retrieves and sets up the appropriate chat model based on the agent's specifications.
+        - Retrieves functions from the specified directories or remote sources.
+        - Wraps the target function, injecting agent, model, and function dependencies as needed.
+        - Logs relevant information and handles exceptions during the setup process.
+
+    Raises:
+        ValueError: If required configurations such as the model are missing.
+        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"}})'
+            )
+
+
+        def wrapper(func):
+
+            agent_kwargs = any(
+                param.name == "agent"
+                for param in inspect.signature(func).parameters.values()
+            )
+            model_kwargs = any(
+                param.name == "model"
+                for param in inspect.signature(func).parameters.values()
+            )
+            functions_kwargs = any(
+                param.name == "functions"
+                for param in inspect.signature(func).parameters.values()
+            )
+
+            @functools.wraps(func)
+            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 await func(*args, **kwargs)
+
+            return wrapped
+
+        return wrapper
+    except Exception as e:
+        logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")
+        raise e

{beamlit-0.0.57rc111 → beamlit-0.0.57rc113}/src/beamlit/functions/common.py

@@ -34,7 +34,7 @@ from beamlit.models import AgentChain
 
 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 +139,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,
@@ -153,8 +153,8 @@ def get_functions(
                                         func = getattr(module, func_name)
                                         if settings.remote:
                                             toolkit = RemoteToolkit(client, slugify(func.__name__))
-                                            toolkit.initialize()
-                                            functions.extend(toolkit.get_tools())
+                                            await toolkit.initialize()
+                                            functions.extend(await toolkit.get_tools())
                                         else:
                                             if asyncio.iscoroutinefunction(func):
                                                 functions.append(
@@ -183,20 +183,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 +207,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-0.0.57rc111 → beamlit-0.0.57rc113}/src/beamlit/functions/local/local.py

@@ -23,7 +23,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 +34,6 @@ class LocalToolKit:
                     spec={
                         "configurations": {
                             "url": self.local_function['url'],
-                            "sse": self.local_function['sse'],
                         },
                         "description": self.local_function['description'] or "",
                     }
@@ -42,8 +41,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-0.0.57rc113/src/beamlit/functions/mcp/client.py

@@ -0,0 +1,96 @@
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+from urllib.parse import urljoin, urlparse
+
+import anyio
+import mcp.types as types
+from anyio.abc import TaskStatus
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from websockets.client import WebSocketClientProtocol
+from websockets.client import connect as ws_connect
+
+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-0.0.57rc111 → beamlit-0.0.57rc113}/src/beamlit/functions/mcp/mcp.py

@@ -14,12 +14,12 @@ 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 .utils import create_schema_model
 
@@ -29,63 +29,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 +96,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 +111,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 +143,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 +160,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-0.0.57rc111 → beamlit-0.0.57rc113}/src/beamlit/functions/remote/remote.py

@@ -115,7 +115,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 +136,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 → beamlit-0.0.57rc113}/src/beamlit/run.py

@@ -5,6 +5,7 @@ import urllib.parse
 from typing import Any
 
 import requests
+
 from beamlit.client import AuthenticatedClient
 from beamlit.common import HTTPError, get_settings