flock-core 0.4.2__py3-none-any.whl → 0.4.5__py3-none-any.whl

flock/evaluators/declarative/declarative_evaluator.py

@@ -52,7 +52,7 @@ class DeclarativeEvaluator(
     _lm_history: list = PrivateAttr(default_factory=list)
 
     async def evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any], tools: list[Any]
+        self, agent: FlockAgent, inputs: dict[str, Any], tools: list[Any], mcp_tools: list[Any] | None = None
     ) -> dict[str, Any]:
         """Evaluate using DSPy, with optional asynchronous streaming."""
         # --- Setup Signature and LM ---
@@ -84,6 +84,7 @@ class DeclarativeEvaluator(
                 _dspy_signature,
                 override_evaluator_type=self.config.override_evaluator_type,
                 tools=tools,
+                mcp_tools=mcp_tools,
                 kwargs=self.config.kwargs,
             )
         except Exception as setup_error:
@@ -106,7 +107,7 @@ class DeclarativeEvaluator(
                     "DSPy task could not be created or is not callable."
                 )
 
-            streaming_task = dspy.streamify(agent_task)
+            streaming_task = dspy.streamify(agent_task, is_async_program=True)
             stream_generator: Generator = streaming_task(**inputs)
             delta_content = ""
 
@@ -139,7 +140,7 @@ class DeclarativeEvaluator(
             logger.info(f"Evaluating agent '{agent.name}' without streaming.")
             try:
                 # Ensure the call is awaited if the underlying task is async
-                result_obj = agent_task(**inputs)
+                result_obj = await agent_task.acall(**inputs)
                 result_dict, cost, lm_history = self._process_result(
                     result_obj, inputs
                 )

flock/mcp/servers/sse/__init__.py

@@ -0,0 +1 @@
+"""Default SSE Server Implementation for Flock."""

flock/mcp/servers/sse/flock_sse_server.py

@@ -0,0 +1,139 @@
+"""This module provides the Flock SSE Server functionality."""
+
+import copy
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Literal
+
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp.client.sse import sse_client
+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 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 SseServerParameters
+
+logger = get_logger("mcp.sse.server")
+tracer = trace.get_tracer(__name__)
+
+
+class FlockSSEConnectionConfig(FlockMCPConnectionConfigurationBase):
+    """Concrete ConnectionConfig for an SSEClient."""
+
+    # 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."
+    )
+
+    connection_parameters: SseServerParameters = Field(
+        ..., description="SSE Server Connection Parameters."
+    )
+
+
+class FlockSSEConfig(FlockMCPConfigurationBase):
+    """Configuration for SSE 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: FlockSSEConnectionConfig = Field(
+        ..., description="Concrete SSE Connection Configuration."
+    )
+
+
+class FlockSSEClient(FlockMCPClientBase):
+    """Client for SSE Servers."""
+
+    config: FlockSSEConfig = Field(..., description="Client configuration.")
+
+    async def create_transport(
+        self,
+        params: SseServerParameters,
+        additional_params: dict[str, Any] | None = None,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """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:
+            override_headers = bool(
+                self.additional_params.get("override_headers", False)
+            )
+            if "headers" in self.additional_params:
+                if override_headers:
+                    param_copy.headers = self.additional_params.get(
+                        "headers", params.headers
+                    )
+                else:
+                    param_copy.headers.update(
+                        self.additional_params.get("headers", {})
+                    )
+            if "read_timeout_seconds" in self.additional_params:
+                param_copy.timeout = self.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(
+                    "sse_read_timeout",
+                    params.sse_read_timeout,
+                )
+            if "url" in self.additional_params:
+                param_copy.url = self.additional_params.get(
+                    "url",
+                    params.url,
+                )
+
+        return sse_client(
+            url=param_copy.url,
+            headers=param_copy.headers,
+            timeout=float(param_copy.timeout),
+            sse_read_timeout=float(param_copy.sse_read_timeout),
+        )
+
+
+class FlockSSEClientManager(FlockMCPClientManagerBase):
+    """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(FlockMCPServerBase):
+    """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

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

flock/mcp/servers/stdio/flock_stdio_server.py

