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.
- beamlit/agents/decorator.py +111 -88
- beamlit/functions/common.py +16 -16
- beamlit/functions/local/local.py +4 -5
- beamlit/functions/mcp/client.py +96 -0
- beamlit/functions/mcp/mcp.py +53 -44
- beamlit/functions/remote/remote.py +8 -6
- beamlit/run.py +1 -0
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc113.dist-info}/METADATA +2 -1
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc113.dist-info}/RECORD +12 -11
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc113.dist-info}/WHEEL +0 -0
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc113.dist-info}/entry_points.txt +0 -0
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc113.dist-info}/licenses/LICENSE +0 -0
beamlit/agents/decorator.py
CHANGED
@@ -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}")
|
beamlit/functions/common.py
CHANGED
@@ -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
|
-
|
190
|
-
|
191
|
-
|
192
|
-
|
193
|
-
|
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
|
|
beamlit/functions/local/local.py
CHANGED
@@ -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()
|
beamlit/functions/mcp/mcp.py
CHANGED
@@ -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,
|
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.
|
35
|
+
self.fallback_url = fallback_url
|
36
36
|
|
37
|
-
async def
|
38
|
-
|
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
|
41
|
-
|
42
|
-
|
43
|
-
|
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
|
-
|
47
|
-
logger.
|
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
|
-
|
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:
|
59
|
-
|
60
|
-
|
61
|
-
|
62
|
-
|
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
|
70
|
-
|
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
|
-
|
77
|
-
|
78
|
-
|
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
|
-
|
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
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: beamlit
|
3
|
-
Version: 0.0.
|
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=
|
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=
|
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=
|
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=
|
137
|
-
beamlit/functions/mcp/
|
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=
|
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.
|
257
|
-
beamlit-0.0.
|
258
|
-
beamlit-0.0.
|
259
|
-
beamlit-0.0.
|
260
|
-
beamlit-0.0.
|
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,,
|
File without changes
|
File without changes
|
File without changes
|