flock-core 0.5.0b28__py3-none-any.whl → 0.5.0b51__py3-none-any.whl

flock/mcp/servers/sse/flock_sse_server.py

@@ -14,15 +14,14 @@ 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 FlockMCPServer
-from flock.core.mcp.mcp_client import FlockMCPClient
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
-from flock.core.mcp.mcp_config import (
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
     FlockMCPConfiguration,
     FlockMCPConnectionConfiguration,
 )
-from flock.core.mcp.types.types import SseServerParameters
+from flock.mcp.types import SseServerParameters
+
 
 logger = get_logger("mcp.sse.server")
 tracer = trace.get_tracer(__name__)
@@ -33,9 +32,7 @@ class FlockSSEConnectionConfig(FlockMCPConnectionConfiguration):
 
     # Only thing we need to override here is the concrete transport_type
     # and connection_parameters fields.
-    transport_type: Literal["sse"] = Field(
-        default="sse", description="Use the sse transport type."
-    )
+    transport_type: Literal["sse"] = Field(default="sse", description="Use the sse transport type.")
 
     connection_parameters: SseServerParameters = Field(
         ..., description="SSE Server Connection Parameters."
@@ -56,7 +53,7 @@ class FlockSSEConfig(FlockMCPConfiguration):
 class FlockSSEClient(FlockMCPClient):
     """Client for SSE Servers."""
 
-    config: FlockSSEConfig = Field(..., description="Client configuration.")
+    config: FlockMCPConfiguration = Field(..., description="Client configuration.")
 
     async def create_transport(
         self,
@@ -73,22 +70,14 @@ class FlockSSEClient(FlockMCPClient):
         param_copy = copy.deepcopy(params)
 
         if additional_params:
-            override_headers = bool(
-                additional_params.get("override_headers", False)
-            )
+            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
-                    )
+                    param_copy.headers = additional_params.get("headers", params.headers)
                 else:
-                    param_copy.headers.update(
-                        additional_params.get("headers", {})
-                    )
+                    param_copy.headers.update(additional_params.get("headers", {}))
             if "read_timeout_seconds" in additional_params:
-                param_copy.timeout =  additional_params.get(
-                    "read_timeout_seconds", params.timeout
-                )
+                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(
@@ -113,34 +102,3 @@ class FlockSSEClient(FlockMCPClient):
             timeout=float(param_copy.timeout),
             sse_read_timeout=float(param_copy.sse_read_timeout),
         )
-
-
-class FlockSSEClientManager(FlockMCPClientManager):
-    """Manager for handling SSE Clients."""
-
-    client_config: FlockSSEConfig = Field(
-        ..., description="Configuration for clients."
-    )
-
-    async def make_client(
-        self, additional_params: dict[str, Any]
-    ) -> FlockSSEClient:
-        """Create a new client instance."""
-        new_client = FlockSSEClient(
-            config=self.client_config, additional_params=additional_params
-        )
-        return new_client
-
-
-class FlockSSEServer(FlockMCPServer):
-    """Class which represents a MCP Server using the SSE Transport type."""
-
-    config: FlockSSEConfig = Field(..., description="Config for the server.")
-
-    async def initialize(self) -> FlockSSEClientManager:
-        """Called when initializing the server."""
-        client_manager = FlockSSEClientManager(
-            client_config=self.config,
-        )
-
-        return client_manager

flock/mcp/servers/stdio/__init__.py

@@ -1 +1 @@
-"""Default Stdio Server Implementation for Flock."""
+"""Default Stdio Server Implementation for flock."""

flock/mcp/servers/stdio/flock_stdio_server.py

@@ -13,15 +13,14 @@ from mcp.types import JSONRPCMessage
 from opentelemetry import trace
 from pydantic import Field
 
-from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_server import FlockMCPServer
-from flock.core.mcp.mcp_client import FlockMCPClient
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
-from flock.core.mcp.mcp_config import (
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
     FlockMCPConfiguration,
     FlockMCPConnectionConfiguration,
 )
-from flock.core.mcp.types.types import StdioServerParameters
+from flock.mcp.types import StdioServerParameters
+
 
 logger = get_logger("mcp.stdio.server")
 tracer = trace.get_tracer(__name__)
@@ -56,7 +55,7 @@ class FlockStdioConfig(FlockMCPConfiguration):
 class FlockStdioClient(FlockMCPClient):
     """Client for Stdio Servers."""
 
-    config: FlockStdioConfig = Field(..., description="Client Configuration.")
+    config: FlockMCPConfiguration = Field(..., description="Client Configuration.")
 
     async def create_transport(
         self,
@@ -78,9 +77,7 @@ class FlockStdioClient(FlockMCPClient):
         if additional_params:
             # If it is present, then modify server parameters based on certain keys.
             if "command" in additional_params:
-                param_copy.command = additional_params.get(
-                    "command", params.command
-                )
+                param_copy.command = additional_params.get("command", params.command)
             if "args" in additional_params:
                 param_copy.args = additional_params.get("args", params.command)
             if "env" in additional_params:
@@ -90,9 +87,7 @@ class FlockStdioClient(FlockMCPClient):
                 param_copy.cwd = additional_params.get("cwd", params.env)
 
             if "encoding" in additional_params:
-                param_copy.encoding = additional_params.get(
-                    "encoding", params.encoding
-                )
+                param_copy.encoding = additional_params.get("encoding", params.encoding)
 
             if "encoding_error_handler" in additional_params:
                 param_copy.encoding_error_handler = additional_params.get(
@@ -101,38 +96,3 @@ class FlockStdioClient(FlockMCPClient):
 
         # stdio_client already is an AsyncContextManager
         return stdio_client(server=param_copy)
-
-
-# Not really needed but kept here as an example.
-class FlockStdioClientManager(FlockMCPClientManager):
-    """Manager for handling Stdio Clients."""
-
-    client_config: FlockStdioConfig = Field(
-        ..., description="Configuration for clients."
-    )
-
-    async def make_client(
-        self, additional_params: dict[str, Any] | None = None
-    ):
-        """Create a new client instance with any additional parameters."""
-        new_client = FlockStdioClient(
-            config=self.client_config,
-            additional_params=additional_params,
-        )
-        return new_client
-
-
-class FlockMCPStdioServer(FlockMCPServer):
-    """Class which represents a MCP Server using the Stdio Transport type.
-
-    This means (most likely) that the server is a locally
-    executed script.
-    """
-
-    config: FlockStdioConfig = Field(..., description="Config for the server.")
-
-    async def initialize(self) -> FlockStdioClientManager:
-        """Called when initializing the server."""
-        client_manager = FlockStdioClientManager(client_config=self.config)
-
-        return client_manager

flock/mcp/servers/streamable_http/flock_streamable_http_server.py

@@ -16,18 +16,17 @@ 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 FlockMCPServer
-from flock.core.mcp.mcp_client import FlockMCPClient
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
-from flock.core.mcp.mcp_config import (
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
     FlockMCPConfiguration,
     FlockMCPConnectionConfiguration,
 )
-from flock.core.mcp.types.types import (
+from flock.mcp.types import (
     StreamableHttpServerParameters,
 )
 
+
 logger = get_logger("mcp.streamable_http.server")
 tracer = trace.get_tracer(__name__)
 
@@ -63,48 +62,38 @@ class FlockStreamableHttpConfig(FlockMCPConfiguration):
 class FlockStreamableHttpClient(FlockMCPClient):
     """Client for StreamableHttpServers."""
 
-    config: FlockStreamableHttpConfig = Field(
-        ..., description="Client configuration."
-    )
+    config: FlockMCPConfiguration = 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,
-]:
+        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)
-            )
+            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
-                    )
+                    param_copy.headers = additional_params.get("headers", params.headers)
                 else:
