flock-core 0.4.543__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/{core/mcp/mcp_config.py → mcp/config.py}

@@ -1,12 +1,12 @@
 """Base Config for MCP Clients."""
 
 import importlib
-from typing import Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
-import httpx
 from pydantic import BaseModel, ConfigDict, Field, create_model
+from typing_extensions import Self
 
-from flock.core.mcp.types.types import (
+from flock.mcp.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
     FlockMessageHandlerMCPCallback,
@@ -18,11 +18,11 @@ from flock.core.mcp.types.types import (
     StreamableHttpServerParameters,
     WebsocketServerParameters,
 )
-from flock.core.serialization.serializable import Serializable
-from flock.core.serialization.serialization_utils import (
-    deserialize_item,
-    serialize_item,
-)
+
+
+if TYPE_CHECKING:
+    import httpx
+
 
 LoggingLevel = Literal[
     "debug",
@@ -36,14 +36,14 @@ LoggingLevel = Literal[
 ]
 
 
-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")
+A = TypeVar("A", bound="FlockMCPCallbackConfiguration")
+B = TypeVar("B", bound="FlockMCPConnectionConfiguration")
+C = TypeVar("C", bound="FlockMCPConfiguration")
+D = TypeVar("D", bound="FlockMCPCachingConfiguration")
+E = TypeVar("E", bound="FlockMCPFeatureConfiguration")
 
 
-class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
+class FlockMCPCachingConfiguration(BaseModel):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -98,19 +98,17 @@ class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
         )
 
     @classmethod
-    def from_dict(cls: type[D], data: dict[str, Any]) -> D:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from a dict."""
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
     @classmethod
-    def with_fields(cls: type[D], **field_definitions) -> type[D]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
+class FlockMCPCallbackConfiguration(BaseModel):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -136,59 +134,21 @@ class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
 
     def to_dict(self, path_type: str = "relative"):
         """Serialize the object."""
-        # we need to register callables.
-        data: dict[str, Any] = {}
-        if self.sampling_callback:
-            sampling_callback_data = serialize_item(self.sampling_callback)
-            data["sampling_callback"] = sampling_callback_data
-
-        if self.list_roots_callback:
-            list_roots_callback_data = serialize_item(self.list_roots_callback)
-            data["list_roots_callback"] = list_roots_callback_data
-
-        if self.logging_callback:
-            logging_callback_data = serialize_item(self.logging_callback)
-            data["logging_callback"] = logging_callback_data
-
-        if self.message_handler:
-            message_handler_data = serialize_item(self.message_handler)
-            data["message_handler"] = message_handler_data
-
-        return data
+        # Callbacks are set programmatically and cannot be serialized
+        return {}
 
     @classmethod
-    def from_dict(cls: type[A], data: dict[str, Any]) -> A:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from a dict."""
-        instance = cls()
-        if data:
-            if "sampling_callback" in data:
-                instance.sampling_callback = deserialize_item(
-                    data["sampling_callback"]
-                )
-            if "list_roots_callback" in data:
-                instance.list_roots_callback = deserialize_item(
-                    data["list_roots_callback"]
-                )
-            if "logging_callback" in data:
-                instance.logging_callback = deserialize_item(
-                    data["logging_callback"]
-                )
-            if "message_handler" in data:
-                instance.message_handler = deserialize_item(
-                    data["message_handler"]
-                )
-
-        return instance
+        return cls()
 
     @classmethod
-    def with_fields(cls: type[A], **field_definitions) -> type[A]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
+class FlockMCPConnectionConfiguration(BaseModel):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -200,17 +160,15 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
         ..., description="Connection parameters for the server."
     )
 
-    transport_type: Literal[
-        "stdio", "websockets", "sse", "streamable_http", "custom"
-    ] = Field(..., description="Type of transport to use.")
+    transport_type: Literal["stdio", "websockets", "sse", "streamable_http", "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."
-    )
+    read_timeout_seconds: float | int = Field(default=60 * 5, description="Read Timeout.")
 
     server_logging_level: LoggingLevel = Field(
         default="error",
@@ -230,14 +188,12 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
             mode="json",
         )
 
-        data["connection_parameters"] = self.connection_parameters.to_dict(
-            path_type=path_type
-        )
+        data["connection_parameters"] = self.connection_parameters.to_dict(path_type=path_type)
 
         return data
 
     @classmethod
