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

flock/core/mcp/types/handlers.py

@@ -0,0 +1,240 @@
+"""Handler functions."""
+
+from collections.abc import Callable
+
+from mcp import CreateMessageRequest
+from mcp.client.session import ClientResponse
+from mcp.shared.context import RequestContext
+from mcp.shared.session import RequestResponder
+from mcp.types import (
+    INTERNAL_ERROR,
+    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,
+    ProgressNotification,
+    ResourceListChangedNotification,
+    ResourceUpdatedNotification,
+    ServerNotification,
+    ToolListChangedNotification,
+)
+
+
+async def handle_incoming_exception(
+    e: Exception,
+    logger_to_use: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Process an incoming exception Message."""
+    server_name = await associated_client.config.name
+
+    # For now, simply log it
+    logger_to_use.error(
+        f"Encountered Exception while communicating with server '{server_name}': {e}"
+    )
+
+
+async def handle_progress_notification(
+    n: ProgressNotification,
+    logger_to_use: FlockLogger,
+    server_name: str,
+) -> None:
+    """Process an incoming progress Notification."""
+    params = n.params
+    progress = params.progress
+    total = params.total or "Unknown"
+    progress_token = params.progressToken
+    metadata = params.meta or {}
+
+    message = f"PROGRESS_NOTIFICATION: Server '{server_name}' reports Progress: {progress}/{total}. (Token: {progress_token}) (Meta Data: {metadata})"
+
+    logger_to_use.info(message)
+
+
+async def handle_cancellation_notification(
+    n: CancelledNotification,
+    logger_to_use: FlockLogger,
+    server_name: str,
+) -> None:
+    """Process an incoming Cancellation Notification."""
+    params = n.params
+    request_id_to_cancel = params.requestId
+    reason = params.reason or "no reason given"
+    metadata = params.meta or {}
+
+    message = f"CANCELLATION_REQUEST: Server '{server_name}' requests to cancel request with id: {request_id_to_cancel}. Reason: {reason}. (Metadata: {metadata})"
+
+    logger_to_use.warning(message)
+
+
+async def handle_resource_update_notification(
+    n: ResourceUpdatedNotification,
+    logger_to_use: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Handle an incoming ResourceUpdatedNotification."""
+    # This also means that the associated client needs to invalidate
+    # its resource_contents_cache
+
+    params = n.params
+    metadata = params.meta or {}
+    uri = params.uri
+
+    message = f"RESOURCE_UPDATE: Server '{associated_client.config.name}' reports change on resoure at: {uri}. (Meta Data: {metadata})"
+
+    logger_to_use.info(message)
+
+    await associated_client.invalidate_resource_contents_cache_entry(key=uri)
+
+
+async def handle_resource_list_changed_notification(
+    n: ResourceListChangedNotification,
+    logger_to_use: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Handle an incoming ResourecListChangedNotification."""
+    # This also means that the associated client needs to invalidate
+    # its resource_contents_cache
+
+    params = n.params or {}
+    metadata = params.meta or {}
+
+    message = f"TOOLS_LIST_CHANGED: Server '{associated_client.config.name}' reports a change in their tools list: {metadata}. Resetting Tools Cache for associated clients."
+
+    logger_to_use.info(message)
+    await associated_client.invalidate_resource_list_cache()
+
+
+async def handle_tool_list_changed_notification(
+    n: ToolListChangedNotification,
+    logger_to_use: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Handle an incoming ToolListChangedNotification."""
+    params = n.params or {}
+    metadata = params.meta or {}
+
+    message = f"TOOLS_LIST_CHANGED: Server '{associated_client.config.name}' reports a change in their tools list: {metadata}. Resetting Tools Cache for associated clients."
+
+    logger_to_use.info(message)
+    await associated_client.invalidate_tool_cache()
+
+
+_SERVER_NOTIFICATION_MAP: dict[type[_MCPServerNotification], Callable] = {
+    ResourceListChangedNotification: handle_resource_list_changed_notification,
+    ResourceUpdatedNotification: handle_resource_update_notification,
+    LoggingMessageNotification: lambda n, log, client: handle_logging_message(
+        params=n,
+        logger=log,
+        server_name=client.config.name,
+    ),
+    ProgressNotification: handle_progress_notification,
+    CancelledNotification: handle_cancellation_notification,
+}
+
+
+async def handle_incoming_server_notification(
+    n: ServerNotification,
+    logger: FlockLogger,
+    client: Any,
+) -> None:
+    """Process an incoming server notification."""
+    handler = _SERVER_NOTIFICATION_MAP.get(type(n.root))
+    if handler:
+        await handler(n.root, logger, client)
+
+
+async def handle_logging_message(
+    params: FlockLoggingMessageNotificationParams,
+    logger: FlockLogger,
+    server_name: str,
+) -> None:
+    """Handle a logging request."""
+    level = params.level
+    method = logger.debug
+    logger_name = params.logger if params.logger else "unknown_remote_logger"
+    metadata = params.meta or {}
+
+    str_level = "DEBUG"
+    prefix = f"Message from Remote MCP Logger '{logger_name}' for server '{server_name}': "
+
+    match level:
+        case "info":
+            method = logger.info
+            str_level = "INFO: "
+        case "notice":
+            method = logger.info
+            str_level = "NOTICE: "
+        case "alert":
+            method = logger.warning
+            str_level = "WARNING: "
+        case "critical":
+            method = logger.error
+            str_level = "CRITICAL: "
+        case "error":
+            method = logger.error
+            str_level = "ERROR: "
+        case "emergency":
+            method = logger.error
+            str_level = "EMERGENCY: "
+
+    full_msg = f"{prefix}{str_level}{params.data} (Meta Data: {metadata})"
+    method(full_msg)
+
+
+async def handle_incoming_request(
+    req: RequestResponder[ServerRequest, ClientResult],
+    logger_to_use: FlockLogger,
+    associated_client: Any,
+) -> None:
+    """Handle generic request."""
+    ctx = RequestContext(
+        request_id=req.request_id,
+        meta=req.request_meta,
+        session=req._session,
+        lifespan_context=None,
+    )
+
+    try:
+        match req.request.root:
+            case CreateMessageRequest(params=req.request.root.params):
+                with req:
+                    # invoke user's sampling callback
+                    # type: ignore
+                    response = await associated_client.sampling_callback(
+                        ctx, req.request.root.params
+                    )
+                    client_resp = ClientResponse.validate_python(response)
+                    await req.respond(client_resp)
+            case ListRootsRequest():
+                with req:
+                    # type: ignore
+                    response = await associated_client.list_roots_callback(ctx)
+                    client_resp = ClientResponse.validate_python(response)
+                    await req.respond(client_resp)
+            case _:
+                # unrecognized -> no-op
+                return
+    except Exception as e:
+        # 1) Log the error and stacktrace
+        logger_to_use.error(
+            f"Error in fallback handle_incoming_request (id={req.request_id}): {e}"
+        )
+        # 2) If the request wasn't already completed, send a JSON-RPC error back
+        if not getattr(req, "_completed", False):
+            with req:
+                err = ErrorData(
+                    code=INTERNAL_ERROR, message=f"Client-side error: {e}"
+                )
+                client_err = ClientResponse.validate_python(err)
+                await req.respond(client_err)
+    return

flock/core/mcp/types/types.py

@@ -0,0 +1,157 @@
+"""Types for Flock's MCP functionality."""
+
+from collections.abc import Awaitable, Callable
+from contextlib import AbstractAsyncContextManager
+from typing import Any
+
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp import (
+    ClientSession,
+    CreateMessageResult,
+    StdioServerParameters as _MCPStdioServerParameters,
+)
+from mcp.shared.context import RequestContext
+from mcp.shared.session import RequestResponder
+from mcp.types import (
+    CancelledNotification as _MCPCancelledNotification,
+    ClientResult,
+    CreateMessageRequestParams,
+    ErrorData,
+    JSONRPCMessage,
+    ListRootsResult,
+    LoggingMessageNotification as _MCPLoggingMessageNotification,
+    LoggingMessageNotificationParams as _MCPLoggingMessageNotificationParams,
+    ProgressNotification as _MCPProgressNotification,
+    PromptListChangedNotification as _MCPPromptListChangedNotification,
+    ResourceListChangedNotification as _MCPResourceListChangedNotification,
+    ResourceUpdatedNotification as _MCPResourceUpdateNotification,
+    Root as _MCPRoot,
+    ServerNotification as _MCPServerNotification,
+    ServerRequest,
+    ToolListChangedNotification as _MCPToolListChangedNotification,
+)
+from pydantic import AnyUrl, BaseModel, ConfigDict, Field
+
+from flock.core.mcp.util.helpers import get_default_env
+
+
+class ServerNotification(_MCPServerNotification):
+    """A notification message sent by the server side."""
+
+
+class CancelledNotification(_MCPCancelledNotification):
+    """Notification, which can be sent bei either side to indicate that it is cancelling a previously issued request."""
+
+
+class ProgressNotification(_MCPProgressNotification):
+    """An out-of band notification used to inform the receiver of a progress update for a long-running request."""
+
+
+class LoggingMessageNotification(_MCPLoggingMessageNotification):
+    """A notification message sent by the server side containing a logging message."""
+
+
+class ResourceUpdatedNotification(_MCPResourceUpdateNotification):
+    """A notification message sent by the server side informing a client about a change in a resource."""
+
+
+class ResourceListChangedNotification(_MCPResourceListChangedNotification):
+    """A notification message sent by the server side informing a client about a change in the list of resources."""
+
+
+class ToolListChangedNotification(_MCPToolListChangedNotification):
+    """A notification message sent by the server side informing a client about a change in the offered tools."""
+
+
+class PromptListChangedNotification(_MCPPromptListChangedNotification):
+    """A notification message sent by the server side informing a client about a change in the list of offered Prompts."""
+
+
+class FlockLoggingMessageNotificationParams(
+    _MCPLoggingMessageNotificationParams
+):
+    """Parameters contained within a Logging Message Notification."""
+
+
+class MCPRoot(_MCPRoot):
+    """Wrapper for mcp.types.Root."""
+
+
+class ServerParameters(BaseModel):
+    """Base Type for server parameters."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+
+class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
+    """Base Type for Stdio Server parameters."""
+
+    env: dict[str, str] | None = Field(
+        default_factory=get_default_env,
+        description="Environment for the MCP Server.",
+    )
+
+
+class WebsocketServerParameters(ServerParameters):
+    """Base Type for Websocket Server params."""
+
+    url: str | AnyUrl = Field(..., description="Url the server listens at.")
+
+
+class SseServerParameters(ServerParameters):
+    """Base Type for SSE Server 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.",
+    )
+
+
+MCPCLientInitFunction = Callable[
+    ...,
+    AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ],
+]
+
+
+FlockSamplingMCPCallback = Callable[
+    [RequestContext, CreateMessageRequestParams],
+    Awaitable[CreateMessageResult | ErrorData],
+]
+
+
+FlockListRootsMCPCallback = Callable[
+    [RequestContext[ClientSession, Any]],
+    Awaitable[ListRootsResult | ErrorData],
+]
+
+FlockLoggingMCPCallback = Callable[
+    [FlockLoggingMessageNotificationParams],
+    Awaitable[None],
+]
+
+FlockMessageHandlerMCPCallback = Callable[
+    [
+        RequestResponder[ServerRequest, ClientResult]
+        | ServerNotification
+        | Exception
+    ],
+    Awaitable[None],
+]

