blaxel 0.64.0__py3-none-any.whl

blaxel/functions/mcp/mcp.py

@@ -0,0 +1,168 @@
+"""
+This module provides functionalities to interact with MCP (Multi-Client Platform) servers.
+It includes classes for managing MCP clients, creating dynamic schemas, and integrating MCP tools into Blaxel.
+"""
+
+import asyncio
+import logging
+import warnings
+from typing import Any, AsyncIterator, Callable
+
+import pydantic
+import pydantic_core
+import requests
+import typing_extensions as t
+from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
+from mcp import ClientSession
+from mcp.types import CallToolResult, ListToolsResult
+
+from blaxel.aimon.settings import get_settings
+from blaxel.authentication import get_authentication_headers
+from blaxel.authentication.authentication import AuthenticatedClient
+from blaxel.functions.mcp.client import websocket_client
+
+from .utils import create_schema_model
+
+settings = get_settings()
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClient:
+    def __init__(self, client: AuthenticatedClient, url: str, fallback_url: str | None = None):
+        self.client = client
+        self.url = url
+        self.fallback_url = fallback_url
+
+    async def list_ws_tools(self, is_fallback: bool = False) -> ListToolsResult:
+        if is_fallback:
+            url = self.fallback_url
+        else:
+            url = self.url
+        try:
+            async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
+                logger.debug("WebSocket connection established")
+                async with ClientSession(read_stream, write_stream) as client:
+                    await client.initialize()
+                    response = await client.list_tools()
+                    logger.debug(f"WebSocket tools: {response}")
+                    return response
+        except Exception as e:
+            logger.error(f"Error listing tools: {e}")
+            logger.debug("WebSocket not available, trying HTTP")
+            return None  # Signal to list_tools() to try HTTP instead
+
+    async def list_tools(self) -> ListToolsResult:
+        logger.debug(f"Listing tools for {self.url}")
+        try:
+            result = await self.list_ws_tools(is_fallback=False)
+            return result
+        except Exception as e: # Fallback to Public endpoint
+            if self.fallback_url:
+                try:
+                    result = await self.list_ws_tools(is_fallback=True)
+                    return result
+                except Exception as e:
+                    raise e
+            else:
+                raise e
+
+
+    async def call_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any] = None,
+        is_fallback: bool = False,
+    ) -> requests.Response | AsyncIterator[CallToolResult]:
+        if is_fallback:
+            url = self.fallback_url
+        else:
+            url = self.url
+        try:
+            async with websocket_client(url, headers=get_authentication_headers(settings)) as (read_stream, write_stream):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    response = await session.call_tool(tool_name, arguments or {})
+                    content = pydantic_core.to_json(response).decode()
+                    if response.isError:
+                        raise ToolException(content)
+                    return content
+        except Exception as e:
+            raise e
+
+
+class MCPTool(BaseTool):
+    """
+    Tool for interacting with MCP server-hosted tools.
+
+    Attributes:
+        client (MCPClient): The MCP client instance.
+        handle_tool_error (bool | str | Callable[[ToolException], str] | None): Error handling strategy.
+    """
+
+    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:
+        try:
+            return await self.client.call_tool(self.name, arguments=kwargs)
+        except Exception as e:
+            if self.client.fallback_url:
+                try:
+                    return await self.client.call_tool(self.name, arguments=kwargs, is_fallback=True) # Fallback to Public endpoint
+                except Exception as e:
+                    raise e
+            else:
+                raise e
+
+    @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):
+    """
+    Toolkit for managing MCP server tools.
+
+    Attributes:
+        client (MCPClient): The MCP client instance.
+        _tools (ListToolsResult | None): Cached list of tools from the MCP server.
+    """
+
+    client: MCPClient
+    """The MCP session used to obtain the tools"""
+
+    _tools: ListToolsResult | None = None
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    async def initialize(self) -> None:
+        """Initialize the session and retrieve tools list"""
+        if self._tools is None:
+            response = await self.client.list_tools()
+            self._tools = response
+
+    @t.override
+    async 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.name, 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
+        ]

blaxel/functions/mcp/utils.py