-                    param_copy.headers.update(
-                        additional_params.get("headers", {})
-                    )
+                    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
-                )
+                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(
@@ -133,37 +122,3 @@ class FlockStreamableHttpClient(FlockMCPClient):
             terminate_on_close=param_copy.terminate_on_close,
             auth=param_copy.auth,
         )
-
-
-class FlockStreamableHttpClientManager(FlockMCPClientManager):
-    """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(FlockMCPServer):
-    """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

@@ -13,15 +13,14 @@ 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 FlockMCPServer
-from flock.core.mcp.mcp_client import FlockMCPClient
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
-from flock.core.mcp.mcp_config import (
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
     FlockMCPConfiguration,
     FlockMCPConnectionConfiguration,
 )
-from flock.core.mcp.types.types import WebsocketServerParameters
+from flock.mcp.types import WebsocketServerParameters
+
 
 logger = get_logger("mcp.ws.server")
 tracer = trace.get_tracer(__name__)
@@ -59,7 +58,7 @@ class FlockWSConfig(FlockMCPConfiguration):
 class FlockWSClient(FlockMCPClient):
     """Client for Websocket servers."""
 
-    config: FlockWSConfig = Field(..., description="Client Configuration")
+    config: FlockMCPConfiguration = Field(..., description="Client Configuration")
 
     # This one we HAVE to specify. This tells Flock
     # how to create the underlying connection.
