beamlit 0.0.57rc111__py3-none-any.whl → 0.0.57rc113__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,6 +4,7 @@ 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
@@ -14,7 +15,7 @@ from langgraph.prebuilt import create_react_agent
14
15
 
15
16
  from beamlit.api.models import get_model, list_models
16
17
  from beamlit.authentication import new_client
17
- from beamlit.common.settings import init
18
+ from beamlit.common.settings import Settings, init
18
19
  from beamlit.errors import UnexpectedStatus
19
20
  from beamlit.functions import get_functions
20
21
  from beamlit.models import Agent, AgentSpec, Metadata
@@ -23,6 +24,104 @@ from .chat import get_chat_model_full
23
24
  from .voice.openai import OpenAIVoiceReactAgent
24
25
 
25
26
 
27
+ async def initialize_agent(
28
+ settings: Settings,
29
+ agent: Agent | dict = None,
30
+ override_model=None,
31
+ override_agent=None,
32
+ override_functions=None,
33
+ remote_functions=None,
34
+ local_functions=None,
35
+ ):
36
+ logger = getLogger(__name__)
37
+ client = new_client()
38
+ chat_model = override_model or None
39
+
40
+ if agent is not None:
41
+ metadata = Metadata(**agent.get("metadata", {}))
42
+ spec = AgentSpec(**agent.get("spec", {}))
43
+ agent = Agent(metadata=metadata, spec=spec)
44
+ if agent.spec.model and chat_model is None:
45
+ try:
46
+ response = get_model.sync_detailed(
47
+ agent.spec.model, client=client
48
+ )
49
+ settings.agent.model = response.parsed
50
+ except UnexpectedStatus as e:
51
+ if e.status_code == 404:
52
+ if e.status_code == 404:
53
+ raise ValueError(f"Model {agent.spec.model} not found")
54
+ raise e
55
+ except Exception as e:
56
+ raise e
57
+
58
+ if settings.agent.model:
59
+ chat_model, provider, model = get_chat_model_full(agent.spec.model, settings.agent.model)
60
+ settings.agent.chat_model = chat_model
61
+ logger.info(f"Chat model configured, using: {provider}:{model}")
62
+
63
+ if override_functions is not None:
64
+ functions = override_functions
65
+ else:
66
+ functions = await get_functions(
67
+ client=client,
68
+ dir=settings.agent.functions_directory,
69
+ remote_functions=remote_functions,
70
+ chain=agent.spec.agent_chain,
71
+ local_functions=local_functions,
72
+ remote_functions_empty=not remote_functions,
73
+ warning=chat_model is not None,
74
+ )
75
+ settings.agent.functions = functions
76
+
77
+ if override_agent is None:
78
+ if chat_model is None:
79
+ models_select = ""
80
+ try:
81
+ models = list_models.sync_detailed(
82
+ client=client
83
+ )
84
+ models = ", ".join([model.metadata.name for model in models.parsed])
85
+ models_select = f"You can select one from your models: {models}"
86
+ except Exception:
87
+ pass
88
+
89
+ raise ValueError(
90
+ f"You must provide a model.\n"
91
+ f"{models_select}\n"
92
+ f"You can create one at {settings.app_url}/{settings.workspace}/global-inference-network/models/create\n"
93
+ "Add it to your agent spec\n"
94
+ "agent={\n"
95
+ ' "metadata": {\n'
96
+ f' "name": "{agent.metadata.name}",\n'
97
+ " },\n"
98
+ ' "spec": {\n'
99
+ ' "model": "MODEL_NAME",\n'
100
+ f' "description": "{agent.spec.description}",\n'
101
+ f' "prompt": "{agent.spec.prompt}",\n'
102
+ " },\n"
103
+ "}")
104
+ if isinstance(chat_model, OpenAIVoiceReactAgent):
105
+ _agent = chat_model
106
+ else:
107
+ memory = MemorySaver()
108
+ if len(functions) == 0:
109
+ raise ValueError("You can define this function in directory "
110
+ f'"{settings.agent.functions_directory}". Here is a sample function you can use:\n\n'
111
+ "from beamlit.functions import function\n\n"
112
+ "@function()\n"
113
+ "def hello_world(query: str):\n"
114
+ " return 'Hello, world!'\n")
115
+ try:
116
+ _agent = create_react_agent(chat_model, functions, checkpointer=memory, state_modifier=agent.spec.prompt or "")
117
+ except AttributeError: # special case for azure-marketplace where it uses the old OpenAI interface (no tools)
118
+ logger.warning("Using the old OpenAI interface for Azure Marketplace, no tools available")
119
+ _agent = create_react_agent(chat_model, [], checkpointer=memory, state_modifier=agent.spec.prompt or "")
120
+
121
+ settings.agent.agent = _agent
122
+ else:
123
+ settings.agent.agent = override_agent
124
+
26
125
  def agent(
27
126
  agent: Agent | dict = None,
28
127
  override_model=None,
@@ -58,17 +157,17 @@ def agent(
58
157
  Re-raises exceptions encountered during model retrieval and agent setup.
59
158
  """
60
159
  logger = getLogger(__name__)
160
+ settings = init()
161
+ _is_initialized = False
61
162
  try:
62
163
  if agent is not None and not isinstance(agent, dict):
63
164
  raise Exception(
64
165
  'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
65
166
  )
66
167
 
67
- client = new_client()
68
- chat_model = override_model or None
69
- settings = init()
70
168
 
71
169
  def wrapper(func):
170
+
72
171
  agent_kwargs = any(
73
172
  param.name == "agent"
74
173
  for param in inspect.signature(func).parameters.values()
@@ -83,99 +182,23 @@ def agent(
83
182
  )
84
183
 
85
184
  @functools.wraps(func)
86
- def wrapped(*args, **kwargs):
185
+ async def wrapped(*args, **kwargs):
186
+ nonlocal _is_initialized
187
+ if not _is_initialized:
188
+ async with asyncio.Lock():
189
+ if not _is_initialized:
190
+ await initialize_agent(settings, agent, override_model, override_agent, override_functions, remote_functions, local_functions)
191
+ _is_initialized = True
87
192
  if agent_kwargs:
88
193
  kwargs["agent"] = settings.agent.agent
89
194
  if model_kwargs:
90
195
  kwargs["model"] = settings.agent.chat_model
91
196
  if functions_kwargs:
92
197
  kwargs["functions"] = settings.agent.functions
93
- return func(*args, **kwargs)
198
+ return await func(*args, **kwargs)
94
199
 
95
200
  return wrapped
96
201
 
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
202
  return wrapper
180
203
  except Exception as e:
181
204
  logger.error(f"Error in agent decorator: {e!s} at line {e.__traceback__.tb_lineno}")
@@ -34,7 +34,7 @@ from beamlit.models import AgentChain
34
34
 
35
35
  logger = getLogger(__name__)
36
36
 
37
- def get_functions(
37
+ async def get_functions(
38
38
  remote_functions: Union[list[str], None] = None,
39
39
  local_functions: Union[list[dict], None] = None,
40
40
  client: Union[AuthenticatedClient, None] = None,
@@ -139,7 +139,7 @@ def get_functions(
139
139
  ):
140
140
  is_kit = keyword.value.value
141
141
  if is_kit and not settings.remote:
142
- kit_functions = get_functions(
142
+ kit_functions = await get_functions(
143
143
  client=client,
144
144
  dir=os.path.join(root),
145
145
  remote_functions_empty=remote_functions_empty,
@@ -153,8 +153,8 @@ def get_functions(
153
153
  func = getattr(module, func_name)
154
154
  if settings.remote:
155
155
  toolkit = RemoteToolkit(client, slugify(func.__name__))
156
- toolkit.initialize()
157
- functions.extend(toolkit.get_tools())
156
+ await toolkit.initialize()
157
+ functions.extend(await toolkit.get_tools())
158
158
  else:
159
159
  if asyncio.iscoroutinefunction(func):
160
160
  functions.append(
@@ -183,20 +183,21 @@ def get_functions(
183
183
  for function in remote_functions:
184
184
  try:
185
185
  toolkit = RemoteToolkit(client, function)
186
- toolkit.initialize()
187
- functions.extend(toolkit.get_tools())
186
+ await toolkit.initialize()
187
+ functions.extend(await toolkit.get_tools())
188
188
  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}")
189
+ if not isinstance(e, RuntimeError):
190
+ logger.debug(
191
+ f"Failed to initialize remote function {function}: {e!s}\n"
192
+ f"Traceback:\n{traceback.format_exc()}"
193
+ )
194
+ logger.warn(f"Failed to initialize remote function {function}: {e!s}")
194
195
  if local_functions:
195
196
  for function in local_functions:
196
197
  try:
197
198
  toolkit = LocalToolKit(client, function)
198
- toolkit.initialize()
199
- functions.extend(toolkit.get_tools())
199
+ await toolkit.initialize()
200
+ functions.extend(await toolkit.get_tools())
200
201
  except Exception as e:
201
202
  logger.debug(
202
203
  f"Failed to initialize local function {function}: {e!s}\n"
@@ -206,8 +207,7 @@ def get_functions(
206
207
 
207
208
  if chain:
208
209
  toolkit = ChainToolkit(client, chain)
209
- toolkit.initialize()
210
- functions.extend(toolkit.get_tools())
211
-
210
+ await toolkit.initialize()
211
+ functions.extend(await toolkit.get_tools())
212
212
  return functions
213
213
 
@@ -23,7 +23,7 @@ class LocalToolKit:
23
23
  _function: Function | None = None
24
24
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
25
25
 
26
- def initialize(self) -> None:
26
+ async def initialize(self) -> None:
27
27
  """Initialize the session and retrieve the local function details."""
28
28
  if self._function is None:
29
29
  try:
@@ -34,7 +34,6 @@ class LocalToolKit:
34
34
  spec={
35
35
  "configurations": {
36
36
  "url": self.local_function['url'],
37
- "sse": self.local_function['sse'],
38
37
  },
39
38
  "description": self.local_function['description'] or "",
40
39
  }
@@ -42,8 +41,8 @@ class LocalToolKit:
42
41
  except Exception as e:
43
42
  raise RuntimeError(f"Failed to initialize local function: {e}")
44
43
 
45
- def get_tools(self) -> list[BaseTool]:
44
+ async def get_tools(self) -> list[BaseTool]:
46
45
  mcp_client = MCPClient(self.client, self._function.spec["configurations"]["url"], sse=self._function.spec["configurations"]["sse"])
47
46
  mcp_toolkit = MCPToolkit(client=mcp_client)
48
- mcp_toolkit.initialize()
49
- return mcp_toolkit.get_tools()
47
+ await mcp_toolkit.initialize()
48
+ return await mcp_toolkit.get_tools()
@@ -0,0 +1,96 @@
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 mcp.types as types
8
+ from anyio.abc import TaskStatus
9
+ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
10
+ from websockets.client import WebSocketClientProtocol
11
+ from websockets.client import connect as ws_connect
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+
16
+ def remove_request_params(url: str) -> str:
17
+ return urljoin(url, urlparse(url).path)
18
+
19
+
20
+ @asynccontextmanager
21
+ async def websocket_client(
22
+ url: str,
23
+ headers: dict[str, Any] | None = None,
24
+ timeout: float = 5,
25
+ ):
26
+ """
27
+ Client transport for WebSocket.
28
+
29
+ The `timeout` parameter controls connection timeout.
30
+ """
31
+ read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
32
+ read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
33
+
34
+ write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
35
+ write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
36
+
37
+ read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
38
+ write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
39
+
40
+ async with anyio.create_task_group() as tg:
41
+ try:
42
+ # Convert http(s):// to ws(s)://
43
+ ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
44
+ logger.debug(f"Connecting to WebSocket endpoint: {remove_request_params(ws_url)}")
45
+
46
+ async with ws_connect(ws_url, extra_headers=headers, open_timeout=timeout) as websocket:
47
+ logger.debug("WebSocket connection established")
48
+
49
+ async def ws_reader(
50
+ websocket: WebSocketClientProtocol,
51
+ task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
52
+ ):
53
+ try:
54
+ task_status.started()
55
+ async for message in websocket:
56
+ logger.debug(f"Received WebSocket message: {message}")
57
+ try:
58
+ parsed_message = types.JSONRPCMessage.model_validate_json(message)
59
+ logger.debug(f"Received server message: {parsed_message}")
60
+ await read_stream_writer.send(parsed_message)
61
+ except Exception as exc:
62
+ logger.error(f"Error parsing server message: {exc}")
63
+ await read_stream_writer.send(exc)
64
+ except Exception as exc:
65
+ logger.error(f"Error in ws_reader: {exc}")
66
+ await read_stream_writer.send(exc)
67
+ finally:
68
+ await read_stream_writer.aclose()
69
+
70
+ async def ws_writer(websocket: WebSocketClientProtocol):
71
+ try:
72
+ async with write_stream_reader:
73
+ async for message in write_stream_reader:
74
+ logger.debug(f"Sending client message: {message}")
75
+ await websocket.send(
76
+ message.model_dump_json(
77
+ by_alias=True,
78
+ exclude_none=True,
79
+ )
80
+ )
81
+ logger.debug("Client message sent successfully")
82
+ except Exception as exc:
83
+ logger.error(f"Error in ws_writer: {exc}")
84
+ finally:
85
+ await write_stream.aclose()
86
+
87
+ await tg.start(ws_reader, websocket)
88
+ tg.start_soon(ws_writer, websocket)
89
+
90
+ try:
91
+ yield read_stream, write_stream
92
+ finally:
93
+ tg.cancel_scope.cancel()
94
+ finally:
95
+ await read_stream_writer.aclose()
96
+ await write_stream.aclose()
@@ -14,12 +14,12 @@ import requests
14
14
  import typing_extensions as t
15
15
  from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
16
16
  from mcp import ClientSession
17
- from mcp.client.sse import sse_client
18
17
  from mcp.types import CallToolResult, ListToolsResult
19
18
 
20
19
  from beamlit.authentication import get_authentication_headers
21
20
  from beamlit.authentication.authentication import AuthenticatedClient
22
21
  from beamlit.common.settings import get_settings
22
+ from beamlit.functions.mcp.client import websocket_client
23
23
 
24
24
  from .utils import create_schema_model
25
25
 
@@ -29,63 +29,65 @@ logger = logging.getLogger(__name__)
29
29
 
30
30
 
31
31
  class MCPClient:
32
- def __init__(self, client: AuthenticatedClient, url: str, sse: bool = False):
32
+ def __init__(self, client: AuthenticatedClient, url: str, fallback_url: str | None = None):
33
33
  self.client = client
34
34
  self.url = url
35
- self._sse = False
35
+ self.fallback_url = fallback_url
36
36
 
37
- async def list_sse_tools(self) -> ListToolsResult:
38
- # Create a new context for each SSE connection
37
+ async def list_ws_tools(self, is_fallback: bool = False) -> ListToolsResult:
38
+ if is_fallback:
39
+ url = self.fallback_url
40
+ else:
41
+ url = self.url
39
42
  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()
43
+ async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
44
+ logger.debug("WebSocket connection established")
45
+ async with ClientSession(read_stream, write_stream) as client:
46
+ await client.initialize()
47
+ response = await client.list_tools()
48
+ logger.debug(f"WebSocket tools: {response}")
44
49
  return response
45
- except Exception:
46
- self._sse = False
47
- logger.info("SSE not available, trying HTTP")
50
+ except Exception as e:
51
+ logger.error(f"Error listing SSE tools: {e}")
52
+ logger.debug("WebSocket not available, trying HTTP")
48
53
  return None # Signal to list_tools() to try HTTP instead
49
54
 
50
- def list_tools(self) -> ListToolsResult:
55
+ async def list_tools(self) -> ListToolsResult:
56
+ logger.debug(f"Listing tools for {self.url}")
51
57
  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
58
+ result = await self.list_ws_tools(is_fallback=False)
57
59
  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())
60
+ except Exception as e: # Fallback to Public endpoint
61
+ if self.fallback_url:
62
+ try:
63
+ result = await self.list_ws_tools(is_fallback=True)
64
+ return result
65
+ except Exception as e:
66
+ raise e
67
+ else:
68
+ raise e
69
+
63
70
 
64
71
  async def call_tool(
65
72
  self,
66
73
  tool_name: str,
67
74
  arguments: dict[str, Any] = None,
75
+ is_fallback: bool = False,
68
76
  ) -> 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):
77
+ if is_fallback:
78
+ url = self.fallback_url
79
+ else:
80
+ url = self.url
81
+ try:
82
+ async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
71
83
  async with ClientSession(read_stream, write_stream) as session:
72
84
  await session.initialize()
73
85
  response = await session.call_tool(tool_name, arguments or {})
74
86
  content = pydantic_core.to_json(response).decode()
75
87
  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
88
+ except Exception as e:
89
+ raise e
90
+
89
91
 
90
92
  class MCPTool(BaseTool):
91
93
  """
@@ -94,7 +96,6 @@ class MCPTool(BaseTool):
94
96
  Attributes:
95
97
  client (MCPClient): The MCP client instance.
96
98
  handle_tool_error (bool | str | Callable[[ToolException], str] | None): Error handling strategy.
97
- sse (bool): Whether to use SSE streaming for responses.
98
99
  """
99
100
 
100
101
  client: MCPClient
@@ -110,7 +111,16 @@ class MCPTool(BaseTool):
110
111
 
111
112
  @t.override
112
113
  async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
113
- return await self.client.call_tool(self.name, arguments=kwargs)
114
+ try:
115
+ return await self.client.call_tool(self.name, arguments=kwargs)
116
+ except Exception as e:
117
+ if self.client.fallback_url:
118
+ try:
119
+ return await self.client.call_tool(self.name, arguments=kwargs, is_fallback=True) # Fallback to Public endpoint
120
+ except Exception as e:
121
+ raise e
122
+ else:
123
+ raise e
114
124
 
115
125
  @t.override
116
126
  @property
@@ -133,14 +143,14 @@ class MCPToolkit(BaseToolkit):
133
143
  _tools: ListToolsResult | None = None
134
144
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
135
145
 
136
- def initialize(self) -> None:
146
+ async def initialize(self) -> None:
137
147
  """Initialize the session and retrieve tools list"""
138
148
  if self._tools is None:
139
- response = self.client.list_tools()
149
+ response = await self.client.list_tools()
140
150
  self._tools = response
141
151
 
142
152
  @t.override
143
- def get_tools(self) -> list[BaseTool]:
153
+ async def get_tools(self) -> list[BaseTool]:
144
154
  if self._tools is None:
145
155
  raise RuntimeError("Must initialize the toolkit first")
146
156
 
@@ -150,7 +160,6 @@ class MCPToolkit(BaseToolkit):
150
160
  name=tool.name,
151
161
  description=tool.description or "",
152
162
  args_schema=create_schema_model(tool.name, tool.inputSchema),
153
- sse=self.sse,
154
163
  )
155
164
  # list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
156
165
  for tool in self._tools.tools
@@ -115,7 +115,7 @@ class RemoteToolkit:
115
115
  _service_name: str | None = None
116
116
  model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
117
117
 
118
- def initialize(self) -> None:
118
+ async def initialize(self) -> None:
119
119
  """Initialize the session and retrieve the remote function details."""
120
120
  if self._function is None:
121
121
  try:
@@ -136,19 +136,21 @@ class RemoteToolkit:
136
136
  f"error: {e.status_code}. Available functions: {', '.join(names)}"
137
137
  )
138
138
 
139
- def get_tools(self) -> list[BaseTool]:
139
+ async def get_tools(self) -> list[BaseTool]:
140
140
  settings = get_settings()
141
141
  if self._function is None:
142
142
  raise RuntimeError("Must initialize the toolkit first")
143
143
 
144
144
  if self._function.spec.integration_connections:
145
+ fallback_url = None
145
146
  url = f"{settings.run_url}/{settings.workspace}/functions/{self._function.metadata.name}"
146
147
  if self._service_name:
148
+ fallback_url = f"https://{self._service_name}.{settings.run_internal_hostname}"
147
149
  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()
150
+ mcp_client = MCPClient(self.client, url, fallback_url)
151
+ mcp_toolkit = MCPToolkit(client=mcp_client, url=url)
152
+ await mcp_toolkit.initialize()
153
+ return await mcp_toolkit.get_tools()
152
154
 
153
155
  if self._function.spec.kit:
154
156
  return [
beamlit/run.py CHANGED
@@ -5,6 +5,7 @@ import urllib.parse
5
5
  from typing import Any
6
6
 
7
7
  import requests
8
+
8
9
  from beamlit.client import AuthenticatedClient
9
10
  from beamlit.common import HTTPError, get_settings
10
11
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: beamlit
3
- Version: 0.0.57rc111
3
+ Version: 0.0.57rc113
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
 
@@ -2,12 +2,12 @@ beamlit/__init__.py,sha256=545gFC-wLLwUktWcOAjUWe_Glha40tBetRTOYSfHnbI,164
2
2
  beamlit/client.py,sha256=PnR6ybZk5dLIJPnDKAf2epHOeQC_7yL0fG4muvphHjA,12695
3
3
  beamlit/errors.py,sha256=gO8GBmKqmSNgAg-E5oT-oOyxztvp7V_6XG7OUTT15q0,546
4
4
  beamlit/py.typed,sha256=8ZJUsxZiuOy1oJeVhsTWQhTG_6pTVHVXk5hJL79ebTk,25
5
- beamlit/run.py,sha256=jR_DFOCqwNCQcaZg8Cf0kUuugC6InsxtmC8UoZrbpdw,4447
5
+ beamlit/run.py,sha256=YMI8iTPB1M9gd_1FG958tw-Prf9EwIcca4jK4igi7MM,4448
6
6
  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=yRxzrLOMhXW0S4ThRD7cm6UqUarw6gdB_ZP2jnZanvM,8311
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
134
+ beamlit/functions/common.py,sha256=v4nmLP9Wotpb5e6hV30XbItgZjr5lwXcF0EOotxRngI,10127
135
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
136
+ beamlit/functions/local/local.py,sha256=F7b_xYDytJIZc0fM1sYMtJgfBCF1cLjquQm83VchTLs,1891
137
+ beamlit/functions/mcp/client.py,sha256=enLo0dzWMBHJEQf6as3UWM8tN3CjUN1YO3UPn67DLac,4072
138
+ beamlit/functions/mcp/mcp.py,sha256=1KMOufh0b8SSXi7BbQd7vu9OClEXDCHUt_nuEOxu8Rc,5950
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=bkDUFiZI8YokqX8Fo76AnKLZF9PcjcDBr37hhxLevs8,6728
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.57rc113.dist-info/METADATA,sha256=Hwx6ZX-alRukrSehva8xagCSY_JA3GhJ3RcSnvY0uWY,3547
258
+ beamlit-0.0.57rc113.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
259
+ beamlit-0.0.57rc113.dist-info/entry_points.txt,sha256=zxhgdn7SP-Otk4rEv7LMPAAa9w4TUCLbu9TJi9-K3xg,115
260
+ beamlit-0.0.57rc113.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
261
+ beamlit-0.0.57rc113.dist-info/RECORD,,