beamlit 0.0.24rc21__tar.gz → 0.0.26rc22__tar.gz

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: beamlit
-Version: 0.0.24rc21
+Version: 0.0.26rc22
 Summary: Add your description here
 Author-email: cploujoux <ch.ploujoux@gmail.com>
 Requires-Python: >=3.12
@@ -12,6 +12,7 @@ Requires-Dist: langchain-community<0.4.0,>=0.3.3
 Requires-Dist: langchain-core<0.4.0,>=0.3.13
 Requires-Dist: langchain-openai<0.3.0,>=0.2.4
 Requires-Dist: langgraph<0.3.0,>=0.2.40
+Requires-Dist: mcp>=1.1.2
 Requires-Dist: pydantic-settings<2.7.0,>=2.6.1
 Requires-Dist: pydantic<2.11.0,>=2.10.3
 Requires-Dist: python-dateutil>=2.8.0

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "beamlit"
-version = "0.0.24.pre-21"
+version = "0.0.26.pre-22"
 description = "Add your description here"
 readme = "README.md"
 authors = [
@@ -25,6 +25,7 @@ dependencies = [
     "langgraph>=0.2.40,<0.3.0",
     "fastapi[standard]>=0.115.4,<0.116.0",
     "asgi-correlation-id>=4.3.4,<5.0.0",
+    "mcp>=1.1.2",
 ]
 [dependency-groups]
 dev = [
@@ -40,4 +41,4 @@ line-length = 100
 
 [tool.ruff.lint]
 select = ["F", "I", "UP"]
-ignore = ["UP006", "UP007", "UP035"]
+ignore = ["UP006", "UP007", "UP035"]

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/agents/chat.py

@@ -7,9 +7,9 @@ from beamlit.models import Model
 logger = getLogger(__name__)
 
 
-def get_base_url(agent_model: Model):
+def get_base_url(name: str):
     settings = get_settings()
-    return f"{settings.run_url}/{settings.workspace}/models/{agent_model.metadata.name}/v1"
+    return f"{settings.run_url}/{settings.workspace}/models/{name}/v1"
 
 
 def get_mistral_chat_model(**kwargs):
@@ -39,15 +39,16 @@ def get_cohere_chat_model(**kwargs):
 
     return ChatCohere(**kwargs)
 
-def get_chat_model(agent_model: Model):
+def get_chat_model(name: str, agent_model: Model):
     settings = get_settings()
     client = new_client()
 
+    environment = (agent_model.metadata and agent_model.metadata.environment) or settings.environment
     headers = get_authentication_headers(settings)
-    headers["X-Beamlit-Environment"] = agent_model.metadata.environment
+    headers["X-Beamlit-Environment"] = environment
 
     jwt = headers.get("X-Beamlit-Authorization", "").replace("Bearer ", "")
-    params = {"environment": agent_model.metadata.environment}
+    params = {"environment": environment}
     chat_classes = {
         "openai": {
             "func": get_openai_chat_model,
@@ -70,7 +71,7 @@ def get_chat_model(agent_model: Model):
             "func": get_xai_chat_model,
             "kwargs": {
                 "api_key": jwt,
-                "xai_api_base": get_base_url(),
+                "xai_api_base": get_base_url(name),
             },
             "remove_kwargs": ["base_url"],
         },
@@ -82,21 +83,27 @@ def get_chat_model(agent_model: Model):
         },
     }
 
-    if agent_model is None:
-        raise ValueError("agent_model not found in configuration")
-    if agent_model.spec.runtime is None:
-        raise ValueError("runtime not found in agent model")
-    if agent_model.spec.runtime.type_ is None:
-        raise ValueError("type not found in runtime")
-    if agent_model.spec.runtime.model is None:
-        raise ValueError("model not found in runtime")
-
-    provider = agent_model.spec.runtime.type_
-    model = agent_model.spec.runtime.model
+    provider = (
+        agent_model.spec
+        and agent_model.spec.runtime
+        and agent_model.spec.runtime.type_
+    )
+    if not provider:
+        logger.warning("Provider not found in agent model, defaulting to OpenAI")
+        provider = "openai"
+
+    model = (
+        agent_model.spec
+        and agent_model.spec.runtime
+        and agent_model.spec.runtime.model
+    )
+    if not model:
+        logger.warning("Model not found in agent model, defaulting to gpt-4o-mini")
+        model = "gpt-4o-mini"
 
     kwargs = {
         "model": model,
-        "base_url": get_base_url(agent_model),
+        "base_url": get_base_url(name),
         "default_query": params,
         "default_headers": headers,
         "api_key": "fake_api_key",
@@ -111,4 +118,4 @@ def get_chat_model(agent_model: Model):
     if "remove_kwargs" in chat_class:
         for key in chat_class["remove_kwargs"]:
             kwargs.pop(key, None)
-    return chat_class["func"](**kwargs)
+    return chat_class["func"](**kwargs), provider, model

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/agents/decorator.py

@@ -8,9 +8,11 @@ from logging import getLogger
 
 from beamlit.api.models import get_model
 from beamlit.authentication import new_client
-from beamlit.common.settings import get_settings, init
+from beamlit.common.settings import init
 from beamlit.errors import UnexpectedStatus
-from beamlit.models import Agent, AgentSpec, Metadata
+from beamlit.functions.mcp.mcp import MCPClient, MCPToolkit
+from beamlit.functions.remote.remote import RemoteToolkit
+from beamlit.models import Agent, AgentMetadata, AgentSpec
 from langchain_core.tools import Tool
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.prebuilt import create_react_agent
@@ -24,6 +26,7 @@ def get_functions(dir="src/functions", from_decorator="function"):
 
     # Walk through all Python files in functions directory and subdirectories
     if not os.path.exists(dir):
+        logger.warn(f"Functions directory {dir} not found")
         return []
     for root, _, files in os.walk(dir):
         for file in files:
@@ -103,6 +106,8 @@ def agent(
     agent: Agent | dict = None,
     override_chat_model=None,
     override_agent=None,
+    mcp_hub=None,
+    remote_functions=None,
 ):
     logger = getLogger(__name__)
     try:
@@ -111,6 +116,7 @@ def agent(
                 'agent must be a dictionary, example: @agent(agent={"metadata": {"name": "my_agent"}})'
             )
 
+        client = new_client()
         chat_model = override_chat_model or None
         settings = init()
 
@@ -132,11 +138,10 @@ def agent(
         settings.agent.functions = functions
 
         if agent is not None:
-            metadata = Metadata(**agent.get("metadata", {}))
+            metadata = AgentMetadata(**agent.get("metadata", {}))
             spec = AgentSpec(**agent.get("spec", {}))
             agent = Agent(metadata=metadata, spec=spec)
             if agent.spec.model and chat_model is None:
-                client = new_client()
                 try:
                     response = get_model.sync_detailed(
                         agent.spec.model, environment=settings.environment, client=client
@@ -158,10 +163,29 @@ def agent(
                     raise e
 
                 if settings.agent.model:
-                    chat_model = get_chat_model(settings.agent.model)
+                    chat_model, provider, model = get_chat_model(agent.spec.model, settings.agent.model)
                     settings.agent.chat_model = chat_model
-                    runtime = settings.agent.model.spec.runtime
-                    logger.info(f"Chat model configured, using: {runtime.type_}:{runtime.model}")
+                    logger.info(f"Chat model configured, using: {provider}:{model}")
+
+        if mcp_hub:
+            for server in mcp_hub:
+                try:
+                    mcp_client = MCPClient(client, server)
+                    toolkit = MCPToolkit(client=mcp_client)
+                    toolkit.initialize()
+                    functions.extend(toolkit.get_tools())
+                except Exception as e:
+                    logger.warn(f"Failed to initialize MCP server {server}: {e!s}")
+
+        if remote_functions:
+
+            for function in remote_functions:
+                try:
+                    toolkit = RemoteToolkit(client, function)
+                    toolkit.initialize()
+                    functions.extend(toolkit.get_tools())
+                except Exception as e:
+                    logger.warn(f"Failed to initialize remote function {function}: {e!s}")
 
         if override_agent is None and len(functions) == 0:
             raise ValueError(

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/authentication/apikey.py

@@ -10,11 +10,11 @@ class ApiKeyProvider(Auth):
 
     def get_headers(self):
         return {
-            "X-Beamlit-Api-Key": self.credentials.api_key,
+            "X-Beamlit-Api-Key": self.credentials.apiKey,
             "X-Beamlit-Workspace": self.workspace_name,
         }
 
     def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
-        request.headers["X-Beamlit-Api-Key"] = self.credentials.api_key
+        request.headers["X-Beamlit-Api-Key"] = self.credentials.apiKey
         request.headers["X-Beamlit-Workspace"] = self.workspace_name
         yield request

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/authentication/authentication.py

@@ -60,7 +60,7 @@ def new_client():
 
 def new_client_with_credentials(config: RunClientWithCredentials):
     provider: Auth = None
-    if config.credentials.api_key:
+    if config.credentials.apiKey:
         provider = ApiKeyProvider(config.credentials, config.workspace)
     elif config.credentials.access_token:
         provider = BearerToken(config.credentials, config.workspace, config.api_url)
@@ -85,7 +85,7 @@ def get_authentication_headers(settings: Settings) -> Dict[str, str]:
         workspace=settings.workspace,
     )
     provider = None
-    if config.credentials.api_key:
+    if config.credentials.apiKey:
         provider = ApiKeyProvider(config.credentials, config.workspace)
     elif config.credentials.access_token:
         provider = BearerToken(config.credentials, config.workspace, config.api_url)

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/authentication/credentials.py

@@ -12,7 +12,7 @@ logger = getLogger(__name__)
 
 @dataclass
 class Credentials:
-    api_key: str = ""
+    apiKey: str = ""
     access_token: str = ""
     refresh_token: str = ""
     expires_in: int = 0
@@ -49,7 +49,7 @@ class Config:
                 {
                     "name": ws.name,
                     "credentials": {
-                        "api_key": ws.credentials.api_key,
+                        "apiKey": ws.credentials.apiKey,
                         "access_token": ws.credentials.access_token,
                         "refresh_token": ws.credentials.refresh_token,
                         "expires_in": ws.credentials.expires_in,
@@ -129,7 +129,7 @@ def load_credentials(workspace_name: str) -> Credentials:
 
 def load_credentials_from_settings(settings: Settings) -> Credentials:
     return Credentials(
-        api_key=settings.authentication.api_key,
+        apiKey=settings.authentication.apiKey,
         client_credentials=settings.authentication.client.credentials,
     )
 
@@ -154,7 +154,7 @@ def create_home_dir_if_missing():
 
 def save_credentials(workspace_name: str, credentials: Credentials):
     create_home_dir_if_missing()
-    if not credentials.access_token and not credentials.api_key:
+    if not credentials.access_token and not credentials.apiKey:
         logger.info("No credentials to save, error")
         return
 

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/common/logger.py

@@ -11,8 +11,10 @@ class ColoredFormatter(logging.Formatter):
     }
 
     def format(self, record):
+        n_spaces = len("CRITICAL") - len(record.levelname)
+        tab = " " * n_spaces
         color = self.COLORS.get(record.levelname, "\033[0m")
-        record.levelname = f"{color}{record.levelname}\033[0m"
+        record.levelname = f"{color}{record.levelname}\033[0m:{tab}"
         return super().format(record)
 
 
@@ -25,5 +27,5 @@ def init(log_level: str):
     logging.getLogger("httpx").propagate = False
 
     handler = logging.StreamHandler()
-    handler.setFormatter(ColoredFormatter("%(levelname)s:\t  %(name)s - %(message)s"))
+    handler.setFormatter(ColoredFormatter(f"%(levelname)s %(name)s - %(message)s"))
     logging.basicConfig(level=log_level, handlers=[handler])

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/common/settings.py

@@ -38,7 +38,7 @@ class SettingsAuthenticationClient(BaseSettings):
 
 
 class SettingsAuthentication(BaseSettings):
-    api_key: Union[None, str] = None
+    apiKey: Union[None, str] = None
     jwt: Union[None, str] = None
     client: SettingsAuthenticationClient = SettingsAuthenticationClient()
 
@@ -64,6 +64,7 @@ class Settings(BaseSettings):
     name: str = Field(default="beamlit-agent")
     base_url: str = Field(default="https://api.beamlit.dev/v0")
     run_url: str = Field(default="https://run.beamlit.dev")
+    mcp_hub_url: str = Field(default="https://mcp-hub-server.beamlit.workers.dev")
     registry_url: str = Field(default="https://serverless-registry-production.beamlit.workers.dev")
     log_level: str = Field(default="INFO")
     agent: SettingsAgent = SettingsAgent()
@@ -117,13 +118,13 @@ def init_agent(
                 functions.append(function)
         settings.agent.functions = functions
 
-    if agent.spec.agent_chain:
-        for chain in agent.spec.agent_chain:
+    if agent.spec.agentChain:
+        for chain in agent.spec.agentChain:
             if chain.enabled:
-                agent_chain = get_agent.sync(chain.name, environment=env, client=client)
+                agentChain = get_agent.sync(chain.name, environment=env, client=client)
                 if chain.description:
-                    agent_chain.spec.description = chain.description
-                agents_chain.append(agent_chain)
+                    agentChain.spec.description = chain.description
+                agents_chain.append(agentChain)
         settings.agent.chain = agents_chain
     if agent.spec.model:
         model = get_model.sync(agent.spec.model, environment=env, client=client)

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/deploy/deploy.py

@@ -6,10 +6,12 @@ from logging import getLogger
 from typing import Literal
 
 from beamlit.common.settings import Settings, get_settings, init
-from beamlit.models import (Agent, EnvironmentMetadata, AgentSpec, AgentChain, Flavor, Function, FunctionSpec,
-                            Runtime)
-from .format import arg_to_dict, format_parameters, format_agent_chain
-from .parser import get_resources, Resource, get_parameters, get_description
+from beamlit.models import (Agent, AgentChain, AgentSpec, EnvironmentMetadata,
+                            Flavor, Function, FunctionSpec, Runtime)
+
+from .format import arg_to_dict, format_agent_chain, format_parameters
+from .parser import Resource, get_description, get_parameters, get_resources
+
 sys.path.insert(0, os.getcwd())
 sys.path.insert(0, os.path.join(os.getcwd(), "src"))
 
@@ -27,8 +29,8 @@ def set_default_values(resource: Resource, deployment: Agent | Function):
     deployment.metadata.environment = settings.environment
     if not deployment.metadata.name:
         deployment.metadata.name = resource.name
-    if not deployment.metadata.display_name:
-        deployment.metadata.display_name = deployment.metadata.name
+    if not deployment.metadata.displayName:
+        deployment.metadata.displayName = deployment.metadata.name
     if not deployment.spec.description:
         deployment.spec.description = get_description(None, resource)
     if not deployment.spec.runtime:
@@ -97,7 +99,7 @@ def get_agent_yaml(
     Generates YAML configuration for an agent deployment.
 
     Args:
-        agent (AgentDeployment): Agent deployment configuration
+        agent (Agent): Agent deployment configuration
         functions (list[tuple[Resource, FunctionDeployment]]): List of associated functions
         settings (Settings): Application settings
 
@@ -109,14 +111,14 @@ apiVersion: beamlit.com/v1alpha1
 kind: Agent
 metadata:
   name: {agent.metadata.name}
-  display_name: {agent.metadata.display_name or agent.metadata.name}
+  displayName: {agent.metadata.displayName or agent.metadata.name}
   environment: {settings.environment}
   workspace: {settings.workspace}
 spec:
   enabled: true
   policies: [{", ".join(agent.spec.policies or [])}]
   functions: [{", ".join([f"{function.metadata.name}" for (_, function) in functions])}]
-  agent_chain: {format_agent_chain(agent.spec.agent_chain)}
+  agentChain: {format_agent_chain(agent.spec.agentChain)}
   model: {agent.spec.model}
   runtime:
     image: {agent.spec.runtime.image}
@@ -143,7 +145,7 @@ apiVersion: beamlit.com/v1alpha1
 kind: Function
 metadata:
   name: {function.metadata.name}
-  display_name: {function.metadata.display_name or function.metadata.name}
+  displayName: {function.metadata.displayName or function.metadata.name}
   environment: {settings.environment}
 spec:
   enabled: true

{beamlit-0.0.24rc21 → beamlit-0.0.26rc22}/src/beamlit/deploy/format.py

@@ -47,21 +47,21 @@ def format_parameters(parameters: list[StoreFunctionParameter]) -> str:
     return "\n".join(formatted)
 
 
-def format_agent_chain(agent_chain: list[AgentChain]) -> str:
+def format_agent_chain(agentChain: list[AgentChain]) -> str:
     """
     Formats agent chain configuration into YAML-compatible string.
 
     Args:
-        agent_chain (list[AgentChain]): List of agent chain configurations
+        agentChain (list[AgentChain]): List of agent chain configurations
 
     Returns:
         str: YAML-formatted string of agent chain
     """
-    if not agent_chain:
+    if not agentChain:
         return "[]"
     formatted = []
 
-    for agent in agent_chain:
+    for agent in agentChain:
         formatted.append(f"""
       - agent: {agent.name}
         enabled: {agent.enabled}""")

beamlit-0.0.26rc22/src/beamlit/functions/mcp/mcp.py

@@ -0,0 +1,125 @@
+import asyncio
+import urllib.parse
+import warnings
+from typing import Any, Callable
+
+import pydantic
+import pydantic_core
+import requests
+import typing_extensions as t
+from beamlit.authentication.authentication import AuthenticatedClient
+from beamlit.common.settings import get_settings
+from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
+from mcp import ListToolsResult
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema as cs
+
+settings = get_settings()
+
+
+def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
+    # Create a new model class that returns our JSON schema.
+    # LangChain requires a BaseModel class.
+    class Schema(pydantic.BaseModel):
+        model_config = pydantic.ConfigDict(extra="allow")
+
+        @t.override
+        @classmethod
+        def __get_pydantic_json_schema__(
+            cls, core_schema: cs.CoreSchema, handler: pydantic.GetJsonSchemaHandler
+        ) -> JsonSchemaValue:
+            return schema
+
+    return Schema
+
+
+class MCPClient:
+    def __init__(self, client: AuthenticatedClient, server_name: str):
+        self.client = client
+        self.server_name = server_name
+        self.headers = {"Api-Key": "1234567890"}
+
+    def list_tools(self) -> requests.Response:
+        client = self.client.get_httpx_client()
+        url = urllib.parse.urljoin(settings.mcp_hub_url, f"{self.server_name}/tools/list")
+
+        response = client.request("GET", url, headers=self.headers)
+        response.raise_for_status()
+        return response
+
+    def call_tool(self, tool_name: str, arguments: dict[str, Any] = None) -> requests.Response:
+        client = self.client.get_httpx_client()
+        url = urllib.parse.urljoin(settings.mcp_hub_url, f"{self.server_name}/tools/call")
+        response = client.request(
+            "POST",
+            url,
+            json={"name": tool_name, "arguments": arguments},
+            headers=self.headers,
+        )
+        response.raise_for_status()
+        return response
+
+class MCPTool(BaseTool):
+    """
+    MCP server tool
+    """
+
+    client: MCPClient
+    handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
+
+    @t.override
+    def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        warnings.warn(
+            "Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
+            stacklevel=1,
+        )
+        return asyncio.run(self._arun(*args, **kwargs))
+
+    @t.override
+    async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        result = self.client.call_tool(self.name, arguments=kwargs)
+        response = result.json()
+        content = pydantic_core.to_json(response["content"]).decode()
+        if response["isError"]:
+            raise ToolException(content)
+        return content
+
+    @t.override
+    @property
+    def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        assert self.args_schema is not None  # noqa: S101
+        return self.args_schema
+
+class MCPToolkit(BaseToolkit):
+    """
+    MCP server toolkit
+    """
+
+    client: MCPClient
+    """The MCP session used to obtain the tools"""
+
+    _tools: ListToolsResult | None = None
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    def initialize(self) -> None:
+        """Initialize the session and retrieve tools list"""
+        if self._tools is None:
+            response = self.client.list_tools()
+            self._tools = ListToolsResult(**response.json())
+
+    @t.override
+    def get_tools(self) -> list[BaseTool]:
+        if self._tools is None:
+            raise RuntimeError("Must initialize the toolkit first")
+
+        return [
+            MCPTool(
+                client=self.client,
+                name=tool.name,
+                description=tool.description or "",
+                args_schema=create_schema_model(tool.inputSchema),
+            )
+            # list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
+            for tool in self._tools.tools
+        ]

beamlit-0.0.26rc22/src/beamlit/functions/remote/remote.py

@@ -0,0 +1,103 @@
+import asyncio
+import urllib.parse
+import warnings
+from typing import Any, Callable
+
+import pydantic
+import pydantic_core
+import typing_extensions as t
+from beamlit.api.functions import get_function
+from beamlit.authentication.authentication import AuthenticatedClient
+from beamlit.models.function import Function
+from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema as cs
+
+
+def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
+    # Create a new model class that returns our JSON schema.
+    # LangChain requires a BaseModel class.
+    class Schema(pydantic.BaseModel):
+        model_config = pydantic.ConfigDict(extra="allow")
+
+        @t.override
+        @classmethod
+        def __get_pydantic_json_schema__(
+            cls, core_schema: cs.CoreSchema, handler: pydantic.GetJsonSchemaHandler
+        ) -> JsonSchemaValue:
+            return schema
+
+    return Schema
+
+
+class RemoteTool(BaseTool):
+    """
+    Remote tool
+    """
+
+    client: AuthenticatedClient
+    handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
+
+    @t.override
+    def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        warnings.warn(
+            "Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy standard tests.",
+            stacklevel=1,
+        )
+        return asyncio.run(self._arun(*args, **kwargs))
+
+    @t.override
+    async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        result = self.client.call_tool(self.name, arguments=kwargs)
+        response = result.json()
+        content = pydantic_core.to_json(response["content"]).decode()
+        if response["isError"]:
+            raise ToolException(content)
+        return content
+
+    @t.override
+    @property
+    def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        assert self.args_schema is not None  # noqa: S101
+        return self.args_schema
+
+class RemoteToolkit(BaseToolkit):
+    """
+    Remote toolkit
+    """
+
+    client: AuthenticatedClient
+    function: str
+    _function: Function | None = None
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    def initialize(self) -> None:
+        """Initialize the session and retrieve tools list"""
+        if self._function is None:
+            self._function = get_function(self.function, client=self.client)
+
+    @t.override
+    def get_tools(self) -> list[BaseTool]:
+        if self._tools is None:
+            raise RuntimeError("Must initialize the toolkit first")
+
+        if self._function.spec.kit:
+            return [
+                RemoteTool(
+                client=self.client,
+                name=func.name,
+                description=func.description or "",
+                args_schema=create_schema_model(func.parameters),
+                )
+                for func in self._function.spec.kit
+            ]
+
+        return [
+            RemoteTool(
+                client=self.client,
+                name=self._function.metadata.name,
+                description=self._function.spec.description or "",
+                args_schema=create_schema_model(self._function.spec.parameters),
+            )
+        ]