@@ -84,36 +83,4 @@ class FlockWSClient(FlockMCPClient):
             # If present, then apply the changes in "url" to the create_transport logic.
             param_copy.url = additional_params.get("url", params.url)
 
-        return websocket_client(
-            url=param_copy.url
-        )  # return the async context manager
-
-
-# not really needed, but kept for type hints and as an example.
-class FlockWSClientManager(FlockMCPClientManager):
-    """Manager for handling websocket clients."""
-
-    client_config: FlockWSConfig = Field(
-        ..., description="Configuration for clients."
-    )
-
-    async def make_client(self, additional_params=None):
-        """Create a new client instance."""
-        new_client = FlockWSClient(
-            config=self.client_config,
-            additional_params=additional_params,
-        )
-        return new_client
-
-
-class FlockWSServer(FlockMCPServer):
-    """Class which represents an MCP Server using the websocket transport type."""
-
-    config: FlockWSConfig = Field(..., description="Config for the server.")
-
-    # Specify the concrete type for the server.
-    async def initialize(self) -> FlockWSClientManager:
-        """Called when initializing the server."""
-        client_manager = FlockWSClientManager(client_config=self.config)
-
-        return client_manager
+        return websocket_client(url=param_copy.url)  # return the async context manager

flock/{core/mcp/flock_mcp_tool.py → mcp/tool.py}

@@ -8,8 +8,10 @@ from mcp import Tool
 from mcp.types import CallToolResult, TextContent, ToolAnnotations
 from opentelemetry import trace
 from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+from flock.logging.logging import get_logger
 
-from flock.core.logging.logging import get_logger
 
 logger = get_logger("mcp.tool")
 tracer = trace.get_tracer(__name__)
@@ -31,17 +33,11 @@ class FlockMCPTool(BaseModel):
 
     name: str = Field(..., description="Name of the tool")
 
-    agent_id: str = Field(
-        ..., description="Associated agent_id. Used for internal tracking."
-    )
+    agent_id: str = Field(..., description="Associated agent_id. Used for internal tracking.")
 
-    run_id: str = Field(
-        ..., description="Associated run_id. Used for internal tracking."
-    )
+    run_id: str = Field(..., description="Associated run_id. Used for internal tracking.")
 
-    description: str | None = Field(
-        ..., description="A human-readable description of the tool"
-    )
+    description: str | None = Field(..., description="A human-readable description of the tool")
 
     input_schema: dict[str, Any] = Field(
         ...,
@@ -53,9 +49,7 @@ class FlockMCPTool(BaseModel):
     )
 
     @classmethod
