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.
- beamlit/agents/decorator.py +113 -91
- beamlit/functions/common.py +16 -17
- beamlit/functions/decorator.py +1 -2
- beamlit/functions/local/local.py +5 -7
- beamlit/functions/mcp/client.py +98 -0
- beamlit/functions/mcp/mcp.py +56 -48
- beamlit/functions/remote/remote.py +9 -8
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/METADATA +2 -1
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/RECORD +12 -11
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/WHEEL +0 -0
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/entry_points.txt +0 -0
- {beamlit-0.0.57rc111.dist-info → beamlit-0.0.57rc112.dist-info}/licenses/LICENSE +0 -0
beamlit/agents/decorator.py
CHANGED
@@ -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}")
|
beamlit/functions/common.py
CHANGED
@@ -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
|
-
|
190
|
-
|
191
|
-
|
192
|
-
|
193
|
-
|
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
|
|
beamlit/functions/decorator.py
CHANGED
beamlit/functions/local/local.py
CHANGED
@@ -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()
|
beamlit/functions/mcp/mcp.py
CHANGED
@@ -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,
|
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.
|
34
|
+
self.fallback_url = fallback_url
|
36
35
|
|
37
|
-
async def
|
38
|
-
|
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
|
41
|
-
|
42
|
-
|
43
|
-
|
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
|
-
|
47
|
-
logger.
|
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
|
-
|
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:
|
59
|
-
|
60
|
-
|
61
|
-
|
62
|
-
|
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
|
70
|
-
|
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
|
-
|
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
|
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
|
-
|
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.
|
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=
|
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=
|
135
|
-
beamlit/functions/decorator.py,sha256=
|
136
|
-
beamlit/functions/local/local.py,sha256=
|
137
|
-
beamlit/functions/mcp/
|
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=
|
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.
|
257
|
-
beamlit-0.0.
|
258
|
-
beamlit-0.0.
|
259
|
-
beamlit-0.0.
|
260
|
-
beamlit-0.0.
|
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,,
|
File without changes
|
File without changes
|
File without changes
|