flock-core 0.4.3__py3-none-any.whl → 0.4.503__py3-none-any.whl

flock/core/flock_factory.py

@@ -1,14 +1,49 @@
 """Factory for creating pre-configured Flock agents."""
 
+import os
 from collections.abc import Callable
-from typing import Any
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import AnyUrl, BaseModel, Field, FileUrl
 
 from flock.core.flock_agent import FlockAgent, SignatureType
 from flock.core.logging.formatters.themes import OutputTheme
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+from flock.core.mcp.mcp_config import (
+    FlockMCPCachingConfigurationBase,
+    FlockMCPCallbackConfigurationBase,
+    FlockMCPFeatureConfigurationBase,
+)
+from flock.core.mcp.types.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    MCPRoot,
+    SseServerParameters,
+    StdioServerParameters,
+    WebsocketServerParameters,
+)
 from flock.evaluators.declarative.declarative_evaluator import (
     DeclarativeEvaluator,
     DeclarativeEvaluatorConfig,
 )
+from flock.mcp.servers.sse.flock_sse_server import (
+    FlockSSEConfig,
+    FlockSSEConnectionConfig,
+    FlockSSEServer,
+)
+from flock.mcp.servers.stdio.flock_stdio_server import (
+    FlockMCPStdioServer,
+    FlockStdioConfig,
+    FlockStdioConnectionConfig,
+)
+from flock.mcp.servers.websockets.flock_websocket_server import (
+    FlockWSConfig,
+    FlockWSConnectionConfig,
+    FlockWSServer,
+)
 from flock.modules.output.output_module import OutputModule, OutputModuleConfig
 from flock.modules.performance.metrics_module import (
     MetricsModule,
@@ -16,9 +51,260 @@ from flock.modules.performance.metrics_module import (
 )
 from flock.workflow.temporal_config import TemporalActivityConfig
 
+LoggingLevel = Literal[
+    "debug",
+    "info",
+    "notice",
+    "warning",
+    "error",
+    "critical",
+    "alert",
+    "emergency",
+]
+
 
 class FlockFactory:
-    """Factory for creating pre-configured Flock agents with common module setups."""
+    """Factory for creating pre-configured Flock agents and pre-configured Flock MCPServers with common module setups."""
+
+    # Classes for type-hints.
+    class StdioParams(BaseModel):
+        """Factory-Params for Stdio-Servers."""
+
+        command: str = Field(
+            ...,
+            description="Command for starting the local script. (e.g. 'uvx', 'bun', 'npx', 'bunx', etc.)",
+        )
+
+        args: list[str] = Field(
+            ...,
+            description="Arguments for starting the local script. (e.g. ['run', './mcp-server.py'])",
+        )
+
+        env: dict[str, Any] | None = Field(
+            default=None,
+            description="Environment variables to pass to the server. (e.g. {'GOOGLE_API_KEY': 'MY_SUPER_SECRET_API_KEY'})",
+        )
+
+        cwd: str | Path | None = Field(
+            default_factory=os.getcwd,
+            description="The working directory to start the script in.",
+        )
+
+        encoding: str = Field(
+            default="utf-8",
+            description="The char-encoding to use when talking to a stdio server. (e.g. 'utf-8', 'ascii', etc.)",
+        )
+
+        encoding_error_handler: Literal["strict", "ignore", "replace"] = Field(
+            default="strict",
+            description="The text encoding error handler. See https://docs.python.org/3/library/codecs.html#codec-base-classes for explanations of possible values",
+        )
+
+    class SSEParams(BaseModel):
+        """Factory-Params for SSE-Servers."""
+
+        url: str | AnyUrl = Field(
+            ...,
+            description="Url the server listens at. (e.g. https://my-mcp-server.io/sse)",
+        )
+
+        headers: dict[str, Any] | None = Field(
+            default=None,
+            description="Additional Headers to pass to the client.",
+        )
+
+        timeout_seconds: float | int = Field(
+            default=5, description="Http Timeout in Seconds."
+        )
+
+        sse_read_timeout_seconds: float | int = Field(
+            default=60 * 5,
+            description="How many seconds to wait for server-sent events until closing the connection. (connections will be automatically re-established.)",
+        )
+
+    class WebsocketParams(BaseModel):
+        """Factory-Params for Websocket Servers."""
+
+        url: str | AnyUrl = Field(
+            ...,
+            description="The url the server listens at. (e.g. ws://my-mcp-server.io/messages)",
+        )
+
+    @staticmethod
+    def create_mcp_server(
+        name: str,
+        connection_params: SSEParams | StdioParams | WebsocketParams,
+        max_retries: int = 3,
+        mount_points: list[str | MCPRoot] | None = None,
+        timeout_seconds: int | float = 10,
+        server_logging_level: LoggingLevel = "error",
+        enable_roots_feature: bool = False,
+        enable_tools_feature: bool = False,
+        enable_sampling_feature: bool = False,
+        enable_prompts_feature: bool = False,
+        sampling_callback: FlockSamplingMCPCallback | None = None,
+        list_roots_callback: FlockListRootsMCPCallback | None = None,
+        logging_callback: FlockLoggingMCPCallback | None = None,
+        message_handler: FlockMessageHandlerMCPCallback | None = None,
+        tool_cache_size: float = 100,
+        tool_cache_ttl: float = 60,
+        resource_contents_cache_size=10,
+        resource_contents_cache_ttl=60 * 5,
+        resource_list_cache_size=100,
+        resource_list_cache_ttl=100,
+        tool_result_cache_size=100,
+        tool_result_cache_ttl=100,
+        description: str | Callable[..., str] | None = None,
+        alert_latency_threshold_ms: int = 30000,
+    ) -> FlockMCPServerBase:
+        """Create a default MCP Server with common modules.
+
+        Allows for creating one of the three default-implementations provided
+        by Flock:
+        - SSE-Server (specify "sse" in type)
+        - Stdio-Server (specify "stdio" in type)
+        - Websockets-Server (specifiy "websockets" in type)
+        """
+        # infer server type from the pydantic model class
+        if isinstance(connection_params, FlockFactory.StdioParams):
+            server_kind = "stdio"
+            concrete_server_cls = FlockMCPStdioServer
+        if isinstance(connection_params, FlockFactory.SSEParams):
+            server_kind = "sse"
+            concrete_server_cls = FlockSSEServer
+        if isinstance(connection_params, FlockFactory.WebsocketParams):
+            server_kind = "websockets"
+            concrete_server_cls = FlockWSServer
+
+        # convert mount points.
+        mounts: list[MCPRoot] = []
+        if mount_points:
+            for item in mount_points:
+                if isinstance(item, MCPRoot):
+                    mounts.append(item)
+                elif isinstance(item, str):
+                    try:
+                        conv = MCPRoot(uri=FileUrl(url=item))
+                        mounts.append(conv)
+                    except Exception:
+                        continue  # ignore
+                else:
+                    continue  # ignore
+
+        # build generic configs
+        feature_config = FlockMCPFeatureConfigurationBase(
+            roots_enabled=enable_roots_feature,
+            tools_enabled=enable_tools_feature,
+            prompts_enabled=enable_prompts_feature,
+            sampling_enabled=enable_sampling_feature,
+        )
+        callback_config = FlockMCPCallbackConfigurationBase(
+            sampling_callback=sampling_callback,
+            list_roots_callback=list_roots_callback,
+            logging_callback=logging_callback,
+            message_handler=message_handler,
+        )
+        caching_config = FlockMCPCachingConfigurationBase(
+            tool_cache_max_size=tool_cache_size,
+            tool_cache_max_ttl=tool_cache_ttl,
+            resource_contents_cache_max_size=resource_contents_cache_size,
+            resource_contents_cache_max_ttl=resource_contents_cache_ttl,
+            resource_list_cache_max_size=resource_list_cache_size,
+            resource_list_cache_max_ttl=resource_list_cache_ttl,
+            tool_result_cache_max_size=tool_result_cache_size,
+            tool_result_cache_max_ttl=tool_result_cache_ttl,
+        )
+        connection_config = None
+        server_config: (
+            FlockStdioConfig | FlockSSEConfig | FlockWSConfig | None
+        ) = None
+
+        # Instantiate correct server + config
+        if server_kind == "stdio":
+            # build stdio config
+            connection_config = FlockStdioConnectionConfig(
+                max_retries=max_retries,
+                connection_parameters=StdioServerParameters(
+                    command=connection_params.command,
+                    args=connection_params.args,
+                    env=connection_params.env,
+                    encoding=connection_params.encoding,
+                    encoding_error_handler=connection_params.encoding_error_handler,
+                    cwd=connection_params.cwd,
+                ),
+                mount_points=mounts,
+                read_timeout_seconds=timeout_seconds,
+                server_logging_level=server_logging_level,
+            )
+            server_config = FlockStdioConfig(
+                name=name,
+                connection_config=connection_config,
+                feature_config=feature_config,
+                caching_config=caching_config,
+                callback_config=callback_config,
+            )
+        elif server_kind == "sse":
+            # build sse config
+            connection_config = FlockSSEConnectionConfig(
+                max_retries=max_retries,
+                connection_parameters=SseServerParameters(
+                    url=connection_params.url,
+                    headers=connection_params.headers,
+                    timeout=connection_params.timeout_seconds,
+                    sse_read_timeout=connection_params.sse_read_timeout_seconds,
+                ),
+                mount_points=mounts,
+                server_logging_level=server_logging_level,
+            )
+
+            server_config = FlockSSEConfig(
+                name=name,
+                connection_config=connection_config,
+                feature_config=feature_config,
+                caching_config=caching_config,
+                callback_config=callback_config,
+            )
+
+        elif server_kind == "websockets":
+            # build websocket config
+            connection_config = FlockWSConnectionConfig(
+                max_retries=max_retries,
+                connection_parameters=WebsocketServerParameters(
+                    url=connection_params.url,
+                ),
+                mount_points=mounts,
+                server_logging_level=server_logging_level,
+            )
+
+            server_config = FlockWSConfig(
+                name=name,
+                connection_config=connection_config,
+                feature_config=feature_config,
+                caching_config=caching_config,
+                callback_config=callback_config,
+            )
+
+        else:
+            raise ValueError(
+                f"Unsupported connection_params type: {type(connection_params)}"
+            )
+
+        if not server_config:
+            raise ValueError(
+                f"Unable to create server configuration for passed params."
+            )
+
+        server = concrete_server_cls(config=server_config)
+
+        metrics_module_config = MetricsModuleConfig(
+            latency_threshold_ms=alert_latency_threshold_ms
+        )
+
+        metrics_module = MetricsModule("metrics", config=metrics_module_config)
+
+        server.add_module(metrics_module)
+
+        return server
 
     @staticmethod
     def create_default_agent(
@@ -28,6 +314,7 @@ class FlockFactory:
         input: SignatureType = None,
         output: SignatureType = None,
         tools: list[Callable[..., Any] | Any] | None = None,
+        servers: list[str | FlockMCPServerBase] | None = None,
         use_cache: bool = True,
         enable_rich_tables: bool = False,
         output_theme: OutputTheme = OutputTheme.abernathy,
@@ -66,6 +353,7 @@ class FlockFactory:
             input=input,
             output=output,
             tools=tools,
+            servers=servers,
             model=model,
             description=description,
             evaluator=evaluator,

flock/core/flock_module.py

@@ -104,3 +104,104 @@ class FlockModule(BaseModel, ABC):
     ) -> None:
         """Called when an error occurs during agent execution."""
         pass
+
+    async def on_pre_server_init(self, server: Any) -> None:
+        """Called before a server initializes."""
+        pass
+
+    async def on_post_server_init(self, server: Any) -> None:
+        """Called after a server initialized."""
+        pass
+
+    async def on_pre_server_terminate(self, server: Any) -> None:
+        """Called before a server terminates."""
+        pass
+
+    async def on_post_server_terminate(self, server: Any) -> None:
+        """Called after a server terminates."""
+        pass
+
+    async def on_server_error(self, server: Any, error: Exception) -> None:
+        """Called when a server errors."""
+        pass
+
+    async def on_connect(
+        self,
+        server: Any,
+        additional_params: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Called before a connection is being established to a mcp server.
+
+        use `server` (type FlockMCPServer) to modify the core behavior of the server.
+        use `additional_params` to 'tack_on' additional configurations (for example additional headers for sse-clients.)
+
+        (For example: modify the server's config)
+        new_config = NewConfigObject(...)
+        server.config = new_config
+
+        Warning:
+            Be very careful when modifying a server's internal state.
+            If you just need to 'tack on' additional information (such as headers)
+            or want to temporarily override certain configurations (such as timeouts)
+            use `additional_params` instead if you can.
+
+        (Or pass additional values downstream:)
+        additional_params["headers"] = { "Authorization": "Bearer 123" }
+        additional_params["read_timeout_seconds"] = 100
+
+
+        Note:
+            `additional_params` resets between mcp_calls.
+            so there is not persistence between individual calls.
+            This choice has been made to allow developers to
+            dynamically switch configurations.
+            (This can be used, for example, to use a module to inject oauth headers for
+            individual users on a call-to-call basis. this also gives you direct control over
+            managing the headers yourself. For example, checking for lifetimes on JWT-Tokens.)
+
+        Note:
+            you can access `additional_params` when you are implementing your own subclasses of
+            FlockMCPClientManager and FlockMCPClient. (with self.additional_params.)
+
+        keys which are processed for `additional_params` in the flock core code are:
+        --- General ---
+
+        "refresh_client": bool -> defaults to False. Indicates whether or not to restart a connection on a call. (can be used when headers oder api-keys change to automatically switch to a new client.)
+        "read_timeout_seconds": float -> How long to wait for a connection to happen.
+
+        --- SSE ---
+
+        "override_headers": bool -> default False. If set to false, additional headers will be appended, if set to True, additional headers will override existing ones.
+        "headers": dict[str, Any] -> Additional Headers injected in sse-clients and ws-clients
+        "sse_read_timeout_seconds": float -> how long until a connection is being terminated for sse-clients.
+        "url": str -> which url the server listens on (allows switching between mcp-servers with modules.)
+
+        --- Stdio ---
+
+        "command": str -> Command to run for stdio-servers.
+        "args": list[str] -> additional paramters for stdio-servers.
+        "env": dict[str, Any] -> Environment-Variables for stdio-servers.
+        "encoding": str -> Encoding to use when talking to stdio-servers.
+        "encoding-error-handler": str -> Encoding error handler to use when talking to stdio-servers.
+
+        --- Websockets ---
+
+        "url": str -> Which url the server listens on (allows switching between mcp-servers with modules.)
+        """
+        pass
+
+    async def on_pre_mcp_call(
+        self,
+        server: Any,
+        arguments: Any | None = None,
+    ) -> None:
+        """Called before any MCP Calls."""
+        pass
+
+    async def on_post_mcp_call(
+        self,
+        server: Any,
+        result: Any | None = None,
+    ) -> None:
+        """Called after any MCP Calls."""
+        pass

flock/core/flock_registry.py

@@ -32,6 +32,7 @@ if TYPE_CHECKING:
     from flock.core.flock_evaluator import FlockEvaluator
     from flock.core.flock_module import FlockModule
     from flock.core.flock_router import FlockRouter
+    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
 
     COMPONENT_BASE_TYPES = (FlockModule, FlockEvaluator, FlockRouter)
 
@@ -43,7 +44,6 @@ else:
     IS_COMPONENT_CHECK_ENABLED = False
 
 # Fallback if core types aren't available during setup
-
 from flock.core.flock_module import FlockModuleConfig
 from flock.core.logging.logging import get_logger
 
@@ -56,7 +56,7 @@ _COMPONENT_CONFIG_MAP: dict[type[BaseModel], type[any]] = {}
 
 
 class FlockRegistry:
-    """Singleton registry for Agents, Callables (functions/methods).
+    """Singleton registry for Agents, Callables (functions/methods) and MCP Servers.
 
     Types (Pydantic/Dataclasses used in signatures), and Component Classes
     (Modules, Evaluators, Routers).
@@ -65,6 +65,7 @@ class FlockRegistry:
     _instance = None
 
     _agents: dict[str, FlockAgent]
+    _servers: dict[str, FlockMCPServerBase]
     _callables: dict[str, Callable]
     _types: dict[str, type]
     _components: dict[str, type]  # For Module, Evaluator, Router classes
@@ -79,6 +80,7 @@ class FlockRegistry:
     def _initialize(self):
         """Initialize the internal dictionaries."""
         self._agents = {}
+        self._servers = {}
         self._callables = {}
         self._types = {}
         self._components = {}
@@ -170,6 +172,35 @@ class FlockRegistry:
             logger.warning(f"Could not determine module/name for object: {obj}")
             return None
 
+    # --- Server Registration ---
+    def register_server(self, server: FlockMCPServerBase) -> None:
+        """Registers a flock mcp server by its name."""
+        if not hasattr(server.config, "name") or not server.config.name:
+            logger.error(
+                "Attempted to register a server without a valid 'name' attribute."
+            )
+            return
+        if (
+            server.config.name in self._servers
+            and self._servers[server.config.name] != server
+        ):
+            logger.warning(
+                f"Server '{server.config.name}' already registered. Overwriting."
+            )
+        self._servers[server.config.name] = server
+        logger.debug(f"Registered server: {server.config.name}")
+
+    def get_server(self, name: str) -> FlockMCPServerBase | None:
+        """Retrieves a registered FlockMCPServer instance by name."""
+        server = self._servers.get(name)
+        if not server:
+            logger.warning(f"Server '{name}' not found in registry.")
+        return server
+
+    def get_all_server_names(self) -> list[str]:
+        """Returns a list of names for all registered servers."""
+        return list(self._servers.keys())
+
     # --- Agent Registration ---
     def register_agent(self, agent: FlockAgent, *, force: bool = False) -> None:
         """Registers a FlockAgent instance by its name.
@@ -361,7 +392,7 @@ class FlockRegistry:
         else:
             # Consider adding dynamic import attempts for types if needed,
             # but explicit registration is generally safer for types.
-            logger.error(f"Type '{type_name}' not found in registry.")
+            logger.warning(f"Type '{type_name}' not found in registry. Will attempt to build it from builtins.")
             raise KeyError(
                 f"Type '{type_name}' not found. Ensure it is registered."
             )
@@ -499,6 +530,8 @@ def get_registry() -> FlockRegistry:
 # Type hinting for decorators to preserve signature
 @overload
 def flock_component(cls: ClassType) -> ClassType: ...  # Basic registration
+
+
 @overload
 def flock_component(
     *, name: str | None = None, config_class: type[ConfigType] | None = None
@@ -542,6 +575,8 @@ def flock_component(
 # Type hinting for decorators
 @overload
 def flock_tool(func: FuncType) -> FuncType: ...
+
+
 @overload
 def flock_tool(
     *, name: str | None = None
@@ -581,6 +616,8 @@ flock_callable = flock_tool
 
 @overload
 def flock_type(cls: ClassType) -> ClassType: ...
+
+
 @overload
 def flock_type(
     *, name: str | None = None

flock/core/flock_server_manager.py

@@ -0,0 +1,136 @@
+"""Manages Server-Lifecycles within the larger lifecycle of Flock."""
+
+import asyncio
+from contextlib import AsyncExitStack
+
+from anyio import Lock
+from pydantic import BaseModel, ConfigDict, Field
+
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+
+
+class FlockServerManager(BaseModel):
+    """Async-context-manager to start/stop a set of Flock MCP servers."""
+
+    servers: list[FlockMCPServerBase] | None = Field(
+        ..., exclude=True, description="The servers to manage."
+    )
+
+    stack: AsyncExitStack | None = Field(
+        default=None,
+        exclude=True,
+        description="Central exit stack for managing the execution context of the servers.",
+    )
+
+    lock: Lock | None = Field(
+        default=None, exclude=True, description="Global lock for mutex access."
+    )
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    def __init__(
+        self,
+        servers: list[FlockMCPServerBase] | None = None,
+        stack: AsyncExitStack | None = None,
+        lock: asyncio.Lock | None = None,
+    ) -> None:
+        """Initialize the FlockServerManager with optional server, stack, and lock references."""
+        super().__init__(
+            servers=servers,
+            stack=stack,
+            lock=lock,
+        )
+
+    def add_server_sync(self, server: FlockMCPServerBase) -> None:
+        """Add a server to be managed by the ServerManager.
+
+        Note:
+          IT IS CRUCIAL THAT THIS METHOD IS NOT CALLED
+          WHEN THE SERVER MANAGER HAS ALREADY BEEN INTIALIZED
+          (with server_manager as manager: ...)
+          OTHERWISE EXECUTION WILL BREAK DOWN.
+        """
+        if self.servers is None:
+            self.servers = []
+
+        self.servers.append(server)
+
+    def remove_server_sync(self, server: FlockMCPServerBase) -> None:
+        """Remove a server from the list of managed servers.
+
+        Note:
+          IT IS CRUCIAL THAT THIS METHOD IS NOT CALLED
+          WHEN THE SERVER MANAGER HAS ALREADY BEEN INITIALIZED
+          (with server_manager as manager: ...)
+          OTHERWISE EXECUTION WILL BREAK DOWN.
+        """
+        if self.servers and server in self.servers:
+            self.servers.remove(server)
+
+    # -- For future use: Allow adding and removal of servers during runtime ---
+    async def add_server_during_runtime(
+        self, server: FlockMCPServerBase
+    ) -> None:
+        """Add a server to the manager and, if already running, start it immediately."""
+        if self.lock is None:
+            self.lock = asyncio.Lock()
+
+        async with self.lock:
+            if self.servers is None:
+                self.servers = []
+
+            self.servers.append(server)
+
+        # If we are already running in async-with, enter the context now
+        if self.stack is not None:
+            await self.stack.enter_async_context(server)
+
+    async def remove_server_during_runtime(
+        self, server: FlockMCPServerBase
+    ) -> None:
+        """Tear down and remove a server from the manager at runtime."""
+        if self.lock is None:
+            self.lock = asyncio.Lock()
+
+        retrieved_server: FlockMCPServerBase | None = None
+
+        async with self.lock:
+            if not self.servers or server not in self.servers:
+                return  # Skip as to not impede application flow
+            else:
+                try:
+                    self.servers.remove(server)
+                    retrieved_server = server
+                except ValueError:
+                    # The server is not present (a little paranoid at this point, but still...)
+                    return
+
+        # tell the server to shut down.
+        if retrieved_server:
+            # trigger the server's own exit hook (this closes its connection_manager, sessions, tools....)
+            await retrieved_server.__aexit__(None, None, None)
+
+    async def __aenter__(self) -> "FlockServerManager":
+        """Enter the asynchronous context for the server manager."""
+        if not self.stack:
+            self.stack = AsyncExitStack()
+
+        if not self.servers:
+            self.servers = []
+
+        if not self.lock:
+            self.lock = asyncio.Lock()
+
+        for srv in self.servers:
+            await self.stack.enter_async_context(srv)
+
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        """Exit the asynchronous context for the server manager."""
+        # Unwind the servers in LIFO order
+        if self.stack is not None:
+            await self.stack.aclose()
+            self.stack = None

flock/core/logging/__init__.py

@@ -1,2 +1,6 @@
 """Flock logging system with Rich integration and structured logging support."""
 
+from .logging import configure_logging
+
+__all__ = ["configure_logging"]
+