flock-core 0.4.512__py3-none-any.whl → 0.4.513__py3-none-any.whl

flock/core/mcp/mcp_config.py

@@ -1,7 +1,9 @@
 """Base Config for MCP Clients."""
 
-from typing import Literal, TypeVar
+import importlib
+from typing import Any, Literal, TypeVar
 
+import httpx
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from flock.core.mcp.types.types import (
@@ -11,6 +13,15 @@ from flock.core.mcp.types.types import (
     FlockSamplingMCPCallback,
     MCPRoot,
     ServerParameters,
+    SseServerParameters,
+    StdioServerParameters,
+    StreamableHttpServerParameters,
+    WebsocketServerParameters,
+)
+from flock.core.serialization.serializable import Serializable
+from flock.core.serialization.serialization_utils import (
+    deserialize_item,
+    serialize_item,
 )
 
 LoggingLevel = Literal[
@@ -32,7 +43,7 @@ D = TypeVar("D", bound="FlockMCPCachingConfigurationBase")
 E = TypeVar("E", bound="FlockMCPFeatureConfigurationBase")
 
 
-class FlockMCPCachingConfigurationBase(BaseModel):
+class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -79,6 +90,18 @@ class FlockMCPCachingConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize the config object."""
+        return self.model_dump(
+            exclude_none=True,
+            mode="json",
+        )
+
+    @classmethod
+    def from_dict(cls: type[D], data: dict[str, Any]) -> D:
+        """Deserialize from a dict."""
+        return cls(**{k: v for k, v in data.items()})
+
     @classmethod
     def with_fields(cls: type[D], **field_definitions) -> type[D]:
         """Create a new config class with additional fields."""
@@ -87,7 +110,7 @@ class FlockMCPCachingConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPCallbackConfigurationBase(BaseModel):
+class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -111,6 +134,52 @@ class FlockMCPCallbackConfigurationBase(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
+    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
+
+    @classmethod
+    def from_dict(cls: type[A], data: dict[str, Any]) -> A:
+        """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
+
     @classmethod
     def with_fields(cls: type[A], **field_definitions) -> type[A]:
         """Create a new config class with additional fields."""
@@ -119,7 +188,7 @@ class FlockMCPCallbackConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPConnectionConfigurationBase(BaseModel):
+class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -131,9 +200,9 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
         ..., description="Connection parameters for the server."
     )
 
-    transport_type: Literal["stdio", "websockets", "sse", "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."
@@ -150,6 +219,77 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
+    def to_dict(self, path_type: str = "relative") -> dict[str, Any]:
+        """Serialize object to a dict."""
+        exclude = ["connection_parameters"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json",
+        )
+
+        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:
+        """Deserialize from dict."""
+        connection_params = data.get("connection_parameters")
+        connection_params_obj = None
+        auth_obj: httpx.Auth | None = None
+        if connection_params:
+            kind = connection_params.get("transport_type", None)
+            auth_spec = connection_params.get("auth", None)
+            if auth_spec:
+                # find the concrete implementation and
+                # instantiate 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()})
+
+                if auth_obj:
+                    connection_params["auth"] = auth_obj
+                else:
+                    # just to be sure
+                    connection_params.pop("auth", None)
+                match kind:
+                    case "stdio":
+                        connection_params_obj = StdioServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "websockets":
+                        connection_params_obj = WebsocketServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "streamable_http":
+                        connection_params_obj = StreamableHttpServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "sse":
+                        connection_params_obj = SseServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case _:
+                        # handle custom server params
+                        connection_params_obj = ServerParameters(
+                            **{k: v for k, v in 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.")
+
     @classmethod
     def with_fields(cls: type[B], **field_definitions) -> type[B]:
         """Create a new config class with additional fields."""
@@ -158,7 +298,7 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPFeatureConfigurationBase(BaseModel):
+class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
@@ -186,6 +326,18 @@ class FlockMCPFeatureConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize the object."""
+        return self.model_dump(
+            mode="json",
+            exclude_none=True,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize from a dict."""
+        return cls(**{k: v for k, v in data.items()})
+
     @classmethod
     def with_fields(cls: type[E], **field_definitions) -> type[E]:
         """Create a new config class with additional fields."""
@@ -194,7 +346,7 @@ class FlockMCPFeatureConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPConfigurationBase(BaseModel):
+class FlockMCPConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for MCP Clients.
 
     Each Client should implement their own config
@@ -229,6 +381,90 @@ class FlockMCPConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative") -> dict[str, Any]:
+        """Serialize the object to a dict."""
+        # each built-in type should serialize, deserialize it self.
+        exclude = [
+            "connection_config",
+            "caching_config",
+            "callback_config",
+            "feature_config",
+        ]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json",
+        )
+
+        # add the core properties
+        data["connection_config"] = self.connection_config.to_dict(path_type)
+        data["caching_config"] = self.caching_config.to_dict(path_type)
+        data["callback_config"] = self.callback_config.to_dict(path_type)
+        data["feature_config"] = self.feature_config.to_dict(path_type)
+
+        return data
+
+    @classmethod
+    def from_dict(cls: type[C], data: dict[str, Any]) -> C:
+        """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"]
+        }
+
+        if connection_config:
+            # Forcing a square into a round hole
+            try:
+                config_field = cls.model_fields["connection_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPConnectionConfigurationBase
+            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')}")
+
+        if caching_config:
+            try:
+                config_field = cls.model_fields["caching_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPCachingConfigurationBase
+            instance_data["caching_config"] = config_cls.from_dict(caching_config)
+        else:
+            instance_data["caching_config"] = FlockMCPCachingConfigurationBase()
+
+        if feature_config:
+            try:
+                config_field = cls.model_fields["feature_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPFeatureConfigurationBase
+            instance_data["feature_config"] = config_cls.from_dict(feature_config)
+        else:
+            instance_data["feature_config"] = FlockMCPFeatureConfigurationBase()
+
+        if callback_config:
+            try:
+                config_field = cls.model_fields["callback_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPCallbackConfigurationBase
+            instance_data["callback_config"] = config_cls.from_dict(callback_config)
+        else:
+            instance_data["callback_config"] = FlockMCPCallbackConfigurationBase()
+
+        return cls(**{k: v for k, v in instance_data.items()})
+
     @classmethod
     def with_fields(cls: type[C], **field_definitions) -> type[C]:
         """Create a new config class with additional fields."""

flock/core/mcp/types/callbacks.py

@@ -8,6 +8,8 @@ from mcp.types import (
     CreateMessageRequestParams,
     ErrorData,
     ListRootsResult,
+    LoggingMessageNotificationParams,
+    ServerNotification,
     ServerRequest,
 )
 
@@ -19,10 +21,6 @@ from flock.core.mcp.types.handlers import (
     handle_incoming_server_notification,
     handle_logging_message,
 )
-from flock.core.mcp.types.types import (
-    FlockLoggingMessageNotificationParams,
-    ServerNotification,
-)
 
 
 async def default_sampling_callback(
@@ -76,7 +74,7 @@ async def default_list_roots_callback(
 
 
 async def default_logging_callback(
-    params: FlockLoggingMessageNotificationParams,
+    params: LoggingMessageNotificationParams,
     logger: FlockLogger,
     server_name: str,
 ) -> None:

flock/core/mcp/types/factories.py

@@ -1,6 +1,6 @@
 """Factories for default MCP Callbacks."""
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from mcp.shared.context import RequestContext
 from mcp.types import (
@@ -8,6 +8,12 @@ from mcp.types import (
 )
 
 from flock.core.logging.logging import FlockLogger, get_logger
+from flock.core.mcp.types.callbacks import (
+    default_list_roots_callback,
+    default_logging_callback,
+    default_message_handler,
+    default_sampling_callback,
+)
 from flock.core.mcp.types.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
@@ -17,18 +23,10 @@ from flock.core.mcp.types.types import (
     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")
+default_logging_callback_logger = get_logger("mcp.callback.logging")
+default_sampling_callback_logger = get_logger("mcp.callback.sampling")
+default_list_roots_callback_logger = get_logger("mcp.callback.roots")
+default_message_handler_logger = get_logger("mcp.callback.message")
 
 
 def default_flock_mcp_logging_callback_factory(
@@ -88,7 +86,7 @@ def default_flock_mcp_message_handler_callback_factory(
     ) -> None:
         await default_message_handler(
             req=n,
-            logger_to_use=logger_to_use,
+            logger=logger_to_use,
             associated_client=associated_client,
         )
 

flock/core/mcp/types/handlers.py

@@ -8,26 +8,23 @@ from mcp.shared.context import RequestContext
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     INTERNAL_ERROR,
+    CancelledNotification,
     ClientResult,
     ErrorData,
     ListRootsRequest,
-    ServerNotification as _MCPServerNotification,
-    ServerRequest,
-)
-
-from flock.core.logging.logging import FlockLogger
-from flock.core.mcp.mcp_client import Any
-from flock.core.mcp.types.types import (
-    CancelledNotification,
-    FlockLoggingMessageNotificationParams,
     LoggingMessageNotification,
+    LoggingMessageNotificationParams,
     ProgressNotification,
     ResourceListChangedNotification,
     ResourceUpdatedNotification,
     ServerNotification,
+    ServerRequest,
     ToolListChangedNotification,
 )
 
+from flock.core.logging.logging import FlockLogger
+from flock.core.mcp.mcp_client import Any
+
 
 async def handle_incoming_exception(
     e: Exception,
@@ -129,11 +126,11 @@ async def handle_tool_list_changed_notification(
     await associated_client.invalidate_tool_cache()
 
 
-_SERVER_NOTIFICATION_MAP: dict[type[_MCPServerNotification], Callable] = {
+_SERVER_NOTIFICATION_MAP: dict[type[ServerNotification], Callable] = {
     ResourceListChangedNotification: handle_resource_list_changed_notification,
     ResourceUpdatedNotification: handle_resource_update_notification,
     LoggingMessageNotification: lambda n, log, client: handle_logging_message(
-        params=n,
+        params=n.params,
         logger=log,
         server_name=client.config.name,
     ),
@@ -154,7 +151,7 @@ async def handle_incoming_server_notification(
 
 
 async def handle_logging_message(
-    params: FlockLoggingMessageNotificationParams,
+    params: LoggingMessageNotificationParams,
     logger: FlockLogger,
     server_name: str,
 ) -> None: