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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -4,25 +4,123 @@ Defines decorators for agent functionalities.
4
4
  """
5
5
 
6
6
  # Import necessary modules
7
+ import asyncio
7
8
  import functools
8
9
  import inspect
9
10
  from logging import getLogger
10
11
  from typing import Callable
11
12
 
12
- from langgraph.checkpoint.memory import MemorySaver
13
- from langgraph.prebuilt import create_react_agent
14
-
15
13
  from beamlit.api.models import get_model, list_models
16
14
  from beamlit.authentication import new_client
17
- from beamlit.common.settings import init
15
+ from beamlit.common.settings import Settings, init
18
16
  from beamlit.errors import UnexpectedStatus
19
17
  from beamlit.functions import get_functions
20
18
  from beamlit.models import Agent, AgentSpec, Metadata
19
+ from langgraph.checkpoint.memory import MemorySaver
20
+ from langgraph.prebuilt import create_react_agent
21
21
 
22
22
  from .chat import get_chat_model_full
23
23
  from .voice.openai import OpenAIVoiceReactAgent
24
24
 
25
25
 
26
+ async def initialize_agent(
27
+ settings: Settings,
28
+ agent: Agent | dict = None,
29
+ override_model=None,
30
+ override_agent=None,
31
+ override_functions=None,
32
+ remote_functions=None,
33
+ local_functions=None,
34
+ ):
35
+ logger = getLogger(__name__)
36
+ client = new_client()
37
+ chat_model = override_model or None
38
+
39
+ if agent is not None:
40
+ metadata = Metadata(**agent.get("metadata", {}))
41
+ spec = AgentSpec(**agent.get("spec", {}))
42
+ agent = Agent(metadata=metadata, spec=spec)
43
+ if agent.spec.model and chat_model is None:
44
+ try:
45
+ response = get_model.sync_detailed(
46
+ agent.spec.model, client=client
47
+ )
48
+ settings.agent.model = response.parsed
49
+ except UnexpectedStatus as e:
50
+ if e.status_code == 404:
51
+ if e.status_code == 404:
52
+ raise ValueError(f"Model {agent.spec.model} not found")
53
+ raise e
54
+ except Exception as e:
55
+ raise e
56
+
57
+ if settings.agent.model:
58
+ chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
59
+ settings.agent.chat_model = chat_model
60
+ logger.info(f"Chat model configured, using: {provider}:{model}")
61
+
62
+ if override_functions is not None:
63
+ functions = override_functions
64
+ else:
65
+ functions = await get_functions(
66
+ client=client,
67
+ dir=settings.agent.functions_directory,
68
+ remote_functions=remote_functions,
69
+ chain=agent.spec.agent_chain,
70
+ local_functions=local_functions,
71
+ remote_functions_empty=not remote_functions,
72
+ warning=chat_model is not None,
73
+ )
74
+ settings.agent.functions = functions
75
+
76
+ if override_agent is None:
77
+ if chat_model is None:
78
+ models_select = ""
79
+ try:
80
+ models = list_models.sync_detailed(
81
+ client=client
82
+ )
83
+ models = ", ".join([model.metadata.name for model in models.parsed])
84
+ models_select = f"You can select one from your models: {models}"
85
+ except Exception:
86
+ pass
87
+
88
+ raise ValueError(
89
+ f"You must provide a model.\n"
90
+ f"{models_select}\n"
91
+ f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
92
+ "Add it to your agent spec\n"
93
+ "agent={\n"
94
+ ' "metadata": {\n'
95
+ f' "name": "{agent.metadata.name}",\n'
96
+ " },\n"
97
+ ' "spec": {\n'
98
+ ' "model": "MODEL_NAME",\n'
99
+ f' "description": "{agent.spec.description}",\n'
100
+ f' "prompt": "{agent.spec.prompt}",\n'
101
+ " },\n"
102
+ "}")
103
+ if isinstance(chat_model, OpenAIVoiceReactAgent):
104
+ _agent = chat_model
105
+ else:
106
+ memory = MemorySaver()
107
+ if len(functions) == 0:
108
+ raise ValueError("You can define this function in directory "
109
+ f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
110
+ "from beamlit.functions import function\n\n"
111
+ "@function()\n"
112
+ "def hello_world(query: str):\n"
113
+ " return 'Hello, world!'\n")
114
+ try:
115
+ _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
116
+ except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
117
+ logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
118
+ _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
119
+
120
+ settings.agent.agent = _agent
121
+ else:
122
+ settings.agent.agent = override_agent
123
+
26
124
  def agent(
27
125
  agent: Agent | dict = None,
28
126
  override_model=None,
@@ -58,17 +156,17 @@ def agent(
58
156
  Re-raises exceptions encountered during model retrieval and agent setup.
59
157
  """