flock/core/mcp/util/helpers.py

@@ -0,0 +1,23 @@
+"""Helper functions for Flock MCP Functionality."""
+
+import hashlib
+import json
+
+from mcp.client.stdio import get_default_environment
+
+
+def get_default_env() -> dict[str, str]:
+    """Returns a default environment object.
+
+    Including only environment-variables
+    deemed safe to inherit.
+    """
+    return get_default_environment()
+
+
+def cache_key_generator(agent_id: str, run_id: str, *args, **kwargs) -> str:
+    """Helper function to generate cache keys for Flock MCP caches."""
+    args_digest = hashlib.md5(
+        json.dumps(kwargs, sort_keys=True).encode()
+    ).hexdigest()
+    return f"{agent_id}:{run_id}:{args_digest}"

flock/core/mixin/dspy_integration.py

@@ -1,10 +1,13 @@
 # src/flock/core/mixin/dspy_integration.py
 """Mixin class for integrating with the dspy library."""
 
+import ast
 import re  # Import re for parsing
 import typing
 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
 
@@ -68,12 +71,15 @@ def _resolve_type_string(type_str: str) -> type:
             # Special handling for Literal
             if BaseType is typing.Literal:
                 # Split literal values, remove quotes, strip whitespace
-                literal_args_raw = split_top_level(args_str)
-                literal_args = tuple(
-                    s.strip().strip("'\"") for s in literal_args_raw
-                )
+                def parse_literal_args(args_str: str) -> tuple[str, ...]:
+                    try:
+                        return tuple(ast.literal_eval(f"[{args_str}]"))
+                    except (SyntaxError, ValueError) as exc:
+                        raise ValueError(f"Cannot parse {args_str!r} as literals") from exc
+
+                literal_args = parse_literal_args(args_str)
                 logger.debug(
-                    f"Parsing Literal arguments: {literal_args_raw} -> {literal_args}"
+                    f"Parsing Literal arguments: {args_str} -> {literal_args}"
                 )
                 resolved_type = typing.Literal[literal_args]  # type: ignore
                 logger.debug(f"Constructed Literal type: {resolved_type}")