-    def from_mcp_tool(
-        cls: type[T], tool: Tool, agent_id: str, run_id: str
-    ) -> T:
+    def from_mcp_tool(cls, tool: Tool, agent_id: str, run_id: str) -> Self:
         """Convert MCP Tool to Flock Tool."""
         return cls(
             name=tool.name,
@@ -67,7 +61,7 @@ class FlockMCPTool(BaseModel):
         )
 
     @classmethod
-    def to_mcp_tool(cls: type[T], instance: T) -> Tool | None:
+    def to_mcp_tool(cls, instance: Self) -> Tool | None:
         """Convert a flock mcp tool into a mcp tool."""
         return Tool(
             name=instance.name,
@@ -77,17 +71,17 @@ class FlockMCPTool(BaseModel):
         )
 
     # Use DSPy's converter for JSON Schema → Tool args to stay aligned with DSPy.
-    def _convert_input_schema_to_tool_args(self, input_schema: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], dict[str, str]]:
+    def _convert_input_schema_to_tool_args(
+        self, input_schema: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, str]]:
         try:
             return convert_input_schema_to_tool_args(input_schema)
         except Exception as e:  # pragma: no cover - defensive
-            logger.error("Failed to convert MCP tool schema to DSPy tool args: %s", e)
+            logger.exception("Failed to convert MCP tool schema to DSPy tool args: %s", e)
             # Fallback to empty definitions to avoid breaking execution
             return {}, {}, {}
 
-    def _convert_mcp_tool_result(
-        self, call_tool_result: CallToolResult
-    ) -> str | list[Any]:
+    def _convert_mcp_tool_result(self, call_tool_result: CallToolResult) -> str | list[Any]:
         text_contents: list[TextContent] = []
         non_text_contents = []
 
@@ -114,9 +108,7 @@ class FlockMCPTool(BaseModel):
 
     def as_dspy_tool(self, server: Any) -> DSPyTool:
         """Wrap this tool as a DSPyTool for downstream."""
-        args, arg_type, args_desc = self._convert_input_schema_to_tool_args(
-            self.input_schema
-        )
+        args, arg_type, args_desc = self._convert_input_schema_to_tool_args(self.input_schema)
 
         async def func(*args, **kwargs):
             with tracer.start_as_current_span(f"tool.{self.name}.call") as span:
@@ -128,9 +120,7 @@ class FlockMCPTool(BaseModel):
                     logger.debug(
                         f"Tool: {self.name}: got client for server '{server_name}' for agent {self.agent_id} on run {self.run_id}"
                     )
-                    logger.debug(
-                        f"Tool: {self.name}: calling server '{server_name}'"
-                    )
+                    logger.debug(f"Tool: {self.name}: calling server '{server_name}'")
                     result = await server.call_tool(
                         agent_id=self.agent_id,
                         run_id=self.run_id,
@@ -142,7 +132,7 @@ class FlockMCPTool(BaseModel):
                     )
                     return self._convert_mcp_tool_result(result)
                 except Exception as e:
-                    logger.error(
+                    logger.exception(
                         f"Tool: Exception ocurred when calling tool '{self.name}': {e}"
                     )
                     span.record_exception(e)

flock/mcp/types/__init__.py

@@ -0,0 +1,42 @@
+"""MCP Type Definitions for Flock-Flow."""
+
+from flock.mcp.types.factories import (
+    default_flock_mcp_list_roots_callback_factory,
+    default_flock_mcp_logging_callback_factory,
+    default_flock_mcp_message_handler_callback_factory,
+    default_flock_mcp_sampling_callback_factory,
+)
+from flock.mcp.types.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockLoggingMessageNotificationParams,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    MCPRoot,
+    ServerNotification,
+    ServerParameters,
+    SseServerParameters,
+    StdioServerParameters,
+    StreamableHttpServerParameters,
+    WebsocketServerParameters,
+)
+
+
+__all__ = [
+    "FlockListRootsMCPCallback",
+    "FlockLoggingMCPCallback",
+    "FlockLoggingMessageNotificationParams",
+    "FlockMessageHandlerMCPCallback",
+    "FlockSamplingMCPCallback",
+    "MCPRoot",
+    "ServerNotification",
+    "ServerParameters",
+    "SseServerParameters",
+    "StdioServerParameters",
+    "StreamableHttpServerParameters",
+    "WebsocketServerParameters",
+    "default_flock_mcp_list_roots_callback_factory",
+    "default_flock_mcp_logging_callback_factory",
+    "default_flock_mcp_message_handler_callback_factory",
+    "default_flock_mcp_sampling_callback_factory",
+]