60
158
  logger = getLogger(__name__)
159
+ settings = init()
160
+ _is_initialized = False
61
161
  try:
62
162
  if agent is not None and not isinstance(agent, dict):
63
163
  raise Exception(
64
164
  'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
65
165
  )
66
166
 
67
- client = new_client()
68
- chat_model = override_model or None
69
- settings = init()
70
167
 
71
168
  def wrapper(func):
169
+
72
170
  agent_kwargs = any(
73
171
  param.name == "agent"
74
172
  for param in inspect.signature(func).parameters.values()
@@ -83,99 +181,23 @@ def agent(
83
181
  )
84
182
 
85
183
  @functools.wraps(func)
86
- def wrapped(*args, **kwargs):
184
+ async def wrapped(*args, **kwargs):
185
+ nonlocal _is_initialized
186
+ if not _is_initialized:
187
+ async with asyncio.Lock():
188
+ if not _is_initialized:
189
+ await initialize_agent(settings, agent, override_model, override_agent, override_functions, remote_functions, local_functions)
190
+ _is_initialized = True
87
191
  if agent_kwargs:
88
192
  kwargs["agent"] = settings.agent.agent
89
193
  if model_kwargs:
90
194
  kwargs["model"] = settings.agent.chat_model
91
195
  if functions_kwargs:
92
196
  kwargs["functions"] = settings.agent.functions
93
- return func(*args, **kwargs)
197
+ return await func(*args, **kwargs)
94
198
 
95
199
  return wrapped
96
200
 
97
- if agent is not None:
98
- metadata = Metadata(**agent.get("metadata", {}))
99
- spec = AgentSpec(**agent.get("spec", {}))
100
- agent = Agent(metadata=metadata, spec=spec)
101
- if agent.spec.model and chat_model is None:
102
- try:
103
- response = get_model.sync_detailed(
104
- agent.spec.model, client=client
105
- )
106
- settings.agent.model = response.parsed
107
- except UnexpectedStatus as e:
108
- raise e
109
- except Exception as e:
110
- raise e
111
-
112
- if settings.agent.model:
113
- chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
114
- settings.agent.chat_model = chat_model
115
- logger.info(f"Chat model configured, using: {provider}:{model}")
116
-
117
- if override_functions is not None:
118
- functions = override_functions
119
- else:
120
- functions = get_functions(
121
- client=client,
122
- dir=settings.agent.functions_directory,
123
- remote_functions=remote_functions,
124
- chain=agent.spec.agent_chain,
125
- local_functions=local_functions,
126
- remote_functions_empty=not remote_functions,
127
- warning=chat_model is not None,
128
- )
129
-
130
- settings.agent.functions = functions
131
-
132
- if override_agent is None:
133
- if chat_model is None:
134
- models_select = ""
135
- try:
136
- models = list_models.sync_detailed(
137
- client=client
138
- )
139
- models = ", ".join([model.metadata.name for model in models.parsed])
140
- models_select = f"You can select one from your models: {models}"
141
- except Exception:
142
- pass
143
-
144
- raise ValueError(
145
- f"You must provide a model.\n"
146
- f"{models_select}\n"
147
- f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
148
- "Add it to your agent spec\n"
149
- "agent={\n"
150
- ' "metadata": {\n'
151
- f' "name": "{agent.metadata.name}",\n'
152
- " },\n"
153
- ' "spec": {\n'
154
- ' "model": "MODEL_NAME",\n'
155
- f' "description": "{agent.spec.description}",\n'
156
- f' "prompt": "{agent.spec.prompt}",\n'
157
- " },\n"
158
- "}")
159
- if isinstance(chat_model, OpenAIVoiceReactAgent):
160
- _agent = chat_model
161
- else:
162
- memory = MemorySaver()
163
- if len(functions) == 0:
164
- raise ValueError("You can define this function in directory "
165
- f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
166
- "from beamlit.functions import function\n\n"
167
- "@function()\n"
168
- "def hello_world(query: str):\n"
169
- " return 'Hello, world!'\n")
170
- try:
171
- _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
172
- except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
173
- logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
174
- _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
175
-
176
- settings.agent.agent = _agent
177
- else:
178
- settings.agent.agent = override_agent
179
201
  return wrapper
180
202
  except Exception as e:
181
203
  logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")
@@ -21,9 +21,6 @@ import traceback
21
21
  from logging import getLogger
22
22
  from typing import Union
23
23
 
24
- from langchain_core.tools import StructuredTool
25
- from langchain_core.tools.base import create_schema_from_function
26
-
27
24
  from beamlit.authentication import new_client
28
25
  from beamlit.client import AuthenticatedClient
29
26
  from beamlit.common import slugify
@@ -31,10 +28,12 @@ from beamlit.common.settings import get_settings
31
28
  from beamlit.functions.local.local import LocalToolKit
32
29
  from beamlit.functions.remote.remote import RemoteToolkit
33
30
  from beamlit.models import AgentChain
31
+ from langchain_core.tools import StructuredTool
32
+ from langchain_core.tools.base import create_schema_from_function
34
33
 
35
34
  logger = getLogger(__name__)
36
35
 
37
- def get_functions(
36
+ async def get_functions(
38
37
  remote_functions: Union[list[str], None] = None,
39
38
  local_functions: Union[list[dict], None] = None,
40
39
  client: Union[AuthenticatedClient, None] = None,
@@ -139,7 +138,7 @@ def get_functions(
139
138
  ):
140
139
  is_kit = keyword.value.value
141
140
  if is_kit and not settings.remote:
142
- kit_functions = get_functions(
141
+ kit_functions = await get_functions(
143
142
  client=client,
144
143
  dir=os.path.join(root),
145
144
  remote_functions_empty=remote_functions_empty,
@@ -183,20 +182,21 @@ def get_functions(
183
182
  for function in remote_functions:
184
183
  try:
185
184
  toolkit = RemoteToolkit(client, function)
186
- toolkit.initialize()
187
- functions.extend(toolkit.get_tools())
185
+ await toolkit.initialize()
186
+ functions.extend(await toolkit.get_tools())
188
187
  except Exception as e:
189
- logger.debug(
190
- f"Failed to initialize remote function {function}: {e!s}\n"
191
- f"Traceback:\n{traceback.format_exc()}"
192
- )
193
- logger.warn(f"Failed to initialize remote function {function}: {e!s}")
188
+ if not isinstance(e, RuntimeError):
189
+ logger.debug(
190
+ f"Failed to initialize remote function {function}: {e!s}\n"
191
+ f"Traceback:\n{traceback.format_exc()}"
192
+ )
193
+ logger.warn(f"Failed to initialize remote function {function}: {e!s}")
194
194
  if local_functions:
195
195
  for function in local_functions:
196
196
  try:
197
197
  toolkit = LocalToolKit(client, function)
198
- toolkit.initialize()
199
- functions.extend(toolkit.get_tools())
198
+ await toolkit.initialize()
199
+ functions.extend(await toolkit.get_tools())
200
200
  except Exception as e:
201
201
  logger.debug(
202
202
  f"Failed to initialize local function {function}: {e!s}\n"
@@ -206,8 +206,7 @@ def get_functions(
206
206
 
207
207
  if chain:
208
208
  toolkit = ChainToolkit(client, chain)
209
- toolkit.initialize()
210
- functions.extend(toolkit.get_tools())
211
-
209
+ await toolkit.initialize()
210
+ functions.extend(await toolkit.get_tools())
212
211
  return functions
213
212
 
@@ -4,9 +4,8 @@ import functools
4
4
  from collections.abc import Callable
5
5
  from logging import getLogger
6
6
 
7
- from fastapi import Request
8
-
9
7
  from beamlit.models import Function, FunctionKit
8
+ from fastapi import Request
10
9
 
11
10
  logger = getLogger(__name__)
12
11
 
@@ -1,11 +1,10 @@
1
1
  from dataclasses import dataclass
2
2
 
3
3
  import pydantic
4
- from langchain_core.tools.base import BaseTool
5
-
6
4
  from beamlit.authentication.authentication import AuthenticatedClient
7
5
  from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
8
6
  from beamlit.models import Function
7
+ from langchain_core.tools.base import BaseTool
9
8
 
10
9
 
11
10
  @dataclass
@@ -23,7 +22,7 @@ class LocalToolKit:
23
22
  _function: Function | None = None
24
23
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
25
24
 
26
- def initialize(self) -> None:
25
+ async def initialize(self) -> None:
27
26
  """Initialize the session and retrieve the local function details."""
28
27
  if self._function is None:
29
28
  try:
@@ -34,7 +33,6 @@ class LocalToolKit:
34
33
  spec={
35
34
  "configurations": {
36
35
  "url": self.local_function['url'],
37
- "sse": self.local_function['sse'],
38
36
  },
39
37
  "description": self.local_function['description'] or "",
40
38
  }
@@ -42,8 +40,8 @@ class LocalToolKit:
42
40
  except Exception as e:
43
41
  raise RuntimeError(f"Failed to initialize local function: {e}")
44
42
 
45
- def get_tools(self) -> list[BaseTool]:
43
+ async def get_tools(self) -> list[BaseTool]:
46
44
  mcp_client = MCPClient(self.client, self._function.spec["configurations"]["url"], sse=self._function.spec["configurations"]["sse"])
47
45
  mcp_toolkit = MCPToolkit(client=mcp_client)
48
- mcp_toolkit.initialize()
49
- return mcp_toolkit.get_tools()
46
+ await mcp_toolkit.initialize()
47
+ return await mcp_toolkit.get_tools()
@@ -0,0 +1,98 @@
1
+ import logging
2
+ from contextlib import asynccontextmanager
3
+ from typing import Any
4
+ from urllib.parse import urljoin, urlparse
5
+
6
+ import anyio
7
+ import httpx
8
+ from anyio.abc import TaskStatus
9
+ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
10
+ from websockets.client import connect as ws_connect
11
+ from websockets.client import WebSocketClientProtocol
12
+
13
+ import mcp.types as types
14
+
15
+ logger = logging.getLogger(__name__)
16
+
17
+
18
+ def remove_request_params(url: str) -> str:
19
+ return urljoin(url, urlparse(url).path)
20
+
21
+
22
+ @asynccontextmanager
23
+ async def websocket_client(
24
+ url: str,
25
+ headers: dict[str, Any] | None = None,
26
+ timeout: float = 5,
27
+ ):
28
+ """
29
+ Client transport for WebSocket.
30
+
31
+ The `timeout` parameter controls connection timeout.
32
+ """
33
+ read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
34
+ read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
35
+
36
+ write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
37
+ write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
38
+
39
+ read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
40
+ write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
41
+
42
+ async with anyio.create_task_group() as tg:
43
+ try:
44
+ # Convert http(s):// to ws(s)://
45
+ ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
46
+ logger.debug(f"Connecting to WebSocket endpoint: {remove_request_params(ws_url)}")
47
+
48
+ async with ws_connect(ws_url, extra_headers=headers, open_timeout=timeout) as websocket:
49
+ logger.debug("WebSocket connection established")
50
+
51
+ async def ws_reader(
52
+ websocket: WebSocketClientProtocol,
53
+ task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
54
+ ):
55
+ try:
56
+ task_status.started()
57
+ async for message in websocket:
58
+ logger.debug(f"Received WebSocket message: {message}")
59
+ try:
60
+ parsed_message = types.JSONRPCMessage.model_validate_json(message)
61
+ logger.debug(f"Received server message: {parsed_message}")
62
+ await read_stream_writer.send(parsed_message)
63
+ except Exception as exc:
64
+ logger.error(f"Error parsing server message: {exc}")
65
+ await read_stream_writer.send(exc)
66
+ except Exception as exc:
67
+ logger.error(f"Error in ws_reader: {exc}")
68
+ await read_stream_writer.send(exc)
69
+ finally:
70
+ await read_stream_writer.aclose()
71
+
72
+ async def ws_writer(websocket: WebSocketClientProtocol):
73
+ try:
74
+ async with write_stream_reader:
75
+ async for message in write_stream_reader:
76
+ logger.debug(f"Sending client message: {message}")
77
+ await websocket.send(
78
+ message.model_dump_json(
79
+ by_alias=True,
80
+ exclude_none=True,
81
+ )
82
+ )
83
+ logger.debug("Client message sent successfully")
84
+ except Exception as exc:
85
+ logger.error(f"Error in ws_writer: {exc}")
86
+ finally:
87
+ await write_stream.aclose()
88
+
89
+ await tg.start(ws_reader, websocket)
90
+ tg.start_soon(ws_writer, websocket)
91
+
92
+ try:
93
+ yield read_stream, write_stream
94
+ finally:
95
+ tg.cancel_scope.cancel()
96
+ finally:
97
+ await read_stream_writer.aclose()
98
+ await write_stream.aclose()
@@ -12,14 +12,13 @@ import pydantic
12
12
  import pydantic_core
13
13
  import requests
14
14
  import typing_extensions as t
15
- from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
16
- from mcp import ClientSession
17
- from mcp.client.sse import sse_client
18
- from mcp.types import CallToolResult, ListToolsResult
19
-
20
15
  from beamlit.authentication import get_authentication_headers
21
16
  from beamlit.authentication.authentication import AuthenticatedClient
22
17
  from beamlit.common.settings import get_settings
18
+ from beamlit.functions.mcp.client import websocket_client
19
+ from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
20
+ from mcp import ClientSession
21
+ from mcp.types import CallToolResult, ListToolsResult
23
22
 
24
23
  from .utils import create_schema_model
25
24
 
@@ -29,63 +28,65 @@ logger = logging.getLogger(__name__)
29
28
 
30
29
 
31
30
  class MCPClient:
32
- def __init__(self, client: AuthenticatedClient, url: str, sse: bool = False):
31
+ def __init__(self, client: AuthenticatedClient, url: str, fallback_url: str | None = None):
33
32
  self.client = client
34
33
  self.url = url
35
- self._sse = False
34
+ self.fallback_url = fallback_url
36
35
 
37
- async def list_sse_tools(self) -> ListToolsResult:
38
- # Create a new context for each SSE connection
36
+ async def list_ws_tools(self, is_fallback: bool = False) -> ListToolsResult:
37
+ if is_fallback:
38
+ url = self.fallback_url
39
+ else:
40
+ url = self.url
39
41
  try:
40
- async with sse_client(f"{self.url}/sse", headers=get_authentication_headers(settings)) as (read_stream, write_stream):
41
- async with ClientSession(read_stream, write_stream) as session:
42
- await session.initialize()
43
- response = await session.list_tools()
42
+ async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
43
+ logger.debug("WebSocket connection established")
44
+ async with ClientSession(read_stream, write_stream) as client:
45
+ await client.initialize()
46
+ response = await client.list_tools()
47
+ logger.debug(f"WebSocket tools: {response}")
44
48
  return response
45
- except Exception:
46
- self._sse = False
47
- logger.info("SSE not available, trying HTTP")
49
+ except Exception as e:
50
+ logger.error(f"Error listing SSE tools: {e}")
51
+ logger.debug("WebSocket not available, trying HTTP")
48
52
  return None # Signal to list_tools() to try HTTP instead
49
53
 
50
- def list_tools(self) -> ListToolsResult:
54
+ async def list_tools(self) -> ListToolsResult:
55
+ logger.debug(f"Listing tools for {self.url}")
51
56
  try:
52
- loop = asyncio.get_event_loop()
53
- result = loop.run_until_complete(self.list_sse_tools())
54
- if result is None: # SSE failed, try HTTP
55
- raise Exception("SSE failed")
56
- self._sse = True
57
+ result = await self.list_ws_tools(is_fallback=False)
57
58
  return result
58
- except Exception: # Fallback to HTTP
59
- client = self.client.get_httpx_client()
60
- response = client.request("GET", f"{self.url}/tools/list")
61
- response.raise_for_status()
62
- return ListToolsResult(**response.json())
59
+ except Exception as e: # Fallback to Public endpoint
60
+ if self.fallback_url:
61
+ try:
62
+ result = await self.list_ws_tools(is_fallback=True)
63
+ return result
64
+ except Exception as e:
65
+ raise e
66
+ else:
67
+ raise e
68
+
63
69
 
64
70
  async def call_tool(
65
71
  self,
66
72
  tool_name: str,
67
73
  arguments: dict[str, Any] = None,
74
+ is_fallback: bool = False,
68
75
  ) -> requests.Response | AsyncIterator[CallToolResult]:
69
- if self._sse:
70
- async with sse_client(f"{self.url}/sse", headers=get_authentication_headers(settings)) as (read_stream, write_stream):
76
+ if is_fallback:
77
+ url = self.fallback_url
78
+ else:
79
+ url = self.url
80
+ try:
81
+ async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
71
82
  async with ClientSession(read_stream, write_stream) as session:
72
83
  await session.initialize()
73
84
  response = await session.call_tool(tool_name, arguments or {})
74
85
  content = pydantic_core.to_json(response).decode()
75
86
  return content
76
- else: # Fallback to HTTP
77
- client = self.client.get_httpx_client()
78
- response = client.request(
79
- "POST",
80
- f"{self.url}/tools/call",
81
- json={"name": tool_name, "arguments": arguments},
82
- )
83
- response.raise_for_status()
84
- result = CallToolResult(response.json())
85
- if result.isError:
86
- raise ToolException(result.content)
87
- content = pydantic_core.to_json(result.content).decode()
88
- return content
87
+ except Exception as e:
88
+ raise e
89
+
89
90
 
90
91
  class MCPTool(BaseTool):
91
92
  """
@@ -94,7 +95,6 @@ class MCPTool(BaseTool):
94
95
  Attributes:
95
96
  client (MCPClient): The MCP client instance.
96
97
  handle_tool_error (bool | str | Callable[[ToolException], str] | None): Error handling strategy.
97
- sse (bool): Whether to use SSE streaming for responses.
98
98
  """
99
99
 
100
100
  client: MCPClient
@@ -110,7 +110,16 @@ class MCPTool(BaseTool):
110
110
 
111
111
  @t.override
112
112
  async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
113
- return await self.client.call_tool(self.name, arguments=kwargs)
113
+ try:
114
+ return await self.client.call_tool(self.name, arguments=kwargs)
115
+ except Exception as e:
116
+ if self.client.fallback_url:
117
+ try:
118
+ return await self.client.call_tool(self.name, arguments=kwargs, is_fallback=True) # Fallback to Public endpoint
119
+ except Exception as e:
120
+ raise e
121
+ else:
122
+ raise e
114
123
 
115
124
  @t.override
116
125
  @property
@@ -133,14 +142,14 @@ class MCPToolkit(BaseToolkit):
133
142
  _tools: ListToolsResult | None = None
134
143
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
135
144
 
136
- def initialize(self) -> None:
145
+ async def initialize(self) -> None:
137
146
  """Initialize the session and retrieve tools list"""
138
147
  if self._tools is None:
139
- response = self.client.list_tools()
148
+ response = await self.client.list_tools()
140
149
  self._tools = response
141
150
 
142
151
  @t.override
143
- def get_tools(self) -> list[BaseTool]:
152
+ async def get_tools(self) -> list[BaseTool]:
144
153
  if self._tools is None:
145
154
  raise RuntimeError("Must initialize the toolkit first")
146
155
 
@@ -150,7 +159,6 @@ class MCPToolkit(BaseToolkit):
150
159
  name=tool.name,
151
160
  description=tool.description or "",
152
161
  args_schema=create_schema_model(tool.name, tool.inputSchema),
153
- sse=self.sse,
154
162
  )
155
163
  # list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
156
164
  for tool in self._tools.tools
@@ -11,8 +11,6 @@ from typing import Callable
11
11
 
12
12
  import pydantic
13
13
  import typing_extensions as t
14
- from langchain_core.tools.base import BaseTool, ToolException
15
-
16
14
  from beamlit.api.functions import get_function, list_functions
17
15
  from beamlit.authentication.authentication import AuthenticatedClient
18
16
  from beamlit.common.settings import get_settings
@@ -20,6 +18,7 @@ from beamlit.errors import UnexpectedStatus
20
18
  from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
21
19
  from beamlit.models import Function, StoreFunctionParameter
22
20
  from beamlit.run import RunClient
21
+ from langchain_core.tools.base import BaseTool, ToolException
23
22
 
24
23
 
25
24
  def create_dynamic_schema(name: str, parameters: list[StoreFunctionParameter]) -> type[pydantic.BaseModel]:
@@ -115,7 +114,7 @@ class RemoteToolkit:
115
114
  _service_name: str | None = None
116
115
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
117
116
 
118
- def initialize(self) -> None:
117
+ async def initialize(self) -> None:
119
118
  """Initialize the session and retrieve the remote function details."""
120
119
  if self._function is None:
121
120
  try:
@@ -136,19 +135,21 @@ class RemoteToolkit:
136
135
  f"error: {e.status_code}. Available functions: {', '.join(names)}"
137
136
  )
138
137
 
139
- def get_tools(self) -> list[BaseTool]:
138
+ async def get_tools(self) -> list[BaseTool]:
140
139
  settings = get_settings()
141
140
  if self._function is None:
142
141
  raise RuntimeError("Must initialize the toolkit first")
143
142
 
144
143
  if self._function.spec.integration_connections:
144
+ fallback_url = None
145
145
  url = f"{settings.run_url}/{settings.workspace}/functions/{self._function.metadata.name}"
146
146
  if self._service_name:
147
+ fallback_url = f"https://{self._service_name}.{settings.run_internal_hostname}"
147
148
  url = f"https://{self._service_name}.{settings.run_internal_hostname}"
148
- mcp_client = MCPClient(self.client, url)
149
- mcp_toolkit = MCPToolkit(client=mcp_client)
150
- mcp_toolkit.initialize()
151
- return mcp_toolkit.get_tools()
149
+ mcp_client = MCPClient(self.client, url, fallback_url)
150
+ mcp_toolkit = MCPToolkit(client=mcp_client, url=url)
151
+ await mcp_toolkit.initialize()
152
+ return await mcp_toolkit.get_tools()
152
153
 
153
154
  if self._function.spec.kit:
154
155
  return [
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: beamlit
3
- Version: 0.0.57rc111
3
+ Version: 0.0.57rc112
4
4
  Summary: Add your description here
5
5
  Author-email: cploujoux <ch.ploujoux@gmail.com>
6
6
  License-File: LICENSE
@@ -43,6 +43,7 @@ Requires-Dist: pyjwt>=2.10.1
43
43
  Requires-Dist: python-dateutil>=2.8.0
44
44
  Requires-Dist: pyyaml<6.1.0,>=6.0.2
45
45
  Requires-Dist: requests<2.33.0,>=2.32.3
46
+ Requires-Dist: websocket>=0.2.1
46
47
  Requires-Dist: websockets>=14.1
47
48
  Description-Content-Type: text/markdown
48
49
 
@@ -7,7 +7,7 @@ beamlit/types.py,sha256=E1hhDh_zXfsSQ0NCt9-uw90_Mr5iIlsdfnfvxv5HarU,1005
7
7
  beamlit/agents/__init__.py,sha256=bWsFaXUbAps3IsL3Prti89m1s714vICXodbQi77h3vY,206
8
8
  beamlit/agents/chain.py,sha256=JsinjAYBr3oaM4heouZaiaV2jMmi779LHAMtD_4P59s,4867
9
9
  beamlit/agents/chat.py,sha256=ufuydptucLNe_Jyr7lQO1WfQ5pe0I5YKh-y0smwWABM,8301
10
- beamlit/agents/decorator.py,sha256=oXhaLXeyuw3gOoiDFoszoilJgJLo9hX9WHL30DSw5go,7754
10
+ beamlit/agents/decorator.py,sha256=pKA9NIaNoToCsjMkvBpOvfQUsb0B20XwdskTonz-isI,8310
11
11
  beamlit/agents/thread.py,sha256=PVP9Gey8fMZHA-Cs8WbfmdSlD7g-n4HKuk1sTVf8yOQ,1087
12
12
  beamlit/agents/voice/openai.py,sha256=-RDBwl16i4TbUhFo5-77Ci3zmI3Y8U2yf2MmvXR2haQ,9479
13
13
  beamlit/agents/voice/utils.py,sha256=tQidyM40Ewuy12wKqpvJLvfJgneQ0sZf50dqnerPGHg,836
@@ -131,12 +131,13 @@ beamlit/deploy/deploy.py,sha256=qpNXHvLj712Uo34MPT8z0AkPNsG03lEPe08p1XpGQhc,1116
131
131
  beamlit/deploy/format.py,sha256=W3ESUHyFv-iZDjVnHOf9YFDDXZSXYIFFbwCoL1GInE0,1162
132
132
  beamlit/deploy/parser.py,sha256=gjRUhOVtfKnc1UNc_FhXsEfj9zrMNuq8W93pNsJBpo0,7586
133
133
  beamlit/functions/__init__.py,sha256=Mnoqpa1dm7TXwjodBbF_40JyD78aXsOYWmqjDSnA1lU,317
134
- beamlit/functions/common.py,sha256=4SK1N0VoQiydUT4frGT9OWMRxgHgMHUcsBhicMkcMmc,9996
135
- beamlit/functions/decorator.py,sha256=iQbLwUo0K83DFJ3ub8O5jKtkbSINnku6GZcKJ9h7-5E,2292
136
- beamlit/functions/local/local.py,sha256=KjkHWBxGlG9fliXnOOwZQEop1g2o10IaLiPo4Zu_XAk,1929
137
- beamlit/functions/mcp/mcp.py,sha256=itd7i5Ej0WSBmw8Z6pbdUL1-NJwDoIQ1edeFyg1o4Co,5814
134
+ beamlit/functions/common.py,sha256=uPmFx8odqHVas8kSfc41B9g_raWtAAzQBNyuRjtDDb0,10114
135
+ beamlit/functions/decorator.py,sha256=P17iKgDPe5VBEgwLPep1yVTD67C8YtRXuoe_kNEfBFg,2291
136
+ beamlit/functions/local/local.py,sha256=GQlp-dlEd465qMwJ-oSkVtLdDGtIYEE4PsVoAzbJJyg,1890
137
+ beamlit/functions/mcp/client.py,sha256=SwGINxSx9a3_LlDAiYuf8NNw8KZf8B5PIOls_wjJlkU,4086
138
+ beamlit/functions/mcp/mcp.py,sha256=T9dUc1PvYO1FXADLu0rZY4FIjUFFQAOuOkiZghLd8eo,5949
138
139
  beamlit/functions/mcp/utils.py,sha256=V7bah6cymdtjJ_LJUrNcHDeApDHA6uXvaGVeFJGKj2U,1850
139
- beamlit/functions/remote/remote.py,sha256=iB8DmJBUj-NchQoPlAPTNLASywIWYlkSlqvHDl-ACYE,6553
140
+ beamlit/functions/remote/remote.py,sha256=O7zq315jHkcvNbrd4AZfsxn4PVdqGvrecKxBR-EwV9U,6727
140
141
  beamlit/models/__init__.py,sha256=042wT7sG_YUmJJAb9edrui-DesMxLjOGVvnEtqc92CY,8916
141
142
  beamlit/models/acl.py,sha256=tH67gsl_BMaviSbTaaIkO1g9cWZgJ6VgAnYVjQSzGZY,3952
142
143
  beamlit/models/agent.py,sha256=DPZZOrDyOAXlh7uvPmhK1jgz8L4UWmaF_SH-OnAdjlw,4162
@@ -253,8 +254,8 @@ beamlit/serve/app.py,sha256=5XZci-R95Zjl97wMtQd1BRtonnkJJ2AeoTVFPKGAOfA,4283
253
254
  beamlit/serve/middlewares/__init__.py,sha256=O7fyfE1DIYmajFY9WWdzxCgeAQWZzJfeUjzHGbpWaAk,309
254
255
  beamlit/serve/middlewares/accesslog.py,sha256=lcu33j4epFSHRBaeTpyt8deNb3kaM3K91-andw4fp80,1112
255
256
  beamlit/serve/middlewares/processtime.py,sha256=3x5w1yQexB0xFNKK6fgLbINxT-eLLunfZ6UDV0bIIF4,944
256
- beamlit-0.0.57rc111.dist-info/METADATA,sha256=5HLIcsxG56ZK1gCw97tpXUpjnlJ2iSlQAk-gs8Na2wg,3515
257
- beamlit-0.0.57rc111.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
258
- beamlit-0.0.57rc111.dist-info/entry_points.txt,sha256=zxhgdn7SP-Otk4rEv7LMPAAa9w4TUCLbu9TJi9-K3xg,115
259
- beamlit-0.0.57rc111.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
260
- beamlit-0.0.57rc111.dist-info/RECORD,,
257
+ beamlit-0.0.57rc112.dist-info/METADATA,sha256=KHt8c7Zlqs1EnvWmSVAEYgO3hasTEWMbR8AnnSmUIw0,3547
258
+ beamlit-0.0.57rc112.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
259
+ beamlit-0.0.57rc112.dist-info/entry_points.txt,sha256=zxhgdn7SP-Otk4rEv7LMPAAa9w4TUCLbu9TJi9-K3xg,115
260
+ beamlit-0.0.57rc112.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
261
+ beamlit-0.0.57rc112.dist-info/RECORD,,