@@ -0,0 +1,138 @@
+"""This module provides the default implementation for MCP servers using the stdio transport."""
+
+import copy
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Literal
+
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp import stdio_client
+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 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 StdioServerParameters
+
+logger = get_logger("mcp.stdio.server")
+tracer = trace.get_tracer(__name__)
+
+
+class FlockStdioConnectionConfig(FlockMCPConnectionConfigurationBase):
+    """Concrete ConnectionConfig for an StdioClient."""
+
+    # Only thing we need to override here is the concrete transport_type
+    # and connection_parameters fields.
+    transport_type: Literal["stdio"] = Field(
+        default="stdio", description="Use the stdio transport type."
+    )
+
+    connection_parameters: StdioServerParameters = Field(
+        ...,
+        description="StdioServerParameters to be used for the stdio transport.",
+    )
+
+
+class FlockStdioConfig(FlockMCPConfigurationBase):
+    """Configuration for Stdio 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: FlockStdioConnectionConfig = Field(
+        ..., description="Concrete Stdio Connection Configuration."
+    )
+
+
+class FlockStdioClient(FlockMCPClientBase):
+    """Client for Stdio Servers."""
+
+    config: FlockStdioConfig = Field(..., description="Client Configuration.")
+
+    async def create_transport(
+        self,
+        params: StdioServerParameters,
+        additional_params: dict[str, Any] | None = None,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """Return an async context manager whose __aenter__ method yields (read_stream, send_stream)."""
+        # additional_params take precedence over passed config, as modules can influence
+        # how to connect to a stdio server.
+
+        # avoid modifying the config of the client as a side-effect.
+        param_copy = copy.deepcopy(params)
+
+        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
+                )
+            if "args" in additional_params:
+                param_copy.args = additional_params.get("args", params.command)
+            if "env" in additional_params:
+                param_copy.env = additional_params.get("env", params.env)
+
+            if "cwd" in additional_params:
+                param_copy.cwd = additional_params.get("cwd", params.env)
+
+            if "encoding" in additional_params:
+                param_copy.encoding = additional_params.get(
+                    "encoding", params.encoding
+                )
+
+            if "encoding_error_handler" in additional_params:
+                param_copy.encoding_error_handler = additional_params.get(
+                    "encoding_error_handler", params.encoding_error_handler
+                )
+
+        # stdio_client already is an AsyncContextManager
+        return stdio_client(server=param_copy)
+
+
+# Not really needed but kept here as an example.
+class FlockStdioClientManager(FlockMCPClientManagerBase):
+    """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(FlockMCPServerBase):
+    """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/websockets/__init__.py

@@ -0,0 +1 @@
+"""Default Websocket Server Implementation."""

flock/mcp/servers/websockets/flock_websocket_server.py

@@ -0,0 +1,119 @@
+"""This module provides the default implementation for MCP servers using the websocket transport."""
+
+import copy
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Literal
+
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp.client.websocket import websocket_client
+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 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 WebsocketServerParameters
+
+logger = get_logger("mcp.ws.server")
+tracer = trace.get_tracer(__name__)
+
+
+# Optional to provide type hints.
+class FlockWSConnectionConfig(FlockMCPConnectionConfigurationBase):
+    """Concrete ConnectionConfig for a WS Client."""
+
+    # Only thing we need to override here is the concrete transport_type
+    # and connection_parameters fields.
+    transport_type: Literal["websockets"] = Field(
+        default="websockets", description="Use the websockets transport type."
+    )
+
+    connection_parameters: WebsocketServerParameters = Field(
+        ...,
+        description="WebsocketServer parameters to be used for the websocket transport.",
+    )
+
+
+# Optional to provide type hints.
+class FlockWSConfig(FlockMCPConfigurationBase):
+    """Configuration for Websocket clients."""
+
+    # The only thing we need to override here is the concrete
+    # connection config. The rest is generic enough to handle
+    # everything else. (This is just here so that type hints work for the
+    # rest of the implementation, we could just omit this override entirely.)
+    connection_config: FlockWSConnectionConfig = Field(
+        ..., description="Concrete WS connection configuration"
+    )
+
+
+class FlockWSClient(FlockMCPClientBase):
+    """Client for Websocket servers."""
+
+    config: FlockWSConfig = Field(..., description="Client Configuration")
+
+    # This one we HAVE to specify. This tells Flock
+    # how to create the underlying connection.
+    async def create_transport(
+        self,
+        params: WebsocketServerParameters,
+        additional_params: dict[str, Any] | None = None,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """Return an async context manager whose __aenter__ method yields a read_stream and a send_stream."""
+        # additional_params take precedence over passed config, as modules
+        # can influece how to connect to a ws server.
+
+        # avoid modifying the underlying config directly
+        param_copy = copy.deepcopy(params)
+
+        if additional_params and "url" in additional_params:
+            # 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(FlockMCPClientManagerBase):
+    """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(FlockMCPServerBase):
+    """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

flock/modules/performance/metrics_module.py

@@ -15,6 +15,7 @@ from flock.core.context.context import FlockContext
 from flock.core.flock_agent import FlockAgent
 from flock.core.flock_module import FlockModule, FlockModuleConfig
 from flock.core.flock_registry import flock_component
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
 
 
 class MetricPoint(BaseModel):
@@ -92,7 +93,10 @@ class MetricsModule(FlockModule):
         MetricsModule._INSTANCE = self
         self._metrics = defaultdict(list)
         self._start_time: float | None = None
+        self._server_start_time: float | None = None
         self._start_memory: int | None = None
+        self._server_start_memory: int | None = None
+        self._client_refreshs: int = 0
 
         # Set up storage
         if self.config.storage_type == "json":
@@ -258,7 +262,7 @@ class MetricsModule(FlockModule):
             # Get all unique metric names from files
             all_metrics = self._load_metrics_from_files()
 
-            for metric_name in all_metrics.keys():
+            for metric_name in all_metrics:
                 stats = self.get_statistics(metric_name)
                 if stats:  # Only include metrics that have data
                     summary["metrics"][metric_name] = stats
@@ -518,3 +522,157 @@ class MetricsModule(FlockModule):
         if instance is None:
             return  # silently ignore if module isn't active
         instance._record_metric(name, value, tags or {})
+
+    # --- MCP Server Lifecycle Hooks ---
+    async def on_server_error(
+        self, server: FlockMCPServerBase, error: Exception
+    ) -> None:
+        """Record server error metrics."""
+        self._record_metric(
+            "errors",
+            1,
+            {
+                "server": server.config.name,
+                "error_type": type(error).__name__,
+            },
+        )
+
+    async def on_pre_server_init(self, server: FlockMCPServerBase):
+        """Initialize metrics collection for server."""
+        self._server_start_time = time.time()
+
+        if self.config.collect_memory:
+            self._server_start_memory = psutil.Process().memory_info().rss
+            self._record_metric(
+                "server_memory",
+                self._server_start_memory,
+                {"server": server.config.name, "phase": "pre_init"},
+            )
+
+    async def on_post_server_init(self, server: FlockMCPServerBase):
+        """Collect metrics after server starts."""
+        if self.config.collect_memory:
+            checkpoint_memory = psutil.Process().memory_info().rss
+            self._record_metric(
+                "server_memory",
+                checkpoint_memory,
+                {"server": server.config.name, "phase": "post_init"},
+            )
+
+    async def on_pre_server_terminate(self, server: FlockMCPServerBase):
+        """Collect metrics before server terminates."""
+        if self.config.collect_memory:
+            checkpoint_memory = psutil.Process().memory_info().rss
+            self._record_metric(
+                "server_memory",
+                checkpoint_memory,
+                {"server": server.config.name, "phase": "pre_terminate"},
+            )
+
+    async def on_post_server_terminate(self, server: FlockMCPServerBase):
+        """Collect metrics after server terminates.
+
+        Clean up and final metric recording.
+        """
+        if self.config.storage_type == "json":
+            # Save aggregated metrics
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            summary_file = os.path.join(
+                self.config.metrics_dir,
+                f"summary_{server.config.name}_{timestamp}.json",
+            )
+
+            # Calculate summary for all metrics
+            summary = {
+                "server": server.config.name,
+                "timestamp": timestamp,
+                "metrics": {},
+            }
+
+            # Get all unique metric names from files
+            all_metrics = self._load_metrics_from_files()
+
+            for metric_name in all_metrics:
+                stats = self.get_statistics(metric_name)
+                if stats:  # Only include metrics that have data
+                    summary["metrics"][metric_name] = stats
+            with open(summary_file, "w") as f:
+                json.dump(summary, f, indent=2)
+
+    async def on_pre_mcp_call(
+        self, server: FlockMCPServerBase, arguments: Any | None = None
+    ):
+        """Record pre-call metrics."""
+        if self.config.collect_cpu:
+            cpu_percent = psutil.Process().cpu_percent()
+            self._record_metric(
+                "cpu",
+                cpu_percent,
+                {"server": server.config.name, "phase": "pre_mcp_call"},
+            )
+        if self.config.collect_memory:
+            current_memory = psutil.Process().memory_info().rss
+            memory_diff = current_memory - self._server_start_memory
+            self._record_metric(
+                "memory",
+                memory_diff,
+                {"server": server.config.name, "phase": "pre_mcp_call"},
+            )
+
+        if isinstance(arguments, dict):
+            self._record_metric(
+                "arguments",
+                len(arguments),
+                {
+                    "server": server.config.name,
+                    "phase": "pre_mcp_call",
+                }.update(arguments),
+            )
+
+    async def on_post_mcp_call(
+        self, server: FlockMCPServerBase, result: Any | None = None
+    ):
+        """Record post-call metrics."""
+        if self.config.collect_timing and self._server_start_time:
+            latency = time.time() - self._server_start_time
+            self._record_metric(
+                "latency", latency, {"server": server.config.name}
+            )
+
+            # Check for alerts
+            if self._should_alert("latency", latency):
+                # In practice, you'd want to integrate with a proper alerting system
+                print(f"ALERT: High latency detected: {latency * 1000:.2f}ms")
+
+        if self.config.collect_cpu:
+            cpu_percent = psutil.Process().cpu_percent()
+            self._record_metric(
+                "cpu",
+                cpu_percent,
+                {"server": server.config.name, "phase": "post_mcp_call"},
+            )
+        if self.config.collect_memory:
+            current_memory = psutil.Process().memory_info().rss
+            memory_diff = current_memory - self._server_start_memory
+            self._record_metric(
+                "memory",
+                memory_diff,
+                {"server": server.config.name, "phase": "post_mcp_call"},
+            )
+
+    async def on_connect(
+        self, server: FlockMCPServerBase, additional_params: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Collect metrics during connect."""
+        # We should track the refresh rate for clients
+        if "refresh_client" in additional_params and additional_params.get(
+            "refresh_client", False
+        ):
+            self._client_refreshs += 1
+            self._record_metric(
+                "client_refreshs",
+                self._client_refreshs,
+                {"server": server.config.name, "phase": "connect"},
+            )
+
+        return additional_params