beamlit 0.0.25__py3-none-any.whl → 0.0.26__py3-none-any.whl
Sign up to get free protection for your applications and to get access to all the features.
- 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.25.dist-info → beamlit-0.0.26.dist-info}/METADATA +2 -1
- {beamlit-0.0.25.dist-info → beamlit-0.0.26.dist-info}/RECORD +39 -35
- {beamlit-0.0.25.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",
|