-    def from_dict(cls: type[B], data: dict[str, Any]) -> B:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from dict."""
         connection_params = data.get("connection_parameters")
         connection_params_obj = None
@@ -248,13 +204,13 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
             if auth_spec:
                 # find the concrete implementation and
                 # instantiate it.
-                    # find the concrete implementation for auth and instatiate it.
+                # find the concrete implementation for auth and instatiate it.
                 impl = auth_spec.get("implementation", None)
                 params = auth_spec.get("params", None)
                 if impl and params:
-                        mod = importlib.import_module(impl["module_path"])
-                        real_cls = getattr(mod, impl["class_name"])
-                        auth_obj = real_cls(**{k: v for k, v in params.items()})
+                    mod = importlib.import_module(impl["module_path"])
+                    real_cls = getattr(mod, impl["class_name"])
+                    auth_obj = real_cls(**dict(params.items()))
 
                 if auth_obj:
                     connection_params["auth"] = auth_obj
@@ -264,60 +220,64 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
                 match kind:
                     case "stdio":
                         connection_params_obj = StdioServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "websockets":
                         connection_params_obj = WebsocketServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "streamable_http":
                         connection_params_obj = StreamableHttpServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "sse":
                         connection_params_obj = SseServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case _:
                         # handle custom server params
-                        connection_params_obj = ServerParameters(
-                            **{k: v for k, v in connection_params.items()}
-                        )
+                        connection_params_obj = ServerParameters(**dict(connection_params.items()))
 
         if connection_params_obj:
             data["connection_parameters"] = connection_params_obj
-            return cls(**{k: v for k, v in data.items()})
-        else:
-            raise ValueError("No connection parameters provided.")
+            return cls(**dict(data.items()))
+        raise ValueError("No connection parameters provided.")
 
     @classmethod
-    def with_fields(cls: type[B], **field_definitions) -> type[B]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
+class FlockMCPFeatureConfiguration(BaseModel):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
-        default=False,
+        default=True,
         description="Whether or not the Roots feature is enabled for this client.",
     )
 
     sampling_enabled: bool = Field(
-        default=False,
+        default=True,
         description="Whether or not the Sampling feature is enabled for this client.",
     )
 
     tools_enabled: bool = Field(
-        default=False,
+        default=True,
         description="Whether or not the Tools feature is enabled for this client.",
     )
 
+    tool_whitelist: list[str] | None = Field(
+        default=None,
+        description="Whitelist of tool names that are enabled for this MCP server. "
+        "If provided, only tools with names in this list will be available "
+        "from this server."
+        "Note: Agent-level tool filtering is generally preferred over "
+        "server-level filtering for better granular control.",
+    )
+
     prompts_enabled: bool = Field(
-        default=False,
+        default=True,
         description="Whether or not the Prompts feature is enabled for this client.",
     )
 
@@ -336,43 +296,39 @@ class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
         """Deserialize from a dict."""
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
     @classmethod
-    def with_fields(cls: type[E], **field_definitions) -> type[E]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPConfigurationBase(BaseModel, Serializable):
+class FlockMCPConfiguration(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."
-    )
+    name: str = Field(..., description="Name of the server the client connects to.")
 
-    connection_config: FlockMCPConnectionConfigurationBase = Field(
+    connection_config: FlockMCPConnectionConfiguration = Field(
         ..., description="MCP Connection Configuration for a client."
     )
 
-    caching_config: FlockMCPCachingConfigurationBase = Field(
-        default_factory=FlockMCPCachingConfigurationBase,
+    caching_config: FlockMCPCachingConfiguration = Field(
+        default_factory=FlockMCPCachingConfiguration,
         description="Configuration for the internal caches of the client.",
     )
 
-    callback_config: FlockMCPCallbackConfigurationBase = Field(
-        default_factory=FlockMCPCallbackConfigurationBase,
+    callback_config: FlockMCPCallbackConfiguration = Field(
+        default_factory=FlockMCPCallbackConfiguration,
         description="Callback configuration for the client.",
     )
 
-    feature_config: FlockMCPFeatureConfigurationBase = Field(
-        default_factory=FlockMCPFeatureConfigurationBase,
+    feature_config: FlockMCPFeatureConfiguration = Field(
+        default_factory=FlockMCPFeatureConfiguration,
         description="Feature configuration for the client.",
     )
 
@@ -407,16 +363,14 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
         return data
 
     @classmethod
-    def from_dict(cls: type[C], data: dict[str, Any]) -> C:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize the class."""
         connection_config = data.pop("connection_config", None)
         caching_config = data.pop("caching_config", None)
         feature_config = data.pop("feature_config", None)
         callback_config = data.pop("callback_config", None)
 
-        instance_data: dict[str, Any] = {
-            "name": data["name"]
-        }
+        instance_data: dict[str, Any] = {"name": data["name"]}
 
         if connection_config:
             # Forcing a square into a round hole
@@ -425,10 +379,12 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPConnectionConfigurationBase
+                config_cls = FlockMCPConnectionConfiguration
             instance_data["connection_config"] = config_cls.from_dict(connection_config)
         else:
-            raise ValueError(f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}")
+            raise ValueError(
+                f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}"
+            )
 
         if caching_config:
             try:
@@ -436,10 +392,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCachingConfigurationBase
+                config_cls = FlockMCPCachingConfiguration
             instance_data["caching_config"] = config_cls.from_dict(caching_config)
         else:
-            instance_data["caching_config"] = FlockMCPCachingConfigurationBase()
+            instance_data["caching_config"] = FlockMCPCachingConfiguration()
 
         if feature_config:
             try:
@@ -447,10 +403,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPFeatureConfigurationBase
+                config_cls = FlockMCPFeatureConfiguration
             instance_data["feature_config"] = config_cls.from_dict(feature_config)
         else:
-            instance_data["feature_config"] = FlockMCPFeatureConfigurationBase()
+            instance_data["feature_config"] = FlockMCPFeatureConfiguration()
 
         if callback_config:
             try:
@@ -458,16 +414,14 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCallbackConfigurationBase
+                config_cls = FlockMCPCallbackConfiguration
             instance_data["callback_config"] = config_cls.from_dict(callback_config)
         else:
-            instance_data["callback_config"] = FlockMCPCallbackConfigurationBase()
+            instance_data["callback_config"] = FlockMCPCallbackConfiguration()
 
-        return cls(**{k: v for k, v in instance_data.items()})
+        return cls(**dict(instance_data.items()))
 
     @classmethod
-    def with_fields(cls: type[C], **field_definitions) -> type[C]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)