@@ -95,7 +101,8 @@ def _resolve_type_string(type_str: str) -> type:
             if BaseType is typing.Optional:
                 if len(resolved_arg_types) != 1:
                     raise ValueError("Optional requires exactly one argument.")
-                resolved_type = typing.Union[resolved_arg_types[0], type(None)]  # type: ignore
+                # type: ignore
+                resolved_type = typing.Union[resolved_arg_types[0], type(None)]
                 logger.debug(
                     f"Constructed Optional type as Union: {resolved_type}"
                 )
@@ -103,7 +110,8 @@ def _resolve_type_string(type_str: str) -> type:
             elif BaseType is typing.Union:
                 if not resolved_arg_types:
                     raise ValueError("Union requires at least one argument.")
-                resolved_type = typing.Union[resolved_arg_types]  # type: ignore
+                # type: ignore
+                resolved_type = typing.Union[resolved_arg_types]
                 logger.debug(f"Constructed Union type: {resolved_type}")
                 return resolved_type
             elif hasattr(
@@ -242,7 +250,8 @@ class DSPyIntegrationMixin:
                 f"Failed to create dynamic type 'dspy_{agent_name}': {e}",
                 exc_info=True,
             )
-            raise TypeError(f"Could not create DSPy signature type: {e}") from e
+            raise TypeError(
+                f"Could not create DSPy signature type: {e}") from e
 
     def _configure_language_model(
         self,
@@ -292,12 +301,15 @@ class DSPyIntegrationMixin:
                 f"Failed to configure DSPy language model '{model}': {e}",
                 exc_info=True,
             )
