beamlit 0.0.24rc21__py3-none-any.whl → 0.0.26rc22__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (39) hide show
  1. beamlit/agents/chat.py +26 -19
  2. beamlit/agents/decorator.py +31 -7
  3. beamlit/authentication/apikey.py +2 -2
  4. beamlit/authentication/authentication.py +2 -2
  5. beamlit/authentication/credentials.py +4 -4
  6. beamlit/common/logger.py +4 -2
  7. beamlit/common/settings.py +7 -6
  8. beamlit/deploy/deploy.py +12 -10
  9. beamlit/deploy/format.py +4 -4
  10. beamlit/functions/mcp/mcp.py +125 -0
  11. beamlit/functions/remote/remote.py +103 -0
  12. beamlit/models/__init__.py +4 -8
  13. beamlit/models/agent.py +18 -0
  14. beamlit/models/agent_history_event.py +2 -2
  15. beamlit/models/api_key.py +2 -2
  16. beamlit/models/continent.py +2 -2
  17. beamlit/models/core_status.py +62 -0
  18. beamlit/models/country.py +2 -2
  19. beamlit/models/environment_metrics.py +2 -21
  20. beamlit/models/function.py +18 -0
  21. beamlit/models/increase_and_rate_metric.py +3 -37
  22. beamlit/models/integration_config.py +45 -0
  23. beamlit/models/integration_connection_secret.py +2 -2
  24. beamlit/models/metrics.py +2 -32
  25. beamlit/models/model.py +18 -0
  26. beamlit/models/model_provider.py +2 -2
  27. beamlit/models/pending_invitation.py +2 -2
  28. beamlit/models/pending_invitation_render.py +6 -6
  29. beamlit/models/pending_invitation_render_workspace.py +2 -2
  30. beamlit/models/resource_environment_metrics.py +10 -124
  31. beamlit/models/resource_metrics.py +4 -40
  32. beamlit/models/runtime.py +1 -1
  33. beamlit/models/store_agent.py +2 -2
  34. beamlit/models/store_configuration.py +4 -4
  35. beamlit/models/store_function.py +2 -2
  36. beamlit/models/workspace.py +2 -2
  37. {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26rc22.dist-info}/METADATA +2 -1
  38. {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26rc22.dist-info}/RECORD +39 -35
  39. {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26rc22.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(agent_model: Model):
10
+ def get_base_url(name: str):
11
11
  settings = get_settings()
12
- return f"{settings.run_url}/{settings.workspace}/models/{agent_model.metadata.name}/v1"
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"] = agent_model.metadata.environment
48
+ headers["X-Beamlit-Environment"] = environment
48
49
 
49
50
  jwt = headers.get("X-Beamlit-Authorization", "").replace("Bearer ", "")
50
- params = {"environment": agent_model.metadata.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
- if agent_model is None:
86
- raise ValueError("agent_model not found in configuration")
87
- if agent_model.spec.runtime is None:
88
- raise ValueError("runtime not found in agent model")
89
- if agent_model.spec.runtime.type_ is None:
90
- raise ValueError("type not found in runtime")
91
- if agent_model.spec.runtime.model is None:
92
- raise ValueError("model not found in runtime")
93
-
94
- provider = agent_model.spec.runtime.type_
95
- model = agent_model.spec.runtime.model
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(agent_model),
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
@@ -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 get_settings, init
11
+ from beamlit.common.settings import init
12
12
  from beamlit.errors import UnexpectedStatus
13
- from beamlit.models import Agent, AgentSpec, Metadata
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 = Metadata(**agent.get("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
- runtime = settings.agent.model.spec.runtime
164
- logger.info(f"Chat model configured, using: {runtime.type_}:{runtime.model}")
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(
@@ -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.api_key,
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.api_key
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.api_key:
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.api_key:
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
- api_key: str = ""
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
- "api_key": ws.credentials.api_key,
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
- api_key=settings.authentication.api_key,
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.api_key:
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:\t %(name)s - %(message)s"))
30
+ handler.setFormatter(ColoredFormatter(f"%(levelname)s %(name)s - %(message)s"))
29
31
  logging.basicConfig(level=log_level, handlers=[handler])
@@ -38,7 +38,7 @@ class SettingsAuthenticationClient(BaseSettings):
38
38
 
39
39
 
40
40
  class SettingsAuthentication(BaseSettings):
41
- api_key: Union[None, str] = None
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.agent_chain:
121
- for chain in agent.spec.agent_chain:
121
+ if agent.spec.agentChain:
122
+ for chain in agent.spec.agentChain:
122
123
  if chain.enabled:
123
- agent_chain = get_agent.sync(chain.name, environment=env, client=client)
124
+ agentChain = get_agent.sync(chain.name, environment=env, client=client)
124
125
  if chain.description:
125
- agent_chain.spec.description = chain.description
126
- agents_chain.append(agent_chain)
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, EnvironmentMetadata, AgentSpec, AgentChain, Flavor, Function, FunctionSpec,
10
- Runtime)
11
- from .format import arg_to_dict, format_parameters, format_agent_chain
12
- from .parser import get_resources, Resource, get_parameters, get_description
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.display_name:
31
- deployment.metadata.display_name = deployment.metadata.name
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 (AgentDeployment): Agent deployment configuration
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
- display_name: {agent.metadata.display_name or agent.metadata.name}
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
- agent_chain: {format_agent_chain(agent.spec.agent_chain)}
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
- display_name: {function.metadata.display_name or function.metadata.name}
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(agent_chain: list[AgentChain]) -> str:
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
- agent_chain (list[AgentChain]): List of agent chain configurations
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 agent_chain:
60
+ if not agentChain:
61
61
  return "[]"
62
62
  formatted = []
63
63
 
64
- for agent in agent_chain:
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
+ ]
@@ -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",