flock-core 0.4.2__py3-none-any.whl → 0.4.5__py3-none-any.whl

flock/core/mcp/mcp_client_manager.py

@@ -0,0 +1,201 @@
+"""Manages a pool of connections for a particular server."""
+
+import copy
+from abc import ABC, abstractmethod
+from asyncio import Lock
+from typing import Any, Generic, TypeVar
+
+from opentelemetry import trace
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+
+from flock.core.logging.logging import get_logger
+from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
+from flock.core.mcp.mcp_client import (
+    FlockMCPClientBase,
+)
+from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
+
+logger = get_logger("core.mcp.connection_manager_base")
+tracer = trace.get_tracer(__name__)
+
+TClient = TypeVar("TClient", bound="FlockMCPClientBase")
+
+
+class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
+    """Handles a Pool of MCPClients of type TClient."""
+
+    client_config: FlockMCPConfigurationBase = Field(
+        ..., description="Configuration for clients."
+    )
+
+    lock: Lock = Field(
+        default_factory=Lock,
+        description="Lock for mutex access.",
+        exclude=True,
+    )
+
+    clients: dict[str, dict[str, FlockMCPClientBase]] = Field(
+        default_factory=dict,
+        exclude=True,
+        description="Internal Store for the clients.",
+    )
+
+    # --- Pydantic v2 Configuratioin ---
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    @abstractmethod
+    async def make_client(
+        self,
+        additional_params: dict[str, Any] | None = None,
+    ) -> type[TClient]:
+        """Instantiate-but don't connect yet-a fresh client of the concrete subtype."""
+        # default implementation
+        pass
+
+    async def get_client(
+        self,
+        agent_id: str,
+        run_id: str,
+        additional_params: dict[str, Any] | None = None,
+    ) -> type[TClient]:
+        """Provides a client from the pool."""
+        # Attempt to get a client from the client store.
+        # clients are stored like this: agent_id -> run_id -> client
+        with tracer.start_as_current_span("client_manager.get_client") as span:
+            span.set_attribute("agent_id", agent_id)
+            span.set_attribute("run_id", run_id)
+            async with self.lock:
+                try:
+                    logger.debug(
+                        f"Attempting to get client for server '{self.client_config.name}'"
+                    )
+                    refresh = False
+                    if additional_params:
+                        refresh = bool(
+                            additional_params.get("refresh_client", False)
+                        )
+                    client = None
+                    run_clients = self.clients.get(agent_id, None)
+                    if run_clients is None or refresh:
+                        # This means, that across all runs, no agent has ever needed a client.
+                        # This also means that we need to create a client.
+                        client = await self.make_client(
+                            additional_params=copy.deepcopy(additional_params)
+                        )
+                        # Insert the freshly created client
+                        self.clients[agent_id] = {}
+                        self.clients[agent_id][run_id] = client
+
+                    else:
+                        # This means there is at least one entry for the agent_id available
+                        # Now, all we need to do is check if the run_id matches the entrie's run_id
+                        client = run_clients.get(run_id, None)
+                        if client is None or refresh:
+                            # Means no client here with the respective run_id
+                            client = await self.make_client(
+                                additional_params=copy.deepcopy(
+                                    additional_params
+                                )
+                            )
+                            # Insert the freshly created client.
+                            self.clients[agent_id][run_id] = client
+
+                    return client
+                except Exception as e:
+                    # Log the exception and raise it so it becomes visible downstream
+                    logger.error(
+                        f"Unexpected Exception ocurred while trying to get client for server '{self.client_config.name}' with agent_id: {agent_id} and run_id: {run_id}: {e}"
+                    )
+                    span.record_exception(e)
+                    raise e
+
+    async def call_tool(
+        self,
+        agent_id: str,
+        run_id: str,
+        name: str,
+        arguments: dict[str, Any],
+        additional_params: dict[str, Any] | None = None,
+    ) -> Any:
+        """Call a tool."""
+        with tracer.start_as_current_span("client_manager.call_tool") as span:
+            span.set_attribute("agent_id", agent_id)
+            span.set_attribute("run_id", run_id)
+            span.set_attribute("tool_name", name)
+            span.set_attribute("arguments", str(arguments))
+            try:
+                client = await self.get_client(
+                    agent_id=agent_id,
+                    run_id=run_id,
+                    additional_params=additional_params,
+                )
+                result = await client.call_tool(
+                    agent_id=agent_id,
+                    run_id=run_id,
+                    name=name,
+                    arguments=arguments,
+                )
+                return result
+            except Exception as e:
+                logger.error(
+                    f"Exception occurred while trying to call tool {name} on server '{self.client_config.name}': {e}"
+                )
+                span.record_exception(e)
+                return None
+
+    async def get_tools(
+        self,
+        agent_id: str,
+        run_id: str,
+        additional_params: dict[str, Any] | None = None,
+    ) -> list[FlockMCPToolBase]:
+        """Retrieves a list of tools for the agents to act on."""
+        with tracer.start_as_current_span("client_manager.get_tools") as span:
+            span.set_attribute("agent_id", agent_id)
+            span.set_attribute("run_id", run_id)
+            try:
+                client = await self.get_client(
+                    agent_id=agent_id,
+                    run_id=run_id,
+                    additional_params=additional_params,
+                )
+                tools: list[FlockMCPToolBase] = await client.get_tools(
+                    agent_id=agent_id, run_id=run_id
+                )
+                return tools
+            except Exception as e:
+                logger.error(
+                    f"Exception occurred while trying to retrieve Tools for server '{self.client_config.name}' with agent_id: {agent_id} and run_id: {run_id}: {e}"
+                )
+                span.record_exception(e)
+                return []
+
+    async def close_all(self) -> None:
+        """Closes all connections in the pool and cancels background tasks."""
+        with tracer.start_as_current_span("client_manager.close_all") as span:
+            async with self.lock:
+                for agent_id, run_dict in self.clients.items():
+                    logger.debug(
+                        f"Shutting down all clients for agent_id: {agent_id}"
+                    )
+                    for run_id, client in run_dict.items():
+                        logger.debug(
+                            f"Shutting down client for agent_id {agent_id} and run_id {run_id}"
+                        )
+                        try:
+                            await client.disconnect()
+                        except Exception as e:
+                            logger.error(
+                                f"Error when trying to disconnect client for server '{self.client_config.name}': {e}"
+                            )
+                            span.record_exception(e)
+                self.clients = {}  # Let the GC take care of the rest.
+                logger.info(
+                    f"All clients disconnected for server '{self.client_config.name}'"
+                )