+            # We need to raise this exception, otherwise Flock will trundle on until it needs dspy.settings.lm and can't find it.
+            raise
 
     def _select_task(
         self,
         signature: Any,
         override_evaluator_type: AgentType,
         tools: list[Any] | None = None,
+        mcp_tools: list[Any] | None = None,
         kwargs: dict[str, Any] = {},
     ) -> Any:
         """Select and instantiate the appropriate DSPy Program/Module."""
@@ -320,19 +332,39 @@ class DSPyIntegrationMixin:
                         f"Item '{tool}' in tools list is not callable, skipping."
                     )
 
+        processed_mcp_tools = []
+        if mcp_tools:
+            for mcp_tool in mcp_tools:
+                if isinstance(mcp_tool, Tool):  # Basic check
+                    processed_mcp_tools.append(mcp_tool)
+                else:
+                    logger.warning(
+                        f"Item '{mcp_tool}' is not a dspy.primitives.Tool, skipping."
+                    )
+
         dspy_program = None
         selected_type = override_evaluator_type
 
         # Determine type if not overridden
         if not selected_type:
             selected_type = (
-                "ReAct" if processed_tools else "Predict"
+                "ReAct" if processed_tools or processed_mcp_tools else "Predict"
             )  # Default logic
 
         logger.debug(
-            f"Selecting DSPy program type: {selected_type} (Tools provided: {bool(processed_tools)})"
+            f"Selecting DSPy program type: {selected_type} (Tools provided: {bool(processed_tools)}) (MCP Tools: {bool(processed_mcp_tools)}"
         )
 
+        # Merge list of native tools and processed tools.
+        # This makes mcp tools appear as native code functions to the llm of the agent.
+        merged_tools = []
+
+        if processed_tools:
+            merged_tools = merged_tools + processed_tools
+
+        if processed_mcp_tools:
+            merged_tools = merged_tools + processed_mcp_tools
+
         try:
             if selected_type == "ChainOfThought":
                 dspy_program = dspy.ChainOfThought(signature, **kwargs)
@@ -340,7 +372,7 @@ class DSPyIntegrationMixin:
                 if not kwargs:
                     kwargs = {"max_iters": 10}
                 dspy_program = dspy.ReAct(
-                    signature, tools=processed_tools or [], **kwargs
+                    signature, tools=merged_tools or [], **kwargs
                 )
             elif selected_type == "Predict":  # Default or explicitly Completion
                 dspy_program = dspy.Predict(signature)
@@ -395,7 +427,8 @@ class DSPyIntegrationMixin:
             final_result = {**inputs, **output_dict}
 
             lm = dspy.settings.get("lm")
