beamlit 0.0.24rc21__py3-none-any.whl → 0.0.26__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/chat.py +26 -19
- beamlit/agents/decorator.py +31 -7
- beamlit/authentication/apikey.py +2 -2
- beamlit/authentication/authentication.py +2 -2
- beamlit/authentication/credentials.py +4 -4
- beamlit/common/logger.py +4 -2
- beamlit/common/settings.py +7 -6
- beamlit/deploy/deploy.py +12 -10
- beamlit/deploy/format.py +4 -4
- beamlit/functions/mcp/mcp.py +125 -0
- beamlit/functions/remote/remote.py +103 -0
- beamlit/models/__init__.py +4 -8
- beamlit/models/agent.py +18 -0
- beamlit/models/agent_history_event.py +2 -2
- beamlit/models/api_key.py +2 -2
- beamlit/models/continent.py +2 -2
- beamlit/models/core_status.py +62 -0
- beamlit/models/country.py +2 -2
- beamlit/models/environment_metrics.py +2 -21
- beamlit/models/function.py +18 -0
- beamlit/models/increase_and_rate_metric.py +3 -37
- beamlit/models/integration_config.py +45 -0
- beamlit/models/integration_connection_secret.py +2 -2
- beamlit/models/metrics.py +2 -32
- beamlit/models/model.py +18 -0
- beamlit/models/model_provider.py +2 -2
- beamlit/models/pending_invitation.py +2 -2
- beamlit/models/pending_invitation_render.py +6 -6
- beamlit/models/pending_invitation_render_workspace.py +2 -2
- beamlit/models/resource_environment_metrics.py +10 -124
- beamlit/models/resource_metrics.py +4 -40
- beamlit/models/runtime.py +1 -1
- beamlit/models/store_agent.py +2 -2
- beamlit/models/store_configuration.py +4 -4
- beamlit/models/store_function.py +2 -2
- beamlit/models/workspace.py +2 -2
- {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26.dist-info}/METADATA +2 -1
- {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26.dist-info}/RECORD +39 -35
- {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26.dist-info}/WHEEL +0 -0
beamlit/agents/chat.py
CHANGED
@@ -7,9 +7,9 @@ from beamlit.models import Model
|
|
7
7
|
logger = getLogger(__name__)
|
8
8
|
|
9
9
|
|
10
|
-
def get_base_url(
|
10
|
+
def get_base_url(name: str):
|
11
11
|
settings = get_settings()
|
12
|
-
return f"{settings.run_url}/{settings.workspace}/models/{
|
12
|
+
return f"{settings.run_url}/{settings.workspace}/models/{name}/v1"
|
13
13
|
|
14
14
|
|
15
15
|
def get_mistral_chat_model(**kwargs):
|
@@ -39,15 +39,16 @@ def get_cohere_chat_model(**kwargs):
|
|
39
39
|
|
40
40
|
return ChatCohere(**kwargs)
|
41
41
|
|
42
|
-
def get_chat_model(agent_model: Model):
|
42
|
+
def get_chat_model(name: str, agent_model: Model):
|
43
43
|
settings = get_settings()
|
44
44
|
client = new_client()
|
45
45
|
|
46
|
+
environment = (agent_model.metadata and agent_model.metadata.environment) or settings.environment
|
46
47
|
headers = get_authentication_headers(settings)
|
47
|
-
headers["X-Beamlit-Environment"] =
|
48
|
+
headers["X-Beamlit-Environment"] = environment
|
48
49
|
|
49
50
|
jwt = headers.get("X-Beamlit-Authorization", "").replace("Bearer ", "")
|
50
|
-
params = {"environment":
|
51
|
+
params = {"environment": environment}
|
51
52
|
chat_classes = {
|
52
53
|
"openai": {
|
53
54
|
"func": get_openai_chat_model,
|
@@ -70,7 +71,7 @@ def get_chat_model(agent_model: Model):
|
|
70
71
|
"func": get_xai_chat_model,
|
71
72
|
"kwargs": {
|
72
73
|
"api_key": jwt,
|
73
|
-
"xai_api_base": get_base_url(),
|
74
|
+
"xai_api_base": get_base_url(name),
|
74
75
|
},
|
75
76
|
"remove_kwargs": ["base_url"],
|
76
77
|
},
|
@@ -82,21 +83,27 @@ def get_chat_model(agent_model: Model):
|
|
82
83
|
},
|
83
84
|
}
|
84
85
|
|
85
|
-
|
86
|
-
|
87
|
-
|
88
|
-
|
89
|
-
|
90
|
-
|
91
|
-
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
|
86
|
+
provider = (
|
87
|
+
agent_model.spec
|
88
|
+
and agent_model.spec.runtime
|
89
|
+
and agent_model.spec.runtime.type_
|
90
|
+
)
|
91
|
+
if not provider:
|
92
|
+
logger.warning("Provider not found in agent model, defaulting to OpenAI")
|
93
|
+
provider = "openai"
|
94
|
+
|
95
|
+
model = (
|
96
|
+
agent_model.spec
|
97
|
+
and agent_model.spec.runtime
|
98
|
+
and agent_model.spec.runtime.model
|
99
|
+
)
|
100
|
+
if not model:
|
101
|
+
logger.warning("Model not found in agent model, defaulting to gpt-4o-mini")
|
102
|
+
model = "gpt-4o-mini"
|
96
103
|
|
97
104
|
kwargs = {
|
98
105
|
"model": model,
|
99
|
-
"base_url": get_base_url(
|
106
|
+
"base_url": get_base_url(name),
|
100
107
|
"default_query": params,
|
101
108
|
"default_headers": headers,
|
102
109
|
"api_key": "fake_api_key",
|
@@ -111,4 +118,4 @@ def get_chat_model(agent_model: Model):
|
|
111
118
|
if "remove_kwargs" in chat_class:
|
112
119
|
for key in chat_class["remove_kwargs"]:
|
113
120
|
kwargs.pop(key, None)
|
114
|
-
return chat_class["func"](**kwargs)
|
121
|
+
return chat_class["func"](**kwargs), provider, model
|
beamlit/agents/decorator.py
CHANGED
@@ -8,9 +8,11 @@ from logging import getLogger
|
|
8
8
|
|
9
9
|
from beamlit.api.models import get_model
|
10
10
|
from beamlit.authentication import new_client
|
11
|
-
from beamlit.common.settings import
|
11
|
+
from beamlit.common.settings import init
|
12
12
|
from beamlit.errors import UnexpectedStatus
|
13
|
-
from beamlit.
|
13
|
+
from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
|
14
|
+
from beamlit.functions.remote.remote import RemoteToolkit
|
15
|
+
from beamlit.models import Agent, AgentMetadata, AgentSpec
|
14
16
|
from langchain_core.tools import Tool
|
15
17
|
from langgraph.checkpoint.memory import MemorySaver
|
16
18
|
from langgraph.prebuilt import create_react_agent
|
@@ -24,6 +26,7 @@ def get_functions(dir="src/functions", from_decorator="function"):
|
|
24
26
|
|
25
27
|
# Walk through all Python files in functions directory and subdirectories
|
26
28
|
if not os.path.exists(dir):
|
29
|
+
logger.warn(f"Functions directory {dir} not found")
|
27
30
|
return []
|
28
31
|
for root, _, files in os.walk(dir):
|
29
32
|
for file in files:
|
@@ -103,6 +106,8 @@ def agent(
|
|
103
106
|
agent: Agent | dict = None,
|
104
107
|
override_chat_model=None,
|
105
108
|
override_agent=None,
|
109
|
+
mcp_hub=None,
|
110
|
+
remote_functions=None,
|
106
111
|
):
|
107
112
|
logger = getLogger(__name__)
|
108
113
|
try:
|
@@ -111,6 +116,7 @@ def agent(
|
|
111
116
|
'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
|
112
117
|
)
|
113
118
|
|
119
|
+
client = new_client()
|
114
120
|
chat_model = override_chat_model or None
|
115
121
|
settings = init()
|
116
122
|
|
@@ -132,11 +138,10 @@ def agent(
|
|
132
138
|
settings.agent.functions = functions
|
133
139
|
|
134
140
|
if agent is not None:
|
135
|
-
metadata =
|
141
|
+
metadata = AgentMetadata(**agent.get("metadata", {}))
|
136
142
|
spec = AgentSpec(**agent.get("spec", {}))
|
137
143
|
agent = Agent(metadata=metadata, spec=spec)
|
138
144
|
if agent.spec.model and chat_model is None:
|
139
|
-
client = new_client()
|
140
145
|
try:
|
141
146
|
response = get_model.sync_detailed(
|
142
147
|
agent.spec.model, environment=settings.environment, client=client
|
@@ -158,10 +163,29 @@ def agent(
|
|
158
163
|
raise e
|
159
164
|
|
160
165
|
if settings.agent.model:
|
161
|
-
chat_model = get_chat_model(settings.agent.model)
|
166
|
+
chat_model, provider, model = get_chat_model(agent.spec.model, settings.agent.model)
|
162
167
|
settings.agent.chat_model = chat_model
|
163
|
-
|
164
|
-
|
168
|
+
logger.info(f"Chat model configured, using: {provider}:{model}")
|
169
|
+
|
170
|
+
if mcp_hub:
|
171
|
+
for server in mcp_hub:
|
172
|
+
try:
|
173
|
+
mcp_client = MCPClient(client, server)
|
174
|
+
toolkit = MCPToolkit(client=mcp_client)
|
175
|
+
toolkit.initialize()
|
176
|
+
functions.extend(toolkit.get_tools())
|
177
|
+
except Exception as e:
|
178
|
+
logger.warn(f"Failed to initialize MCP server {server}: {e!s}")
|
179
|
+
|
180
|
+
if remote_functions:
|
181
|
+
|
182
|
+
for function in remote_functions:
|
183
|
+
try:
|
184
|
+
toolkit = RemoteToolkit(client, function)
|
185
|
+
toolkit.initialize()
|
186
|
+
functions.extend(toolkit.get_tools())
|
187
|
+
except Exception as e:
|
188
|
+
logger.warn(f"Failed to initialize remote function {function}: {e!s}")
|
165
189
|
|
166
190
|
if override_agent is None and len(functions) == 0:
|
167
191
|
raise ValueError(
|
beamlit/authentication/apikey.py
CHANGED
@@ -10,11 +10,11 @@ class ApiKeyProvider(Auth):
|
|
10
10
|
|
11
11
|
def get_headers(self):
|
12
12
|
return {
|
13
|
-
"X-Beamlit-Api-Key": self.credentials.
|
13
|
+
"X-Beamlit-Api-Key": self.credentials.apiKey,
|
14
14
|
"X-Beamlit-Workspace": self.workspace_name,
|
15
15
|
}
|
16
16
|
|
17
17
|
def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
|
18
|
-
request.headers["X-Beamlit-Api-Key"] = self.credentials.
|
18
|
+
request.headers["X-Beamlit-Api-Key"] = self.credentials.apiKey
|
19
19
|
request.headers["X-Beamlit-Workspace"] = self.workspace_name
|
20
20
|
yield request
|
@@ -60,7 +60,7 @@ def new_client():
|
|
60
60
|
|
61
61
|
def new_client_with_credentials(config: RunClientWithCredentials):
|
62
62
|
provider: Auth = None
|
63
|
-
if config.credentials.
|
63
|
+
if config.credentials.apiKey:
|
64
64
|
provider = ApiKeyProvider(config.credentials, config.workspace)
|
65
65
|
elif config.credentials.access_token:
|
66
66
|
provider = BearerToken(config.credentials, config.workspace, config.api_url)
|
@@ -85,7 +85,7 @@ def get_authentication_headers(settings: Settings) -> Dict[str, str]:
|
|
85
85
|
workspace=settings.workspace,
|
86
86
|
)
|
87
87
|
provider = None
|
88
|
-
if config.credentials.
|
88
|
+
if config.credentials.apiKey:
|
89
89
|
provider = ApiKeyProvider(config.credentials, config.workspace)
|
90
90
|
elif config.credentials.access_token:
|
91
91
|
provider = BearerToken(config.credentials, config.workspace, config.api_url)
|
@@ -12,7 +12,7 @@ logger = getLogger(__name__)
|
|
12
12
|
|
13
13
|
@dataclass
|
14
14
|
class Credentials:
|
15
|
-
|
15
|
+
apiKey: str = ""
|
16
16
|
access_token: str = ""
|
17
17
|
refresh_token: str = ""
|
18
18
|
expires_in: int = 0
|
@@ -49,7 +49,7 @@ class Config:
|
|
49
49
|
{
|
50
50
|
"name": ws.name,
|
51
51
|
"credentials": {
|
52
|
-
"
|
52
|
+
"apiKey": ws.credentials.apiKey,
|
53
53
|
"access_token": ws.credentials.access_token,
|
54
54
|
"refresh_token": ws.credentials.refresh_token,
|
55
55
|
"expires_in": ws.credentials.expires_in,
|
@@ -129,7 +129,7 @@ def load_credentials(workspace_name: str) -> Credentials:
|
|
129
129
|
|
130
130
|
def load_credentials_from_settings(settings: Settings) -> Credentials:
|
131
131
|
return Credentials(
|
132
|
-
|
132
|
+
apiKey=settings.authentication.apiKey,
|
133
133
|
client_credentials=settings.authentication.client.credentials,
|
134
134
|
)
|
135
135
|
|
@@ -154,7 +154,7 @@ def create_home_dir_if_missing():
|
|
154
154
|
|
155
155
|
def save_credentials(workspace_name: str, credentials: Credentials):
|
156
156
|
create_home_dir_if_missing()
|
157
|
-
if not credentials.access_token and not credentials.
|
157
|
+
if not credentials.access_token and not credentials.apiKey:
|
158
158
|
logger.info("No credentials to save, error")
|
159
159
|
return
|
160
160
|
|
beamlit/common/logger.py
CHANGED
@@ -11,8 +11,10 @@ class ColoredFormatter(logging.Formatter):
|
|
11
11
|
}
|
12
12
|
|
13
13
|
def format(self, record):
|
14
|
+
n_spaces = len("CRITICAL") - len(record.levelname)
|
15
|
+
tab = " " * n_spaces
|
14
16
|
color = self.COLORS.get(record.levelname, "\033[0m")
|
15
|
-
record.levelname = f"{color}{record.levelname}\033[0m"
|
17
|
+
record.levelname = f"{color}{record.levelname}\033[0m:{tab}"
|
16
18
|
return super().format(record)
|
17
19
|
|
18
20
|
|
@@ -25,5 +27,5 @@ def init(log_level: str):
|
|
25
27
|
logging.getLogger("httpx").propagate = False
|
26
28
|
|
27
29
|
handler = logging.StreamHandler()
|
28
|
-
handler.setFormatter(ColoredFormatter("%(levelname)s
|
30
|
+
handler.setFormatter(ColoredFormatter(f"%(levelname)s %(name)s - %(message)s"))
|
29
31
|
logging.basicConfig(level=log_level, handlers=[handler])
|
beamlit/common/settings.py
CHANGED
@@ -38,7 +38,7 @@ class SettingsAuthenticationClient(BaseSettings):
|
|
38
38
|
|
39
39
|
|
40
40
|
class SettingsAuthentication(BaseSettings):
|
41
|
-
|
41
|
+
apiKey: Union[None, str] = None
|
42
42
|
jwt: Union[None, str] = None
|
43
43
|
client: SettingsAuthenticationClient = SettingsAuthenticationClient()
|
44
44
|
|
@@ -64,6 +64,7 @@ class Settings(BaseSettings):
|
|
64
64
|
name: str = Field(default="beamlit-agent")
|
65
65
|
base_url: str = Field(default="https://api.beamlit.dev/v0")
|
66
66
|
run_url: str = Field(default="https://run.beamlit.dev")
|
67
|
+
mcp_hub_url: str = Field(default="https://mcp-hub-server.beamlit.workers.dev")
|
67
68
|
registry_url: str = Field(default="https://serverless-registry-production.beamlit.workers.dev")
|
68
69
|
log_level: str = Field(default="INFO")
|
69
70
|
agent: SettingsAgent = SettingsAgent()
|
@@ -117,13 +118,13 @@ def init_agent(
|
|
117
118
|
functions.append(function)
|
118
119
|
settings.agent.functions = functions
|
119
120
|
|
120
|
-
if agent.spec.
|
121
|
-
for chain in agent.spec.
|
121
|
+
if agent.spec.agentChain:
|
122
|
+
for chain in agent.spec.agentChain:
|
122
123
|
if chain.enabled:
|
123
|
-
|
124
|
+
agentChain = get_agent.sync(chain.name, environment=env, client=client)
|
124
125
|
if chain.description:
|
125
|
-
|
126
|
-
agents_chain.append(
|
126
|
+
agentChain.spec.description = chain.description
|
127
|
+
agents_chain.append(agentChain)
|
127
128
|
settings.agent.chain = agents_chain
|
128
129
|
if agent.spec.model:
|
129
130
|
model = get_model.sync(agent.spec.model, environment=env, client=client)
|
beamlit/deploy/deploy.py
CHANGED
@@ -6,10 +6,12 @@ from logging import getLogger
|
|
6
6
|
from typing import Literal
|
7
7
|
|
8
8
|
from beamlit.common.settings import Settings, get_settings, init
|
9
|
-
from beamlit.models import (Agent,
|
10
|
-
Runtime)
|
11
|
-
|
12
|
-
from .
|
9
|
+
from beamlit.models import (Agent, AgentChain, AgentSpec, EnvironmentMetadata,
|
10
|
+
Flavor, Function, FunctionSpec, Runtime)
|
11
|
+
|
12
|
+
from .format import arg_to_dict, format_agent_chain, format_parameters
|
13
|
+
from .parser import Resource, get_description, get_parameters, get_resources
|
14
|
+
|
13
15
|
sys.path.insert(0, os.getcwd())
|
14
16
|
sys.path.insert(0, os.path.join(os.getcwd(), "src"))
|
15
17
|
|
@@ -27,8 +29,8 @@ def set_default_values(resource: Resource, deployment: Agent | Function):
|
|
27
29
|
deployment.metadata.environment = settings.environment
|
28
30
|
if not deployment.metadata.name:
|
29
31
|
deployment.metadata.name = resource.name
|
30
|
-
if not deployment.metadata.
|
31
|
-
deployment.metadata.
|
32
|
+
if not deployment.metadata.displayName:
|
33
|
+
deployment.metadata.displayName = deployment.metadata.name
|
32
34
|
if not deployment.spec.description:
|
33
35
|
deployment.spec.description = get_description(None, resource)
|
34
36
|
if not deployment.spec.runtime:
|
@@ -97,7 +99,7 @@ def get_agent_yaml(
|
|
97
99
|
Generates YAML configuration for an agent deployment.
|
98
100
|
|
99
101
|
Args:
|
100
|
-
agent (
|
102
|
+
agent (Agent): Agent deployment configuration
|
101
103
|
functions (list[tuple[Resource, FunctionDeployment]]): List of associated functions
|
102
104
|
settings (Settings): Application settings
|
103
105
|
|
@@ -109,14 +111,14 @@ apiVersion: beamlit.com/v1alpha1
|
|
109
111
|
kind: Agent
|
110
112
|
metadata:
|
111
113
|
name: {agent.metadata.name}
|
112
|
-
|
114
|
+
displayName: {agent.metadata.displayName or agent.metadata.name}
|
113
115
|
environment: {settings.environment}
|
114
116
|
workspace: {settings.workspace}
|
115
117
|
spec:
|
116
118
|
enabled: true
|
117
119
|
policies: [{", ".join(agent.spec.policies or [])}]
|
118
120
|
functions: [{", ".join([f"{function.metadata.name}" for (_, function) in functions])}]
|
119
|
-
|
121
|
+
agentChain: {format_agent_chain(agent.spec.agentChain)}
|
120
122
|
model: {agent.spec.model}
|
121
123
|
runtime:
|
122
124
|
image: {agent.spec.runtime.image}
|
@@ -143,7 +145,7 @@ apiVersion: beamlit.com/v1alpha1
|
|
143
145
|
kind: Function
|
144
146
|
metadata:
|
145
147
|
name: {function.metadata.name}
|
146
|
-
|
148
|
+
displayName: {function.metadata.displayName or function.metadata.name}
|
147
149
|
environment: {settings.environment}
|
148
150
|
spec:
|
149
151
|
enabled: true
|
beamlit/deploy/format.py
CHANGED
@@ -47,21 +47,21 @@ def format_parameters(parameters: list[StoreFunctionParameter]) -> str:
|
|
47
47
|
return "\n".join(formatted)
|
48
48
|
|
49
49
|
|
50
|
-
def format_agent_chain(
|
50
|
+
def format_agent_chain(agentChain: list[AgentChain]) -> str:
|
51
51
|
"""
|
52
52
|
Formats agent chain configuration into YAML-compatible string.
|
53
53
|
|
54
54
|
Args:
|
55
|
-
|
55
|
+
agentChain (list[AgentChain]): List of agent chain configurations
|
56
56
|
|
57
57
|
Returns:
|
58
58
|
str: YAML-formatted string of agent chain
|
59
59
|
"""
|
60
|
-
if not
|
60
|
+
if not agentChain:
|
61
61
|
return "[]"
|
62
62
|
formatted = []
|
63
63
|
|
64
|
-
for agent in
|
64
|
+
for agent in agentChain:
|
65
65
|
formatted.append(f"""
|
66
66
|
- agent: {agent.name}
|
67
67
|
enabled: {agent.enabled}""")
|
@@ -0,0 +1,125 @@
|
|
1
|
+
import asyncio
|
2
|
+
import urllib.parse
|
3
|
+
import warnings
|
4
|
+
from typing import Any, Callable
|
5
|
+
|
6
|
+
import pydantic
|
7
|
+
import pydantic_core
|
8
|
+
import requests
|
9
|
+
import typing_extensions as t
|
10
|
+
from beamlit.authentication.authentication import AuthenticatedClient
|
11
|
+
from beamlit.common.settings import get_settings
|
12
|
+
from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
|
13
|
+
from mcp import ListToolsResult
|
14
|
+
from pydantic.json_schema import JsonSchemaValue
|
15
|
+
from pydantic_core import core_schema as cs
|
16
|
+
|
17
|
+
settings = get_settings()
|
18
|
+
|
19
|
+
|
20
|
+
def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
|
21
|
+
# Create a new model class that returns our JSON schema.
|
22
|
+
# LangChain requires a BaseModel class.
|
23
|
+
class Schema(pydantic.BaseModel):
|
24
|
+
model_config = pydantic.ConfigDict(extra="allow")
|
25
|
+
|
26
|
+
@t.override
|
27
|
+
@classmethod
|
28
|
+
def __get_pydantic_json_schema__(
|
29
|
+
cls, core_schema: cs.CoreSchema, handler: pydantic.GetJsonSchemaHandler
|
30
|
+
) -> JsonSchemaValue:
|
31
|
+
return schema
|
32
|
+
|
33
|
+
return Schema
|
34
|
+
|
35
|
+
|
36
|
+
class MCPClient:
|
37
|
+
def __init__(self, client: AuthenticatedClient, server_name: str):
|
38
|
+
self.client = client
|
39
|
+
self.server_name = server_name
|
40
|
+
self.headers = {"Api-Key": "1234567890"}
|
41
|
+
|
42
|
+
def list_tools(self) -> requests.Response:
|
43
|
+
client = self.client.get_httpx_client()
|
44
|
+
url = urllib.parse.urljoin(settings.mcp_hub_url, f"{self.server_name}/tools/list")
|
45
|
+
|
46
|
+
response = client.request("GET", url, headers=self.headers)
|
47
|
+
response.raise_for_status()
|
48
|
+
return response
|
49
|
+
|
50
|
+
def call_tool(self, tool_name: str, arguments: dict[str, Any] = None) -> requests.Response:
|
51
|
+
client = self.client.get_httpx_client()
|
52
|
+
url = urllib.parse.urljoin(settings.mcp_hub_url, f"{self.server_name}/tools/call")
|
53
|
+
response = client.request(
|
54
|
+
"POST",
|
55
|
+
url,
|
56
|
+
json={"name": tool_name, "arguments": arguments},
|
57
|
+
headers=self.headers,
|
58
|
+
)
|
59
|
+
response.raise_for_status()
|
60
|
+
return response
|
61
|
+
|
62
|
+
class MCPTool(BaseTool):
|
63
|
+
"""
|
64
|
+
MCP server tool
|
65
|
+
"""
|
66
|
+
|
67
|
+
client: MCPClient
|
68
|
+
handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
|
69
|
+
|
70
|
+
@t.override
|
71
|
+
def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
72
|
+
warnings.warn(
|
73
|
+
"Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
|
74
|
+
stacklevel=1,
|
75
|
+
)
|
76
|
+
return asyncio.run(self._arun(*args, **kwargs))
|
77
|
+
|
78
|
+
@t.override
|
79
|
+
async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
80
|
+
result = self.client.call_tool(self.name, arguments=kwargs)
|
81
|
+
response = result.json()
|
82
|
+
content = pydantic_core.to_json(response["content"]).decode()
|
83
|
+
if response["isError"]:
|
84
|
+
raise ToolException(content)
|
85
|
+
return content
|
86
|
+
|
87
|
+
@t.override
|
88
|
+
@property
|
89
|
+
def tool_call_schema(self) -> type[pydantic.BaseModel]:
|
90
|
+
assert self.args_schema is not None # noqa: S101
|
91
|
+
return self.args_schema
|
92
|
+
|
93
|
+
class MCPToolkit(BaseToolkit):
|
94
|
+
"""
|
95
|
+
MCP server toolkit
|
96
|
+
"""
|
97
|
+
|
98
|
+
client: MCPClient
|
99
|
+
"""The MCP session used to obtain the tools"""
|
100
|
+
|
101
|
+
_tools: ListToolsResult | None = None
|
102
|
+
|
103
|
+
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
|
104
|
+
|
105
|
+
def initialize(self) -> None:
|
106
|
+
"""Initialize the session and retrieve tools list"""
|
107
|
+
if self._tools is None:
|
108
|
+
response = self.client.list_tools()
|
109
|
+
self._tools = ListToolsResult(**response.json())
|
110
|
+
|
111
|
+
@t.override
|
112
|
+
def get_tools(self) -> list[BaseTool]:
|
113
|
+
if self._tools is None:
|
114
|
+
raise RuntimeError("Must initialize the toolkit first")
|
115
|
+
|
116
|
+
return [
|
117
|
+
MCPTool(
|
118
|
+
client=self.client,
|
119
|
+
name=tool.name,
|
120
|
+
description=tool.description or "",
|
121
|
+
args_schema=create_schema_model(tool.inputSchema),
|
122
|
+
)
|
123
|
+
# list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
|
124
|
+
for tool in self._tools.tools
|
125
|
+
]
|
@@ -0,0 +1,103 @@
|
|
1
|
+
import asyncio
|
2
|
+
import urllib.parse
|
3
|
+
import warnings
|
4
|
+
from typing import Any, Callable
|
5
|
+
|
6
|
+
import pydantic
|
7
|
+
import pydantic_core
|
8
|
+
import typing_extensions as t
|
9
|
+
from beamlit.api.functions import get_function
|
10
|
+
from beamlit.authentication.authentication import AuthenticatedClient
|
11
|
+
from beamlit.models.function import Function
|
12
|
+
from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
|
13
|
+
from pydantic.json_schema import JsonSchemaValue
|
14
|
+
from pydantic_core import core_schema as cs
|
15
|
+
|
16
|
+
|
17
|
+
def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
|
18
|
+
# Create a new model class that returns our JSON schema.
|
19
|
+
# LangChain requires a BaseModel class.
|
20
|
+
class Schema(pydantic.BaseModel):
|
21
|
+
model_config = pydantic.ConfigDict(extra="allow")
|
22
|
+
|
23
|
+
@t.override
|
24
|
+
@classmethod
|
25
|
+
def __get_pydantic_json_schema__(
|
26
|
+
cls, core_schema: cs.CoreSchema, handler: pydantic.GetJsonSchemaHandler
|
27
|
+
) -> JsonSchemaValue:
|
28
|
+
return schema
|
29
|
+
|
30
|
+
return Schema
|
31
|
+
|
32
|
+
|
33
|
+
class RemoteTool(BaseTool):
|
34
|
+
"""
|
35
|
+
Remote tool
|
36
|
+
"""
|
37
|
+
|
38
|
+
client: AuthenticatedClient
|
39
|
+
handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
|
40
|
+
|
41
|
+
@t.override
|
42
|
+
def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
43
|
+
warnings.warn(
|
44
|
+
"Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
|
45
|
+
stacklevel=1,
|
46
|
+
)
|
47
|
+
return asyncio.run(self._arun(*args, **kwargs))
|
48
|
+
|
49
|
+
@t.override
|
50
|
+
async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
51
|
+
result = self.client.call_tool(self.name, arguments=kwargs)
|
52
|
+
response = result.json()
|
53
|
+
content = pydantic_core.to_json(response["content"]).decode()
|
54
|
+
if response["isError"]:
|
55
|
+
raise ToolException(content)
|
56
|
+
return content
|
57
|
+
|
58
|
+
@t.override
|
59
|
+
@property
|
60
|
+
def tool_call_schema(self) -> type[pydantic.BaseModel]:
|
61
|
+
assert self.args_schema is not None # noqa: S101
|
62
|
+
return self.args_schema
|
63
|
+
|
64
|
+
class RemoteToolkit(BaseToolkit):
|
65
|
+
"""
|
66
|
+
Remote toolkit
|
67
|
+
"""
|
68
|
+
|
69
|
+
client: AuthenticatedClient
|
70
|
+
function: str
|
71
|
+
_function: Function | None = None
|
72
|
+
|
73
|
+
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
|
74
|
+
|
75
|
+
def initialize(self) -> None:
|
76
|
+
"""Initialize the session and retrieve tools list"""
|
77
|
+
if self._function is None:
|
78
|
+
self._function = get_function(self.function, client=self.client)
|
79
|
+
|
80
|
+
@t.override
|
81
|
+
def get_tools(self) -> list[BaseTool]:
|
82
|
+
if self._tools is None:
|
83
|
+
raise RuntimeError("Must initialize the toolkit first")
|
84
|
+
|
85
|
+
if self._function.spec.kit:
|
86
|
+
return [
|
87
|
+
RemoteTool(
|
88
|
+
client=self.client,
|
89
|
+
name=func.name,
|
90
|
+
description=func.description or "",
|
91
|
+
args_schema=create_schema_model(func.parameters),
|
92
|
+
)
|
93
|
+
for func in self._function.spec.kit
|
94
|
+
]
|
95
|
+
|
96
|
+
return [
|
97
|
+
RemoteTool(
|
98
|
+
client=self.client,
|
99
|
+
name=self._function.metadata.name,
|
100
|
+
description=self._function.spec.description or "",
|
101
|
+
args_schema=create_schema_model(self._function.spec.parameters),
|
102
|
+
)
|
103
|
+
]
|
beamlit/models/__init__.py
CHANGED
@@ -13,6 +13,7 @@ from .configuration import Configuration
|
|
13
13
|
from .continent import Continent
|
14
14
|
from .core_spec import CoreSpec
|
15
15
|
from .core_spec_configurations import CoreSpecConfigurations
|
16
|
+
from .core_status import CoreStatus
|
16
17
|
from .country import Country
|
17
18
|
from .create_api_key_for_service_account_body import CreateApiKeyForServiceAccountBody
|
18
19
|
from .create_workspace_service_account_body import CreateWorkspaceServiceAccountBody
|
@@ -30,6 +31,7 @@ from .function_release import FunctionRelease
|
|
30
31
|
from .function_spec import FunctionSpec
|
31
32
|
from .get_workspace_service_accounts_response_200_item import GetWorkspaceServiceAccountsResponse200Item
|
32
33
|
from .increase_and_rate_metric import IncreaseAndRateMetric
|
34
|
+
from .integration_config import IntegrationConfig
|
33
35
|
from .integration_connection import IntegrationConnection
|
34
36
|
from .integration_connection_config import IntegrationConnectionConfig
|
35
37
|
from .integration_connection_secret import IntegrationConnectionSecret
|
@@ -71,13 +73,7 @@ from .resource_deployment_metrics_query_per_second_per_region_per_code import (
|
|
71
73
|
)
|
72
74
|
from .resource_environment_metrics import ResourceEnvironmentMetrics
|
73
75
|
from .resource_environment_metrics_inference_per_region import ResourceEnvironmentMetricsInferencePerRegion
|
74
|
-
from .resource_environment_metrics_inference_per_second_per_region import (
|
75
|
-
ResourceEnvironmentMetricsInferencePerSecondPerRegion,
|
76
|
-
)
|
77
76
|
from .resource_environment_metrics_query_per_region_per_code import ResourceEnvironmentMetricsQueryPerRegionPerCode
|
78
|
-
from .resource_environment_metrics_query_per_second_per_region_per_code import (
|
79
|
-
ResourceEnvironmentMetricsQueryPerSecondPerRegionPerCode,
|
80
|
-
)
|
81
77
|
from .resource_log import ResourceLog
|
82
78
|
from .resource_metrics import ResourceMetrics
|
83
79
|
from .runtime import Runtime
|
@@ -116,6 +112,7 @@ __all__ = (
|
|
116
112
|
"Continent",
|
117
113
|
"CoreSpec",
|
118
114
|
"CoreSpecConfigurations",
|
115
|
+
"CoreStatus",
|
119
116
|
"Country",
|
120
117
|
"CreateApiKeyForServiceAccountBody",
|
121
118
|
"CreateWorkspaceServiceAccountBody",
|
@@ -133,6 +130,7 @@ __all__ = (
|
|
133
130
|
"FunctionSpec",
|
134
131
|
"GetWorkspaceServiceAccountsResponse200Item",
|
135
132
|
"IncreaseAndRateMetric",
|
133
|
+
"IntegrationConfig",
|
136
134
|
"IntegrationConnection",
|
137
135
|
"IntegrationConnectionConfig",
|
138
136
|
"IntegrationConnectionSecret",
|
@@ -170,9 +168,7 @@ __all__ = (
|
|
170
168
|
"ResourceDeploymentMetricsQueryPerSecondPerRegionPerCode",
|
171
169
|
"ResourceEnvironmentMetrics",
|
172
170
|
"ResourceEnvironmentMetricsInferencePerRegion",
|
173
|
-
"ResourceEnvironmentMetricsInferencePerSecondPerRegion",
|
174
171
|
"ResourceEnvironmentMetricsQueryPerRegionPerCode",
|
175
|
-
"ResourceEnvironmentMetricsQueryPerSecondPerRegionPerCode",
|
176
172
|
"ResourceLog",
|
177
173
|
"ResourceMetrics",
|
178
174
|
"Runtime",
|