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.
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.26.dist-info}/METADATA +2 -1
  38. {beamlit-0.0.24rc21.dist-info → beamlit-0.0.26.dist-info}/RECORD +39 -35
  39. {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(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",