flock/{core/mcp → mcp}/types/callbacks.py

@@ -1,5 +1,7 @@
 """MCP Callbacks."""
 
+from typing import Any
+
 from mcp.shared.context import RequestContext
 from mcp.shared.session import RequestResponder
 from mcp.types import (
@@ -13,9 +15,8 @@ from mcp.types import (
     ServerRequest,
 )
 
-from flock.core.logging.logging import FlockLogger
-from flock.core.mcp.mcp_client import Any
-from flock.core.mcp.types.handlers import (
+from flock.logging.logging import FlockLogger
+from flock.mcp.types.handlers import (
     handle_incoming_exception,
     handle_incoming_request,
     handle_incoming_server_notification,
@@ -27,14 +28,12 @@ async def default_sampling_callback(
     ctx: RequestContext, params: CreateMessageRequestParams, logger: FlockLogger
 ) -> ErrorData:
     """Default Callback for Sampling."""
-    logger.info(f"Rejecting Sampling Request.")
+    logger.info("Rejecting Sampling Request.")
     return ErrorData(code=INVALID_REQUEST, message="Sampling not supported.")
 
 
 async def default_message_handler(
-    req: RequestResponder[ServerRequest, ClientResult]
-    | ServerNotification
-    | Exception,
+    req: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
     logger: FlockLogger,
     associated_client: Any,
 ) -> None:
@@ -65,12 +64,9 @@ async def default_list_roots_callback(
 ) -> ListRootsResult | ErrorData:
     """Default List Roots Callback."""
     if associated_client.config.feature_config.roots_enabled:
-        current_roots = await associated_client.get_current_roots()
+        current_roots = await associated_client.get_roots()
         return ListRootsResult(roots=current_roots)
-    else:
-        return ErrorData(
-            code=INVALID_REQUEST, message="List roots not supported."
-        )
+    return ErrorData(code=INVALID_REQUEST, message="List roots not supported.")
 
 
 async def default_logging_callback(
@@ -79,6 +75,4 @@ async def default_logging_callback(
     server_name: str,
 ) -> None:
     """Default Logging Handling Callback."""
-    await handle_logging_message(
-        params=params, logger=logger, server_name=server_name
-    )
+    await handle_logging_message(params=params, logger=logger, server_name=server_name)

flock/{core/mcp → mcp}/types/factories.py

@@ -7,14 +7,14 @@ from mcp.types import (
     CreateMessageRequestParams,
 )
 
-from flock.core.logging.logging import FlockLogger, get_logger
-from flock.core.mcp.types.callbacks import (
+from flock.logging.logging import FlockLogger, get_logger
+from flock.mcp.types.callbacks import (
     default_list_roots_callback,
     default_logging_callback,
     default_message_handler,
     default_sampling_callback,
 )
-from flock.core.mcp.types.types import (
+from flock.mcp.types.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
     FlockLoggingMessageNotificationParams,
@@ -23,6 +23,7 @@ from flock.core.mcp.types.types import (
     ServerNotification,
 )
 
+
 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")
@@ -45,6 +46,8 @@ def default_flock_mcp_logging_callback_factory(
             server_name=associated_client.config.name,
         )
 
+    return _method
+
 
 def default_flock_mcp_sampling_callback_factory(
     associated_client: Any,
@@ -60,9 +63,7 @@ def default_flock_mcp_sampling_callback_factory(
         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
-        )
+        await default_sampling_callback(ctx=ctx, params=params, logger=logger_to_use)
 
     return _method