-            cost = sum([x["cost"] for x in lm.history if x["cost"] is not None])
+            cost = sum([x["cost"]
+                       for x in lm.history if x["cost"] is not None])
             lm_history = lm.history
 
             return final_result, cost, lm_history

flock/core/serialization/flock_serializer.py

@@ -17,7 +17,8 @@ from pydantic import BaseModel, create_model
 from flock.core.flock_registry import get_registry
 from flock.core.logging.logging import get_logger
 from flock.core.serialization.serialization_utils import (
-    extract_pydantic_models_from_type_string,  # Assuming this handles basic serialization needs
+    # Assuming this handles basic serialization needs
+    extract_pydantic_models_from_type_string,
 )
 
 if TYPE_CHECKING:
@@ -52,9 +53,40 @@ class FlockSerializer:
         )
 
         data["agents"] = {}
+        data["mcp_servers"] = {}
         custom_types = {}
         components = {}
 
+        for name, server_instance in flock_instance._servers.items():
+            try:
+                # Servers handle their own serialization via their to_dict method
+                server_data = server_instance.to_dict(path_type=path_type)
+                data["mcp_servers"][name] = server_data
+
+                # --- Extract Component Information ---
+
+                # Modules
+                if "modules" in server_data:
+                    for module_name, module_data in server_data[
+                        "modules"
+                    ].items():
+                        if module_data and "type" in module_data:
+                            component_type = module_data["type"]
+                            if component_type not in components:
+                                logger.debug(
+                                    f"Adding module component '{component_type}' from module '{module_name}' in server '{name}'"
+                                )
+                                components[component_type] = (
+                                    FlockSerializer._get_component_definition(
+                                        component_type, path_type
+                                    )
+                                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to serialize server '{name}' within Flock: {e}",
+                    exc_info=True,
+                )
+
         for name, agent_instance in flock_instance._agents.items():
             try:
                 logger.debug(f"Serializing agent '{name}'")
@@ -260,6 +292,9 @@ class FlockSerializer:
         # Import concrete types needed for instantiation
         from flock.core.flock import Flock  # Import the actual class
         from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServerBase as ConcreteFlockMCPServer,
+        )
 
         logger.debug(
             f"Deserializing Flock from dict. Provided keys: {list(data.keys())}"
@@ -290,6 +325,8 @@ class FlockSerializer:
             FlockSerializer._check_dependencies(data.pop("dependencies"))
 
         agents_data = data.pop("agents", {})
+        server_data = data.pop("mcp_servers", {})
+        logger.info(f"Found {len(server_data)} servers to deserialize")
         logger.info(f"Found {len(agents_data)} agents to deserialize")
 
         try:
@@ -309,6 +346,20 @@ class FlockSerializer:
                 f"Failed to initialize Flock from dict: {e}"
             ) from e
 
+        # Deserialize and add server AFTER Flock instance exists and BEFORE Agents have been added
+        for name, server_data in server_data.items():
+            try:
+                logger.debug(f"Deserializing server '{name}'")
+                server_data.setdefault("name", name)
+                server_instance = ConcreteFlockMCPServer.from_dict(server_data)
+                flock_instance.add_server(server_instance)
+                logger.debug(f"Successfully added server '{name}' to Flock")
+            except Exception as e:
+                logger.error(
+                    f"Failed to deserialize/add server '{name}': {e}",
+                    exc_info=True,
+                )
+
         # Deserialize and add agents AFTER Flock instance exists
         for name, agent_data in agents_data.items():
             try:

flock/core/util/spliter.py

@@ -188,6 +188,9 @@ if __name__ == "__main__":
 
 
     SAMPLE_18 = "ticket_info : str, reasoning : str, search_queries : list[str], relevant_documents: dict[str, float] | dict of pdf_ids as keys and scores as values"
+
+
+    SAMPLE_19 = "title, headings: list[str], entities_and_metadata: list[dict[str, str]], type:Literal['news', 'blog', 'opinion piece', 'tweet']"
     # ➜ [('münze', 'str', "Deutsche Münzbezeichnung, engl. 'coin'")]
 
     for title, spec in [
@@ -209,6 +212,7 @@ if __name__ == "__main__":
         ("Sample-16", SAMPLE_16),
         ("Sample-17", SAMPLE_17),
         ("Sample-18", SAMPLE_18),
+        ("Sample-19", SAMPLE_19),
     ]:
         print(f"\n{title}")
         for row in parse_schema(spec):