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

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:

flock/core/mcp/types/types.py

@@ -1,9 +1,13 @@
 """Types for Flock's MCP functionality."""
 
+import importlib
+import inspect
+import os
 from collections.abc import Awaitable, Callable
 from contextlib import AbstractAsyncContextManager
-from typing import Any
+from typing import Any, Literal
 
+import httpx
 from anyio.streams.memory import (
     MemoryObjectReceiveStream,
     MemoryObjectSendStream,
@@ -36,6 +40,7 @@ from mcp.types import (
 from pydantic import AnyUrl, BaseModel, ConfigDict, Field
 
 from flock.core.mcp.util.helpers import get_default_env
+from flock.core.serialization.serializable import Serializable
 
 
 class ServerNotification(_MCPServerNotification):
@@ -80,17 +85,40 @@ class MCPRoot(_MCPRoot):
     """Wrapper for mcp.types.Root."""
 
 
-class ServerParameters(BaseModel):
+class ServerParameters(BaseModel, Serializable):
     """Base Type for server parameters."""
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
     )
 
+    transport_type: Literal["stdio", "websockets", "sse", "streamable_http"] = Field(
+        ...,
+        description="which type of transport these connection params are used for."
+    )
+
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize."""
+        return self.model_dump(
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json"
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize."""
+        return cls(**data)
+
 
 class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
     """Base Type for Stdio Server parameters."""
 
+    transport_type: Literal["stdio"] = Field(
+        default="stdio",
+        description="Use stdio params."
+    )
+
     env: dict[str, str] | None = Field(
         default_factory=get_default_env,
         description="Environment for the MCP Server.",
@@ -100,12 +128,121 @@ class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
 class WebsocketServerParameters(ServerParameters):
     """Base Type for Websocket Server params."""
 
+    transport_type: Literal["websockets"] = Field(
+        default="websockets",
+        description="Use websocket params."
+    )
+
     url: str | AnyUrl = Field(..., description="Url the server listens at.")
 
+class StreamableHttpServerParameters(ServerParameters):
+    """Base Type for StreamableHttp params."""
+
+    transport_type: Literal["streamable_http"] = Field(
+        default="streamable_http",
+        description="Use streamable http params."
+    )
+
+    url: str | AnyUrl = Field(
+        ...,
+        description="The url the server listens at."
+    )
+
+    headers: dict[str, Any] | None = Field(
+        default=None,
+        description="Additional headers to pass to the client."
+    )
+
+    timeout: float | int = Field(
+        default=5,
+        description="Http Timeout",
+    )
+
+    sse_read_timeout: float | int = Field(
+        default=60*5,
+        description="How long the client will wait before disconnecting from the server."
+    )
+
+    terminate_on_close: bool = Field(
+        default=True,
+        description="Terminate connection on close"
+    )
+
+    auth: httpx.Auth | None = Field(
+        default=None,
+        description="Httpx Auth Scheme"
+    )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize the object from a dict."""
+        # find and import the concrete implementation for
+        # the auth object
+        auth_obj: httpx.Auth | None = None
+        auth_impl = data.pop("auth", None)
+        if auth_impl:
+            # find the concrete implementation
+            impl = auth_impl.pop("implementation", None)
+            params = auth_impl.pop("params", None)
+            if impl:
+                mod = importlib.import_module(impl["module_path"])
+                real_cls = getattr(mod, impl["classname"])
+                if params:
+                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                else:
+                    # assume that the implementation handles it.
+                    auth_obj = real_cls()
+            else:
+                raise ValueError("No concrete implementation for auth provided.")
+
+        data["auth"] = auth_obj
+        return cls(**{k: v for k, v in data.items()})
+
+    def to_dict(self, path_type = "relative"):
+        """Serialize the object."""
+        exclude = ["auth"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+        )
+
+        # inject implentation info for auth
+        if self.auth is not None:
+            file_path = inspect.getsourcefile(type(self.auth))
+            if path_type == "relative":
+                file_path = os.path.relpath(file_path)
+            try:
+                # params should be primitive types, keeping with the
+                # declarative approach of flock.
+                params = {
+                    k: getattr(self.auth, k)
+                    for k in getattr(self.auth, "__dict__", {})
+                    if not k.startswith("_")
+                }
+            except Exception:
+                params = None
+
+        data["auth"] = {
+            "implementation": {
+                "class_name": type(self.auth).__name__,
+                "module_path": type(self.auth).__module__,
+                "file_path": file_path,
+            },
+            "params": params,
+        }
+
+        return data
 
 class SseServerParameters(ServerParameters):
     """Base Type for SSE Server params."""
 
+    transport_type: Literal["sse"] = Field(
+        default="sse",
+        description="Use sse server params."
+    )
+
     url: str | AnyUrl = Field(..., description="The url the server listens at.")
 
     headers: dict[str, Any] | None = Field(
@@ -119,6 +256,72 @@ class SseServerParameters(ServerParameters):
         description="How long the client will wait before disconnecting from the server.",
     )
 
+    auth: httpx.Auth | None = Field(
+        default=None,
+        description="Httpx Auth Scheme."
+    )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize the object from a dict."""
+        # find and import the concrete implementation for
+        # the auth object.
+        auth_obj: httpx.Auth | None = None
+        auth_impl = data.pop("auth", None) # get the specs for the auth class
+        if auth_impl:
+            # find the concrete implementation
+            impl = auth_impl.pop("implementation", None)
+            params = auth_impl.pop("params", None)
+            if impl:
+                mod = importlib.import_module(impl["module_path"])
+                real_cls = getattr(mod, impl["class_name"])
+                if params:
+                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                else:
+                    # assume that implementation handles it
+                    auth_obj = real_cls()
+            else:
+                raise ValueError("No concrete implementation for auth provided.")
+
+        data["auth"] = auth_obj
+        return cls(**{k: v for k, v in data.items()})
+
+    def to_dict(self, path_type = "relative"):
+        """Serialize the object."""
+        exclude = ["auth"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+        )
+
+        # inject implentation info for auth
+        if self.auth is not None:
+            file_path = inspect.getsourcefile(type(self.auth))
+            if path_type == "relative":
+                file_path = os.path.relpath(file_path)
+            try:
+                # params should be primitive types, keeping with the
+                # declarative approach of flock.
+                params = {
+                    k: getattr(self.auth, k)
+                    for k in getattr(self.auth, "__dict__", {})
+                    if not k.startswith("_")
+                }
+            except Exception:
+                params = None
+
+        data["auth"] = {
+            "implementation": {
+                "class_name": type(self.auth).__name__,
+                "module_path": type(self.auth).__module__,
+                "file_path": file_path,
+            },
+            "params": params,
+        }
+
+        return data
 
 MCPCLientInitFunction = Callable[
     ...,

flock/core/mixin/dspy_integration.py

@@ -9,7 +9,7 @@ from typing import Any, Literal
 from dspy import Tool
 
 from flock.core.logging.logging import get_logger
-from flock.core.util.spliter import split_top_level
+from flock.core.util.splitter import split_top_level
 
 # Import split_top_level (assuming it's moved or copied appropriately)
 # Option 1: If moved to a shared util

flock/core/util/input_resolver.py

@@ -1,7 +1,7 @@
 """Utility functions for resolving input keys to their corresponding values."""
 
 from flock.core.context.context import FlockContext
-from flock.core.util.spliter import split_top_level
+from flock.core.util.splitter import split_top_level
 
 
 def get_callable_members(obj):

flock/mcp/servers/sse/flock_sse_server.py

@@ -4,12 +4,13 @@ import copy
 from contextlib import AbstractAsyncContextManager
 from typing import Any, Literal
 
+import httpx
 from anyio.streams.memory import (
     MemoryObjectReceiveStream,
     MemoryObjectSendStream,
 )
 from mcp.client.sse import sse_client
-from mcp.types import JSONRPCMessage
+from mcp.shared.message import SessionMessage
 from opentelemetry import trace
 from pydantic import Field
 
@@ -63,45 +64,51 @@ class FlockSSEClient(FlockMCPClientBase):
         additional_params: dict[str, Any] | None = None,
     ) -> AbstractAsyncContextManager[
         tuple[
-            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-            MemoryObjectSendStream[JSONRPCMessage],
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
         ]
     ]:
         """Return an async context manager whose __aenter__ method yields (read_stream, send_stream)."""
         # avoid modifying the config of the client as a side-effect.
         param_copy = copy.deepcopy(params)
 
-        if self.additional_params:
+        if additional_params:
             override_headers = bool(
-                self.additional_params.get("override_headers", False)
+                additional_params.get("override_headers", False)
             )
-            if "headers" in self.additional_params:
+            if "headers" in additional_params:
                 if override_headers:
-                    param_copy.headers = self.additional_params.get(
+                    param_copy.headers = additional_params.get(
                         "headers", params.headers
                     )
                 else:
                     param_copy.headers.update(
-                        self.additional_params.get("headers", {})
+                        additional_params.get("headers", {})
                     )
-            if "read_timeout_seconds" in self.additional_params:
-                param_copy.timeout = self.additional_params.get(
+            if "read_timeout_seconds" in additional_params:
+                param_copy.timeout =  additional_params.get(
                     "read_timeout_seconds", params.timeout
                 )
 
-            if "sse_read_timeout" in self.additional_params:
-                param_copy.sse_read_timeout = self.additional_params.get(
+            if "sse_read_timeout" in additional_params:
+                param_copy.sse_read_timeout = additional_params.get(
                     "sse_read_timeout",
                     params.sse_read_timeout,
                 )
-            if "url" in self.additional_params:
-                param_copy.url = self.additional_params.get(
+            if "url" in additional_params:
+                param_copy.url = additional_params.get(
                     "url",
                     params.url,
                 )
 
+            if "auth" in additional_params and isinstance(
+                additional_params.get("auth"), httpx.Auth
+            ):
+                param_copy.auth = additional_params.get("auth", param_copy.auth)
+
         return sse_client(
             url=param_copy.url,
+            auth=param_copy.auth,
             headers=param_copy.headers,
             timeout=float(param_copy.timeout),
             sse_read_timeout=float(param_copy.sse_read_timeout),

flock/mcp/servers/streamable_http/flock_streamable_http_server.py

@@ -0,0 +1,169 @@
+"""This module provides the Flock Streamable-Http functionality."""
+
+import copy
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager
+from datetime import timedelta
+from typing import Any, Literal
+
+import httpx
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.message import SessionMessage
+from opentelemetry import trace
+from pydantic import Field
+
+from flock.core.logging.logging import get_logger
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.mcp_client import FlockMCPClientBase
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManagerBase
+from flock.core.mcp.mcp_config import (
+    FlockMCPConfigurationBase,
+    FlockMCPConnectionConfigurationBase,
+)
+from flock.core.mcp.types.types import (
+    StreamableHttpServerParameters,
+)
+
+logger = get_logger("mcp.streamable_http.server")
+tracer = trace.get_tracer(__name__)
+
+GetSessionIdCallback = Callable[[], str | None]
+
+
+class FlockStreamableHttpConnectionConfig(FlockMCPConnectionConfigurationBase):
+    """Concrete ConnectionConfig for a StreamableHttpClient."""
+
+    # Only thing we need to override here is the concrete transport_type
+    # and connection parameter fields.
+    transport_type: Literal["streamable_http"] = Field(
+        default="streamable_http",
+        description="Use the streamable_http Transport type.",
+    )
+
+    connection_parameters: StreamableHttpServerParameters = Field(
+        ..., description="Streamable HTTP Server Connection Parameters."
+    )
+
+
+class FlockStreamableHttpConfig(FlockMCPConfigurationBase):
+    """Configuration for Streamable HTTP Clients."""
+
+    # The only thing we need to override here is the
+    # concrete connection config.
+    # The rest is generic enough to handle everything else.
+    connection_config: FlockStreamableHttpConnectionConfig = Field(
+        ..., description="Concrete StreamableHttp Connection Configuration."
+    )
+
+
+class FlockStreamableHttpClient(FlockMCPClientBase):
+    """Client for StreamableHttpServers."""
+
+    config: FlockStreamableHttpConfig = Field(
+        ..., description="Client configuration."
+    )
+
+    async def create_transport(
+        self,
+        params: StreamableHttpServerParameters,
+        additional_params: dict[str, Any] | None = None,
+    ) -> AbstractAsyncContextManager[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        GetSessionIdCallback,
+    ],
+    None,
+]:
+        """Return an async context manager whose __aenter__ method yields (read_stream, send_stream)."""
+        param_copy = copy.deepcopy(params)
+
+        if additional_params:
+            override_headers = bool(
+                additional_params.get("override_headers", False)
+            )
+
+            if "headers" in additional_params:
+                if override_headers:
+                    param_copy.headers = additional_params.get(
+                        "headers", params.headers
+                    )
+                else:
+                    param_copy.headers.update(
+                        additional_params.get("headers", {})
+                    )
+            if "auth" in additional_params and isinstance(
+                additional_params.get("auth"), httpx.Auth
+            ):
+                param_copy.auth = additional_params.get("auth", param_copy.auth)
+
+            if "read_timeout_seconds" in additional_params:
+                param_copy.timeout = additional_params.get(
+                    "read_timeout_seconds", params.timeout
+                )
+
+            if "sse_read_timeout" in additional_params:
+                param_copy.sse_read_timeout = additional_params.get(
+                    "sse_read_timeout",
+                    params.sse_read_timeout,
+                )
+            if "url" in additional_params:
+                param_copy.url = additional_params.get(
+                    "url",
+                    params.url,
+                )
+
+            if "terminate_on_close" in additional_params:
+                param_copy.terminate_on_close = bool(
+                    additional_params.get("terminate_on_close", True)
+                )
+
+        timeout_http = timedelta(seconds=param_copy.timeout)
+        sse_timeout = timedelta(seconds=param_copy.sse_read_timeout)
+
+        return streamablehttp_client(
+            url=param_copy.url,
+            headers=param_copy.headers,
+            timeout=timeout_http,
+            sse_read_timeout=sse_timeout,
+            terminate_on_close=param_copy.terminate_on_close,
+            auth=param_copy.auth,
+        )
+
+
+class FlockStreamableHttpClientManager(FlockMCPClientManagerBase):
+    """Manager for handling StreamableHttpClients."""
+
+    client_config: FlockStreamableHttpConfig = Field(
+        ..., description="Configuration for clients."
+    )
+
+    async def make_client(
+        self, additional_params: dict[str, Any] | None = None
+    ) -> FlockStreamableHttpClient:
+        """Create a new client instance."""
+        new_client = FlockStreamableHttpClient(
+            config=self.client_config,
+            additional_params=additional_params,
+        )
+        return new_client
+
+
+class FlockStreamableHttpServer(FlockMCPServerBase):
+    """Class which represents a MCP Server using the streamable Http Transport type."""
+
+    config: FlockStreamableHttpConfig = Field(
+        ..., description="Config for the server."
+    )
+
+    async def initialize(self) -> FlockStreamableHttpClientManager:
+        """Called when initializing the server."""
+        client_manager = FlockStreamableHttpClientManager(
+            client_config=self.config
+        )
+
+        return client_manager

flock/mcp/servers/websockets/flock_websocket_server.py

@@ -9,7 +9,7 @@ from anyio.streams.memory import (
     MemoryObjectSendStream,
 )
 from mcp.client.websocket import websocket_client
-from mcp.types import JSONRPCMessage
+from mcp.shared.message import SessionMessage
 from opentelemetry import trace
 from pydantic import Field
 
@@ -69,8 +69,8 @@ class FlockWSClient(FlockMCPClientBase):
         additional_params: dict[str, Any] | None = None,
     ) -> AbstractAsyncContextManager[
         tuple[
-            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-            MemoryObjectSendStream[JSONRPCMessage],
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
         ]
     ]:
         """Return an async context manager whose __aenter__ method yields a read_stream and a send_stream."""