flock/core/mcp/mcp_config.py

@@ -0,0 +1,237 @@
+"""Base Config for MCP Clients."""
+
+from typing import Literal, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from flock.core.mcp.types.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    MCPRoot,
+    ServerParameters,
+)
+
+LoggingLevel = Literal[
+    "debug",
+    "info",
+    "notice",
+    "warning",
+    "error",
+    "critical",
+    "alert",
+    "emergency",
+]
+
+
+A = TypeVar("A", bound="FlockMCPCallbackConfigurationBase")
+B = TypeVar("B", bound="FlockMCPConnectionConfigurationBase")
+C = TypeVar("C", bound="FlockMCPConfigurationBase")
+D = TypeVar("D", bound="FlockMCPCachingConfigurationBase")
+E = TypeVar("E", bound="FlockMCPFeatureConfigurationBase")
+
+
+class FlockMCPCachingConfigurationBase(BaseModel):
+    """Configuration for Caching in Clients."""
+
+    tool_cache_max_size: float = Field(
+        default=100, description="Maximum number of items in the Tool Cache."
+    )
+
+    tool_cache_max_ttl: float = Field(
+        default=60,
+        description="Max TTL for items in the tool cache in seconds.",
+    )
+
+    resource_contents_cache_max_size: float = Field(
+        default=10,
+        description="Maximum number of entries in the Resource Contents cache.",
+    )
+
+    resource_contents_cache_max_ttl: float = Field(
+        default=60 * 5,
+        description="Maximum number of items in the Resource Contents cache.",
+    )
+
+    resource_list_cache_max_size: float = Field(
+        default=10,
+        description="Maximum number of entries in the Resource List Cache.",
+    )
+
+    resource_list_cache_max_ttl: float = Field(
+        default=100,
+        description="Maximum TTL for entries in the Resource List Cache.",
+    )
+
+    tool_result_cache_max_size: float = Field(
+        default=1000,
+        description="Maximum number of entries in the Tool Result Cache.",
+    )
+
+    tool_result_cache_max_ttl: float = Field(
+        default=20,
+        description="Maximum TTL in seconds for entries in the Tool Result Cache.",
+    )
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    @classmethod
+    def with_fields(cls: type[D], **field_definitions) -> type[D]:
+        """Create a new config class with additional fields."""
+        return create_model(
+            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
+        )
+
+
+class FlockMCPCallbackConfigurationBase(BaseModel):
+    """Base Configuration Class for Callbacks for Clients."""
+
+    sampling_callback: FlockSamplingMCPCallback | None = Field(
+        default=None,
+        description="Callback for handling sampling requests from an external server.",
+    )
+
+    list_roots_callback: FlockListRootsMCPCallback | None = Field(
+        default=None, description="Callback for handling list roots requests."
+    )
+
+    logging_callback: FlockLoggingMCPCallback | None = Field(
+        default=None,
+        description="Callback for handling logging messages from an external server.",
+    )
+
+    message_handler: FlockMessageHandlerMCPCallback | None = Field(
+        default=None,
+        description="Callback for handling messages not covered by other callbacks.",
+    )
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+
+    @classmethod
+    def with_fields(cls: type[A], **field_definitions) -> type[A]:
+        """Create a new config class with additional fields."""
+        return create_model(
+            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
+        )
+
+
+class FlockMCPConnectionConfigurationBase(BaseModel):
+    """Base Configuration Class for Connection Parameters for a client."""
+
+    max_retries: int = Field(
+        default=3,
+        description="How many times to attempt to establish the connection before giving up.",
+    )
+
+    connection_parameters: ServerParameters = Field(
+        ..., description="Connection parameters for the server."
+    )
+
+    transport_type: Literal["stdio", "websockets", "sse", "custom"] = Field(
+        ..., description="Type of transport to use."
+    )
+
+    mount_points: list[MCPRoot] | None = Field(
+        default=None, description="Initial Mountpoints to operate under."
+    )
+
+    read_timeout_seconds: float | int = Field(
+        default=60 * 5, description="Read Timeout."
+    )
+
+    server_logging_level: LoggingLevel = Field(
+        default="error",
+        description="The logging level for logging events from the remote server.",
+    )
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+
+    @classmethod
+    def with_fields(cls: type[B], **field_definitions) -> type[B]:
+        """Create a new config class with additional fields."""
+        return create_model(
+            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
+        )
+
+
+class FlockMCPFeatureConfigurationBase(BaseModel):
+    """Base Configuration Class for switching MCP Features on and off."""
+
+    roots_enabled: bool = Field(
+        default=False,
+        description="Whether or not the Roots feature is enabled for this client.",
+    )
+
+    sampling_enabled: bool = Field(
+        default=False,
+        description="Whether or not the Sampling feature is enabled for this client.",
+    )
+
+    tools_enabled: bool = Field(
+        default=False,
+        description="Whether or not the Tools feature is enabled for this client.",
+    )
+
+    prompts_enabled: bool = Field(
+        default=False,
+        description="Whether or not the Prompts feature is enabled for this client.",
+    )
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    @classmethod
+    def with_fields(cls: type[E], **field_definitions) -> type[E]:
+        """Create a new config class with additional fields."""
+        return create_model(
+            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
+        )
+
+
+class FlockMCPConfigurationBase(BaseModel):
+    """Base Configuration Class for MCP Clients.
+
+    Each Client should implement their own config
+    model by inheriting from this class.
+    """
+
+    name: str = Field(
+        ..., description="Name of the server the client connects to."
+    )
+
+    connection_config: FlockMCPConnectionConfigurationBase = Field(
+        ..., description="MCP Connection Configuration for a client."
+    )
+
+    caching_config: FlockMCPCachingConfigurationBase = Field(
+        default_factory=FlockMCPCachingConfigurationBase,
+        description="Configuration for the internal caches of the client.",
+    )
+
+    callback_config: FlockMCPCallbackConfigurationBase = Field(
+        default_factory=FlockMCPCallbackConfigurationBase,
+        description="Callback configuration for the client.",
+    )
+
+    feature_config: FlockMCPFeatureConfigurationBase = Field(
+        default_factory=FlockMCPFeatureConfigurationBase,
+        description="Feature configuration for the client.",
+    )
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    @classmethod
+    def with_fields(cls: type[C], **field_definitions) -> type[C]:
+        """Create a new config class with additional fields."""
+        return create_model(
+            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
+        )

flock/core/mcp/types/__init__.py

@@ -0,0 +1 @@
+"""MCP Types package."""

flock/core/mcp/types/callbacks.py

@@ -0,0 +1,86 @@
+"""MCP Callbacks."""
+
+from mcp.shared.context import RequestContext
+from mcp.shared.session import RequestResponder
+from mcp.types import (
+    INVALID_REQUEST,
+    ClientResult,
+    CreateMessageRequestParams,
+    ErrorData,
+    ListRootsResult,
+    ServerRequest,
+)
+
+from flock.core.logging.logging import FlockLogger
+from flock.core.mcp.mcp_client import Any
+from flock.core.mcp.types.handlers import (
+    handle_incoming_exception,
+    handle_incoming_request,
+    handle_incoming_server_notification,
+    handle_logging_message,
+)
+from flock.core.mcp.types.types import (
+    FlockLoggingMessageNotificationParams,
+    ServerNotification,
+)
+
+
+async def default_sampling_callback(
+    ctx: RequestContext, params: CreateMessageRequestParams, logger: FlockLogger
+) -> ErrorData:
+    """Default Callback for Sampling."""
+    logger.info(f"Rejecting Sampling Request.")
+    return ErrorData(code=INVALID_REQUEST, message="Sampling not supported.")
+
+
+async def default_message_handler(
+    req: RequestResponder[ServerRequest, ClientResult]
+    | ServerNotification
+    | Exception,
+    logger: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Default Message Handler."""
+    if isinstance(req, Exception):
+        await handle_incoming_exception(
+            e=req,
+            logger_to_use=logger,
+            associated_client=associated_client,
+        )
+    elif isinstance(req, ServerNotification):
+        await handle_incoming_server_notification(
+            n=req,
+            logger=logger,
+            client=associated_client,
+        )
+    elif isinstance(req, RequestResponder[ServerRequest, ClientResult]):
+        await handle_incoming_request(
+            req=req,
+            logger_to_use=logger,
+            associated_client=associated_client,
+        )
+
+
+async def default_list_roots_callback(
+    associated_client: Any,
+    logger: FlockLogger,
+) -> ListRootsResult | ErrorData:
+    """Default List Roots Callback."""
+    if associated_client.config.feature_config.roots_enabled:
+        current_roots = await associated_client.get_current_roots()
+        return ListRootsResult(roots=current_roots)
+    else:
+        return ErrorData(
+            code=INVALID_REQUEST, message="List roots not supported."
+        )
+
+
+async def default_logging_callback(
+    params: FlockLoggingMessageNotificationParams,
+    logger: FlockLogger,
+    server_name: str,
+) -> None:
+    """Default Logging Handling Callback."""
+    await handle_logging_message(
+        params=params, logger=logger, server_name=server_name
+    )

flock/core/mcp/types/factories.py

@@ -0,0 +1,111 @@
+"""Factories for default MCP Callbacks."""
+
+from typing import TYPE_CHECKING, Any
+
+from mcp.shared.context import RequestContext
+from mcp.types import (
+    CreateMessageRequestParams,
+)
+
+from flock.core.logging.logging import FlockLogger, get_logger
+from flock.core.mcp.types.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockLoggingMessageNotificationParams,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    ServerNotification,
+)
+
+if TYPE_CHECKING:
+    from flock.core.mcp.types.callbacks import (
+        default_list_roots_callback,
+        default_logging_callback,
+        default_message_handler,
+        default_sampling_callback,
+    )
+
+default_logging_callback_logger = get_logger("core.mcp.callback.logging")
+default_sampling_callback_logger = get_logger("core.mcp.callback.sampling")
+default_list_roots_callback_logger = get_logger("core.mcp.callback.sampling")
+default_message_handler_logger = get_logger("core.mcp.callback.message")
+
+
+def default_flock_mcp_logging_callback_factory(
+    associated_client: Any,
+    logger: FlockLogger | None = None,
+) -> FlockLoggingMCPCallback:
+    """Creates a fallback for handling incoming logging requests."""
+    logger_to_use = logger if logger else default_logging_callback_logger
+
+    async def _method(
+        params: FlockLoggingMessageNotificationParams,
+    ) -> None:
+        return await default_logging_callback(
+            params=params,
+            logger=logger_to_use,
+            server_name=associated_client.config.name,
+        )
+
+
+def default_flock_mcp_sampling_callback_factory(
+    associated_client: Any,
+    logger: FlockLogger | None = None,
+) -> FlockSamplingMCPCallback:
+    """Creates a fallback for handling incoming sampling requests."""
+    logger_to_use = logger if logger else default_sampling_callback_logger
+
+    async def _method(
+        ctx: RequestContext,
+        params: CreateMessageRequestParams,
+    ):
+        logger_to_use.info(
+            f"SAMPLING_REQUEST: server '{associated_client.config.name}' sent a sampling request: {params}"
+        )
+        await default_sampling_callback(
+            ctx=ctx, params=params, logger=logger_to_use
+        )
+
+    return _method
+
+
+def default_flock_mcp_message_handler_callback_factory(
+    associated_client: Any,
+    logger: FlockLogger | None = None,
+) -> FlockMessageHandlerMCPCallback:
+    """Creates a fallback for handling incoming messages.
+
+    Note:
+        Incoming Messages differ from incoming requests.
+        Requests can do things like list roots, create messages through sampling etc.
+        While Incoming Messages mainly consist of miscellanious information
+        sent by the server.
+    """
+    logger_to_use = logger if logger else default_message_handler_logger
+
+    async def _method(
+        n: ServerNotification,
+    ) -> None:
+        await default_message_handler(
+            req=n,
+            logger_to_use=logger_to_use,
+            associated_client=associated_client,
+        )
+
+    return _method
+
+
+def default_flock_mcp_list_roots_callback_factory(
+    associated_client: Any,
+    logger: FlockLogger | None = None,
+) -> FlockListRootsMCPCallback:
+    """Creates a fallback for a list roots callback for a client."""
+    logger_to_use = logger or default_list_roots_callback_logger
+
+    async def _method(*args, **kwargs):
+        return await default_list_roots_callback(
+            associated_client=associated_client,
+            logger=logger_to_use,
+        )
+
+    return _method