@@ -0,0 +1,56 @@
+"""
+This module provides functionalities to interact with MCP (Multi-Client Platform) servers.
+It includes classes for managing MCP clients, creating dynamic schemas, and integrating MCP tools into Blaxel.
+"""
+import pydantic
+import typing_extensions as t
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema as cs
+
+TYPE_MAP = {
+    "integer": int,
+    "number": float,
+    "array": list,
+    "object": dict,
+    "boolean": bool,
+    "string": str,
+    "null": type(None),
+}
+
+FIELD_DEFAULTS = {
+    int: 0,
+    float: 0.0,
+    list: [],
+    bool: False,
+    str: "",
+    type(None): None,
+}
+
+def configure_field(name: str, type_: dict[str, t.Any], required: list[str]) -> tuple[type, t.Any]:
+    field_type = TYPE_MAP[type_["type"]]
+    default_ = FIELD_DEFAULTS.get(field_type) if name not in required else ...
+    return field_type, default_
+
+def create_schema_model(name: str, schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
+    # Create a new model class that returns our JSON schema.
+    # LangChain requires a BaseModel class.
+    class SchemaBase(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
+
+    # Since this langchain patch, we need to synthesize pydantic fields from the schema
+    # https://github.com/langchain-ai/langchain/commit/033ac417609297369eb0525794d8b48a425b8b33
+    required = schema.get("required", [])
+    fields: dict[str, t.Any] = {
+        name: configure_field(name, type_, required) for name, type_ in schema["properties"].items()
+    }
+
+    return pydantic.create_model(f"{name}Schema", __base__=SchemaBase, **fields)
+
+

blaxel/functions/remote/remote.py

@@ -0,0 +1,183 @@
+"""
+This module provides functionalities to integrate remote functions into Blaxel.
+It includes classes for creating dynamic schemas based on function parameters and managing remote toolkits.
+"""
+
+import asyncio
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Callable
+
+import pydantic
+import typing_extensions as t
+from langchain_core.tools.base import BaseTool, ToolException
+
+from blaxel.aimon.settings import get_settings
+from blaxel.api.functions import get_function, list_functions
+from blaxel.authentication.authentication import AuthenticatedClient
+from blaxel.errors import UnexpectedStatus
+from blaxel.functions.mcp.mcp import MCPClient, MCPToolkit
+from blaxel.models import Function, StoreFunctionParameter
+from blaxel.run import RunClient
+
+
+def create_dynamic_schema(name: str, parameters: list[StoreFunctionParameter]) -> type[pydantic.BaseModel]:
+    """
+    Creates a dynamic Pydantic schema based on function parameters.
+
+    Args:
+        name (str): The name of the schema.
+        parameters (list[StoreFunctionParameter]): List of parameter objects.
+
+    Returns:
+        type[pydantic.BaseModel]: The dynamically created Pydantic model.
+    """
+    field_definitions = {}
+    for param in parameters:
+        field_type = str
+        if param.type_ == "number":
+            field_type = float
+        elif param.type_ == "integer":
+            field_type = int
+        elif param.type_ == "boolean":
+            field_type = bool
+
+        field_definitions[param.name] = (
+            field_type,
+            pydantic.Field(description=param.description or "")
+        )
+    return pydantic.create_model(
+        f"{name}Schema",
+        **field_definitions
+    )
+
+
+class RemoteTool(BaseTool):
+    """
+    Tool for interacting with remote functions.
+
+    Attributes:
+        client (RunClient): The client used to execute remote function calls.
+        resource_name (str): The name of the remote resource.
+        kit (bool): Indicates whether the tool is part of a function kit.
+        handle_tool_error (bool | str | Callable[[ToolException], str] | None): Error handling strategy.
+    """
+
+    client: RunClient
+    resource_name: str
+    kit: bool = False
+    handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
+    service_name: str | None = None
+    cloud: bool = False
+    @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:
+        body = {**kwargs}
+        if self.kit:
+            body["name"] = self.name
+        result = self.client.run(
+            "function",
+            self.resource_name,
+            "POST",
+            cloud=self.cloud,
+            service_name=self.service_name,
+            json=body,
+        )
+        return result.text
+
+    @t.override
+    @property
+    def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        assert self.args_schema is not None  # noqa: S101
+        return self.args_schema
+
+@dataclass
+class RemoteToolkit:
+    """
+    Toolkit for managing remote function tools.
+
+    Attributes:
+        client (AuthenticatedClient): The authenticated client instance.
+        function (str): The name of the remote function to integrate.
+        _function (Function | None): Cached Function object after initialization.
+    """
+    client: AuthenticatedClient
+    function: str
+    _function: Function | None = None
+    _service_name: str | None = None
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    async def initialize(self) -> None:
+        """Initialize the session and retrieve the remote function details."""
+        if self._function is None:
+            try:
+                response = get_function.sync_detailed(self.function, client=self.client)
+                function_name = self.function.upper().replace("-", "_")
+                if os.getenv(f"BL_FUNCTION_{function_name}_SERVICE_NAME"):
+                    self._service_name = os.getenv(f"BL_FUNCTION_{function_name}_SERVICE_NAME")
+                self._function = response.parsed
+            except UnexpectedStatus as e:
+                functions = list_functions.sync_detailed(
+                    client=self.client,
+                ).parsed
+                names = [
+                    f.metadata.name
+                    for f in functions
+                ]
+                raise RuntimeError(
+                    f"error: {e.status_code}. Available functions: {', '.join(names)}"
+                )
+
+    async def get_tools(self) -> list[BaseTool]:
+        settings = get_settings()
+        if self._function is None:
+            raise RuntimeError("Must initialize the toolkit first")
+
+        if self._function.spec.integration_connections:
+            fallback_url = None
+            url = f"{settings.run_url}/{settings.workspace}/functions/{self._function.metadata.name}"
+            if self._service_name:
+                fallback_url = f"https://{self._service_name}.{settings.run_internal_hostname}"
+                url = f"https://{self._service_name}.{settings.run_internal_hostname}"
+            mcp_client = MCPClient(self.client, url, fallback_url)
+            mcp_toolkit = MCPToolkit(client=mcp_client, url=url)
+            await mcp_toolkit.initialize()
+            return await mcp_toolkit.get_tools()
+
+        if self._function.spec.kit:
+            return [
+                RemoteTool(
+                    client=RunClient(self.client),
+                    name=func.name,
+                    resource_name=self._function.metadata.name,
+                    kit=True,
+                    description=func.description or "",
+                    args_schema=create_dynamic_schema(func.name, func.parameters),
+                    cloud=settings.cloud,
+                    service_name=self._service_name,
+                )
+                for func in self._function.spec.kit
+            ]
+
+        return [
+            RemoteTool(
+                client=RunClient(self.client),
+                name=self._function.metadata.name,
+                resource_name=self._function.metadata.name,
+                description=self._function.spec.description or "",
+                args_schema=create_dynamic_schema(
+                    self._function.metadata.name,
+                    self._function.spec.parameters
+                ),
+                cloud=settings.cloud,
+                service_name=self._service_name,
+            )
+        ]

blaxel/models/__init__.py

@@ -0,0 +1,233 @@
+"""Contains all the data models used in inputs/outputs"""
+
+from .acl import ACL
+from .agent import Agent
+from .agent_chain import AgentChain
+from .agent_spec import AgentSpec
+from .api_key import ApiKey
+from .configuration import Configuration
+from .continent import Continent
+from .core_event import CoreEvent
+from .core_spec import CoreSpec
+from .core_spec_configurations import CoreSpecConfigurations
+from .country import Country
+from .create_api_key_for_service_account_body import CreateApiKeyForServiceAccountBody
+from .create_workspace_service_account_body import CreateWorkspaceServiceAccountBody
+from .create_workspace_service_account_response_200 import CreateWorkspaceServiceAccountResponse200
+from .delete_workspace_service_account_response_200 import DeleteWorkspaceServiceAccountResponse200
+from .entrypoint import Entrypoint
+from .entrypoint_env import EntrypointEnv
+from .flavor import Flavor
+from .form import Form
+from .form_config import FormConfig
+from .form_oauthomitempty import FormOauthomitempty
+from .form_secrets import FormSecrets
+from .function import Function
+from .function_kit import FunctionKit
+from .function_spec import FunctionSpec
+from .get_trace_ids_response_200 import GetTraceIdsResponse200
+from .get_trace_logs_response_200 import GetTraceLogsResponse200
+from .get_trace_response_200 import GetTraceResponse200
+from .get_workspace_service_accounts_response_200_item import (
+    GetWorkspaceServiceAccountsResponse200Item,
+)
+from .histogram_bucket import HistogramBucket
+from .histogram_stats import HistogramStats
+from .integration_connection import IntegrationConnection
+from .integration_connection_spec import IntegrationConnectionSpec
+from .integration_connection_spec_config import IntegrationConnectionSpecConfig
+from .integration_connection_spec_secret import IntegrationConnectionSpecSecret
+from .integration_model import IntegrationModel
+from .integration_repository import IntegrationRepository
+from .invite_workspace_user_body import InviteWorkspaceUserBody
+from .knowledgebase import Knowledgebase
+from .knowledgebase_spec import KnowledgebaseSpec
+from .knowledgebase_spec_options import KnowledgebaseSpecOptions
+from .last_n_requests_metric import LastNRequestsMetric
+from .latency_metric import LatencyMetric
+from .location_response import LocationResponse
+from .mcp_definition import MCPDefinition
+from .mcp_definition_entrypoint import MCPDefinitionEntrypoint
+from .mcp_definition_form import MCPDefinitionForm
+from .metadata import Metadata
+from .metadata_labels import MetadataLabels
+from .metric import Metric
+from .metrics import Metrics
+from .metrics_models import MetricsModels
+from .metrics_request_total_per_code import MetricsRequestTotalPerCode
+from .metrics_rps_per_code import MetricsRpsPerCode
+from .model import Model
+from .model_private_cluster import ModelPrivateCluster
+from .model_spec import ModelSpec
+from .o_auth import OAuth
+from .owner_fields import OwnerFields
+from .pending_invitation import PendingInvitation
+from .pending_invitation_accept import PendingInvitationAccept
+from .pending_invitation_render import PendingInvitationRender
+from .pending_invitation_render_invited_by import PendingInvitationRenderInvitedBy
+from .pending_invitation_render_workspace import PendingInvitationRenderWorkspace
+from .pending_invitation_workspace_details import PendingInvitationWorkspaceDetails
+from .pod_template_spec import PodTemplateSpec
+from .policy import Policy
+from .policy_location import PolicyLocation
+from .policy_max_tokens import PolicyMaxTokens
+from .policy_spec import PolicySpec
+from .private_cluster import PrivateCluster
+from .private_location import PrivateLocation
+from .repository import Repository
+from .request_duration_over_time_metric import RequestDurationOverTimeMetric
+from .request_duration_over_time_metrics import RequestDurationOverTimeMetrics
+from .request_total_by_origin_metric import RequestTotalByOriginMetric
+from .request_total_by_origin_metric_request_total_by_origin import (
+    RequestTotalByOriginMetricRequestTotalByOrigin,
+)
+from .request_total_by_origin_metric_request_total_by_origin_and_code import (
+    RequestTotalByOriginMetricRequestTotalByOriginAndCode,
+)
+from .request_total_metric import RequestTotalMetric
+from .request_total_metric_request_total_per_code import RequestTotalMetricRequestTotalPerCode
+from .request_total_metric_rps_per_code import RequestTotalMetricRpsPerCode
+from .resource_log import ResourceLog
+from .resource_metrics import ResourceMetrics
+from .resource_metrics_request_total_per_code import ResourceMetricsRequestTotalPerCode
+from .resource_metrics_rps_per_code import ResourceMetricsRpsPerCode
+from .revision_configuration import RevisionConfiguration
+from .revision_metadata import RevisionMetadata
+from .runtime import Runtime
+from .runtime_startup_probe import RuntimeStartupProbe
+from .serverless_config import ServerlessConfig
+from .spec_configuration import SpecConfiguration
+from .store_agent import StoreAgent
+from .store_agent_labels import StoreAgentLabels
+from .store_configuration import StoreConfiguration
+from .store_configuration_option import StoreConfigurationOption
+from .store_function import StoreFunction
+from .store_function_kit import StoreFunctionKit
+from .store_function_labels import StoreFunctionLabels
+from .store_function_parameter import StoreFunctionParameter
+from .time_fields import TimeFields
+from .token_rate_metric import TokenRateMetric
+from .token_rate_metrics import TokenRateMetrics
+from .token_total_metric import TokenTotalMetric
+from .trace_ids_response import TraceIdsResponse
+from .update_workspace_service_account_body import UpdateWorkspaceServiceAccountBody
+from .update_workspace_service_account_response_200 import UpdateWorkspaceServiceAccountResponse200
+from .update_workspace_user_role_body import UpdateWorkspaceUserRoleBody
+from .websocket_channel import WebsocketChannel
+from .workspace import Workspace
+from .workspace_labels import WorkspaceLabels
+from .workspace_user import WorkspaceUser
+
+__all__ = (
+    "ACL",
+    "Agent",
+    "AgentChain",
+    "AgentSpec",
+    "ApiKey",
+    "Configuration",
+    "Continent",
+    "CoreEvent",
+    "CoreSpec",
+    "CoreSpecConfigurations",
+    "Country",
+    "CreateApiKeyForServiceAccountBody",
+    "CreateWorkspaceServiceAccountBody",
+    "CreateWorkspaceServiceAccountResponse200",
+    "DeleteWorkspaceServiceAccountResponse200",
+    "Entrypoint",
+    "EntrypointEnv",
+    "Flavor",
+    "Form",
+    "FormConfig",
+    "FormOauthomitempty",
+    "FormSecrets",
+    "Function",
+    "FunctionKit",
+    "FunctionSpec",
+    "GetTraceIdsResponse200",
+    "GetTraceLogsResponse200",
+    "GetTraceResponse200",
+    "GetWorkspaceServiceAccountsResponse200Item",
+    "HistogramBucket",
+    "HistogramStats",
+    "IntegrationConnection",
+    "IntegrationConnectionSpec",
+    "IntegrationConnectionSpecConfig",
+    "IntegrationConnectionSpecSecret",
+    "IntegrationModel",
+    "IntegrationRepository",
+    "InviteWorkspaceUserBody",
+    "Knowledgebase",
+    "KnowledgebaseSpec",
+    "KnowledgebaseSpecOptions",
+    "LastNRequestsMetric",
+    "LatencyMetric",
+    "LocationResponse",
+    "MCPDefinition",
+    "MCPDefinitionEntrypoint",
+    "MCPDefinitionForm",
+    "Metadata",
+    "MetadataLabels",
+    "Metric",
+    "Metrics",
+    "MetricsModels",
+    "MetricsRequestTotalPerCode",
+    "MetricsRpsPerCode",
+    "Model",
+    "ModelPrivateCluster",
+    "ModelSpec",
+    "OAuth",
+    "OwnerFields",
+    "PendingInvitation",
+    "PendingInvitationAccept",
+    "PendingInvitationRender",
+    "PendingInvitationRenderInvitedBy",
+    "PendingInvitationRenderWorkspace",
+    "PendingInvitationWorkspaceDetails",
+    "PodTemplateSpec",
+    "Policy",
+    "PolicyLocation",
+    "PolicyMaxTokens",
+    "PolicySpec",
+    "PrivateCluster",
+    "PrivateLocation",
+    "Repository",
+    "RequestDurationOverTimeMetric",
+    "RequestDurationOverTimeMetrics",
+    "RequestTotalByOriginMetric",
+    "RequestTotalByOriginMetricRequestTotalByOrigin",
+    "RequestTotalByOriginMetricRequestTotalByOriginAndCode",
+    "RequestTotalMetric",
+    "RequestTotalMetricRequestTotalPerCode",
+    "RequestTotalMetricRpsPerCode",
+    "ResourceLog",
+    "ResourceMetrics",
+    "ResourceMetricsRequestTotalPerCode",
+    "ResourceMetricsRpsPerCode",
+    "RevisionConfiguration",
+    "RevisionMetadata",
+    "Runtime",
+    "RuntimeStartupProbe",
+    "ServerlessConfig",
+    "SpecConfiguration",
+    "StoreAgent",
+    "StoreAgentLabels",
+    "StoreConfiguration",
+    "StoreConfigurationOption",
+    "StoreFunction",
+    "StoreFunctionKit",
+    "StoreFunctionLabels",
+    "StoreFunctionParameter",
+    "TimeFields",
+    "TokenRateMetric",
+    "TokenRateMetrics",
+    "TokenTotalMetric",
+    "TraceIdsResponse",
+    "UpdateWorkspaceServiceAccountBody",
+    "UpdateWorkspaceServiceAccountResponse200",
+    "UpdateWorkspaceUserRoleBody",
+    "WebsocketChannel",
+    "Workspace",
+    "WorkspaceLabels",
+    "WorkspaceUser",
+)