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

flock/core/__init__.py

@@ -14,6 +14,12 @@ from flock.core.flock_registry import (
     flock_type,
     get_registry,
 )
+from flock.core.mcp.flock_mcp_server import (
+    FlockMCPServerBase,
+)
+from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
+from flock.core.mcp.mcp_client import FlockMCPClientBase
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManagerBase
 
 __all__ = [
     "Flock",
@@ -22,6 +28,11 @@ __all__ = [
     "FlockEvaluator",
     "FlockEvaluatorConfig",
     "FlockFactory",
+    "FlockMCPClientBase",
+    "FlockMCPClientManagerBase",
+    "FlockMCPServerBase",
+    "FlockMCPServerConfig",
+    "FlockMCPToolBase",
     "FlockModule",
     "FlockModuleConfig",
     "FlockRegistry",

flock/core/flock.py

@@ -22,6 +22,9 @@ _R = TypeVar("_R")
 from box import Box
 from temporalio import workflow
 
+from flock.core.flock_server_manager import FlockServerManager
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+
 with workflow.unsafe.imports_passed_through():
     from datasets import Dataset  # type: ignore
 
@@ -135,7 +138,14 @@ class Flock(BaseModel, Serializable):
     # Marked with underscore to indicate it's managed internally and accessed via property
     _agents: dict[str, FlockAgent]
     _start_agent_name: str | None = None  # For potential pre-configuration
-    _start_input: dict = {}  # Instance attribute overwritten in __init__; kept for typing clarity
+    _start_input: dict = {}  # For potential pre-configuration
+
+    # Internal server storage - not part of the Pydantic model for direct serialization
+    _servers: dict[str, FlockMCPServerBase]
+
+    # Async context-manager for startup and teardown of servers
+    # Not part of the pydantic model
+    _mgr: FlockServerManager
 
     # Pydantic v2 model config
     model_config = {
@@ -175,6 +185,7 @@ class Flock(BaseModel, Serializable):
         enable_temporal: bool = False,
         enable_logging: bool | list[str] = False,
         agents: list[FlockAgent] | None = None,
+        servers: list[FlockMCPServerBase] | None = None,
         temporal_config: TemporalWorkflowConfig | None = None,
         temporal_start_in_process_worker: bool = True,
         **kwargs,
@@ -198,12 +209,31 @@ class Flock(BaseModel, Serializable):
 
         # Initialize runtime attributes AFTER super().__init__()
         self._agents = {}
+        self._servers = {}
         self._start_agent_name = None
         self._start_input = {}
+        self._mgr = FlockServerManager()
 
         # Set up logging based on the enable_logging flag
         self._configure_logging(enable_logging) # Pass original value to _configure_logging
 
+        # Register passed servers
+        # (need to be registered first so that agents can retrieve them from the registry)
+        # This will also add them to the managed list of self._mgr
+        if servers:
+            from flock.core.mcp.flock_mcp_server import (
+                FlockMCPServerBase as ConcreteFlockMCPServer,
+            )
+
+            for server in servers:
+                if isinstance(server, ConcreteFlockMCPServer):
+                    self.add_server(server)
+                else:
+                    logger.warning(
+                        f"Item provided in 'servers' list is not a FlockMCPServer: {type(server)}"
+                    )
+
+
         # Register passed agents
         if agents:
             from flock.core.flock_agent import (
@@ -330,12 +360,48 @@ class Flock(BaseModel, Serializable):
             set_baggage("session_id", session_id)
             logger.debug(f"Generated new session_id: {session_id}")
 
-    def add_agent(self, agent: FlockAgent) -> FlockAgent:
-        """Adds an agent instance to this Flock configuration and registry."""
-        from flock.core.flock_agent import (
-            FlockAgent as ConcreteFlockAgent,  # Local import
+    def add_server(self, server: FlockMCPServerBase) -> FlockMCPServerBase:
+        """Adds a server instance to this Flock configuration and registry as well as set it up to be managed by self._mgr."""
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServerBase as ConcreteFlockMCPServer,
         )
 
+        if not isinstance(server, ConcreteFlockMCPServer):
+            raise TypeError("Provided object is not a FlockMCPServer instance.")
+        if not server.config.name:
+            raise ValueError("Server must have a name.")
+
+        if server.config.name in self.servers:
+            raise ValueError(
+                f"Server with this name already exists. Name: '{server.config.name}'"
+            )
+
+        self._servers[server.config.name] = server
+        FlockRegistry.register_server(server)  # Register globally.
+
+        # Make sure that the server is also added to
+        # the server_list managed by FlockServerManager
+        if not self._mgr:
+            self._mgr = FlockServerManager()
+
+        # Prepare server to be managed by the FlockServerManager
+        logger.info(f"Adding server '{server.config.name}' to managed list.")
+        self._mgr.add_server_sync(server=server)
+        logger.info(f"Server '{server.config.name}' is now on managed list.")
+
+        logger.info(
+            f"Server '{server.config.name}' added to Flock '{self.name}'"
+        )
+        return server
+
+    def add_agent(self, agent: FlockAgent) -> FlockAgent:
+        """Adds an agent instance to this Flock configuration and registry.
+
+        This also registers all servers attached to the agent, if they have not been registered
+        beforehand.
+        """
+        from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+
         if not isinstance(agent, ConcreteFlockAgent):
             raise TypeError("Provided object is not a FlockAgent instance.")
         if not agent.name:
@@ -372,6 +438,11 @@ class Flock(BaseModel, Serializable):
         """Returns the dictionary of agents managed by this Flock instance."""
         return self._agents
 
+    @property
+    def servers(self) -> dict[str, FlockMCPServerBase]:
+        """Returns the dictionary of servers managed by this Flock instance."""
+        return self._servers
+
     def run(
         self,
         start_agent: FlockAgent | str | None = None,
@@ -380,6 +451,7 @@ class Flock(BaseModel, Serializable):
         run_id: str = "",
         box_result: bool = True,
         agents: list[FlockAgent] | None = None,
+        servers: list[FlockMCPServerBase] | None = None,
         memo: dict[str, Any] | None = None
     ) -> Box | dict:
         return self._run_sync(
@@ -390,6 +462,7 @@ class Flock(BaseModel, Serializable):
                 run_id=run_id,
                 box_result=box_result,
                 agents=agents,
+                servers=servers,
                 memo=memo,
             )
         )
@@ -403,15 +476,28 @@ class Flock(BaseModel, Serializable):
         run_id: str = "",
         box_result: bool = True,
         agents: list[FlockAgent] | None = None,
+        servers: list[FlockMCPServerBase] | None = None,
         memo: dict[str, Any] | None = None,
     ) -> Box | dict:
         """Entry point for running an agent system asynchronously."""
-        from flock.core.flock_agent import (
-            FlockAgent as ConcreteFlockAgent,  # Local import
+        # Import here to allow forward reference resolution
+        from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServerBase as ConcreteFlockServer,
         )
 
         with tracer.start_as_current_span("flock.run_async") as span:
-            # Add passed agents first
+            # Add passed servers so that agents have access to them.
+            if servers:
+                for server_obj in servers:
+                    if isinstance(server_obj, ConcreteFlockServer):
+                        self.add_server(server=server_obj)
+                    else:
+                        logger.warning(
+                            f"Item in 'servers' list is not a FlockMCPServer: {type(server_obj)}"
+                        )
+
+            # Add passed agents
             if agents:
                 for agent_obj in agents:
                     if isinstance(agent_obj, ConcreteFlockAgent):
@@ -434,7 +520,7 @@ class Flock(BaseModel, Serializable):
 
             # Default to first agent if only one exists and none specified
             if not start_agent_name and len(self._agents) == 1:
-                start_agent_name = list(self._agents.keys())[0]
+                start_agent_name = next(iter(self._agents.keys()))
             elif not start_agent_name:
                 raise ValueError(
                     "No start_agent specified and multiple/no agents exist in the Flock instance."
@@ -497,44 +583,60 @@ class Flock(BaseModel, Serializable):
                         self.temporal_config.model_dump(mode="json"),
                     )
 
-                logger.info(
-                    "Starting agent execution",
-                    agent=start_agent_name,
-                    enable_temporal=self.enable_temporal,
-                )
-
-                # Execute workflow
-                if not self.enable_temporal:
-                    result = await run_local_workflow(
-                        run_context, box_result=False # Boxing handled below
-                    )
-                else:
-                    result = await run_temporal_workflow(
-                        self, # Pass the Flock instance
-                        run_context,
-                        box_result=False, # Boxing handled below
-                        memo=memo,
+                # At this point, initial setup is done
+                # and flock is ready to execute it's agent_workflow.
+                # Befor that happens, the ServerManager needs to
+                # get the Servers up and running (Populate pools, build connections, start scripts, etc.)
+                async with self._mgr:
+                    # Enter the manager's async context,
+                    # running it's __aenter__ method and starting all registered servers
+                    # after this block ends, self._mgr's __aexit__ will be called
+                    # all servers will be torn down.
+                    logger.info(
+                        f"Entering managed server context. Servers starting up."
                     )
 
-                span.set_attribute("result.type", str(type(result)))
-                result_str = str(result)
-                span.set_attribute(
-                    "result.preview",
-                    result_str[:1000]
-                    + ("..." if len(result_str) > 1000 else ""),
-                )
+                    logger.info(
+                        "Starting agent execution",
+                        agent=start_agent_name,
+                        enable_temporal=self.enable_temporal,
+                    )
 
-                if box_result:
-                    try:
-                        logger.debug("Boxing final result.")
-                        return Box(result)
-                    except ImportError: # Should not happen as Box is a direct dependency
-                        logger.warning(
-                            "Box library not installed (should be direct dependency), returning raw dict."
+                    # Execute workflow
+                    if not self.enable_temporal:
+                        result = await run_local_workflow(
+                            run_context, box_result=False # Boxing handled below
+                        )
+                    else:
+                        result = await run_temporal_workflow(
+                            self, # Pass the Flock instance
+                            run_context,
+                            box_result=False, # Boxing handled below
+                            memo=memo,
                         )
+
+                    span.set_attribute("result.type", str(type(result)))
+                    result_str = str(result)
+                    span.set_attribute(
+                        "result.preview",
+                        result_str[:1000]
+                        + ("..." if len(result_str) > 1000 else ""),
+                    )
+
+                    if box_result:
+                        try:
+                            logger.debug("Boxing final result.")
+                            return Box(result)
+                        except ImportError:
+                            logger.warning(
+                                "Box library not installed, returning raw dict."
+                            )
+                            return result
+                    else:
                         return result
-                else:
-                    return result
+
+                        # The context of self._mgr ends here, meaning, that servers will
+                        # be cleaned up and shut down.
 
             except Exception as e:
                 logger.error(

flock/core/flock_agent.py

@@ -4,11 +4,13 @@
 import asyncio
 import json
 import os
+import uuid
 from abc import ABC
 from collections.abc import Callable
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, TypeVar
 
+from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
 from flock.core.serialization.json_encoder import FlockJSONEncoder
 from flock.workflow.temporal_config import TemporalActivityConfig
 
@@ -59,7 +61,13 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
     Inherits from Pydantic BaseModel, ABC, DSPyIntegrationMixin, and Serializable.
     """
 
+    agent_id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Internal, Unique UUID4 for this agent instance. No need to set it manually. Used for MCP features.",
+    )
+
     name: str = Field(..., description="Unique identifier for the agent.")
+
     model: str | None = Field(
         None,
         description="The model identifier to use (e.g., 'openai/gpt-4o'). If None, uses Flock's default.",
@@ -88,6 +96,11 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
             description="List of callable tools the agent can use. These must be registered.",
         )
     )
+    servers: list[str | FlockMCPServerBase] | None = Field(
+        default=None,
+        description="List of MCP Servers the agent can use to enhance its capabilities. These must be registered.",
+    )
+
     write_to_file: bool = Field(
         default=False,
         description="Write the agent's output to a file.",
@@ -132,9 +145,11 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
         input: SignatureType = None,
         output: SignatureType = None,
         tools: list[Callable[..., Any]] | None = None,
+        servers: list[str | FlockMCPServerBase] | None = None,
         evaluator: "FlockEvaluator | None" = None,
         handoff_router: "FlockRouter | None" = None,
-        modules: dict[str, "FlockModule"] | None = None,  # Use dict for modules
+        # Use dict for modules
+        modules: dict[str, "FlockModule"] | None = None,
         write_to_file: bool = False,
         wait_for_input: bool = False,
         temporal_activity_config: TemporalActivityConfig | None = None,
@@ -147,6 +162,7 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
             input=input,  # Store the raw input spec
             output=output,  # Store the raw output spec
             tools=tools,
+            servers=servers,
             write_to_file=write_to_file,
             wait_for_input=wait_for_input,
             evaluator=evaluator,
@@ -335,6 +351,37 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                     # For now, assume evaluator handles tool resolution if necessary
                     registered_tools = self.tools
 
+                # Retrieve available mcp_tools if the evaluator needs them
+                mcp_tools = []
+                if self.servers:
+                    from flock.core.flock_registry import get_registry
+
+                    FlockRegistry = get_registry()  # Get the registry
+                    for server in self.servers:
+                        registered_server: FlockMCPServerBase | None = None
+                        server_tools = []
+                        if isinstance(server, FlockMCPServerBase):
+                            # check if registered
+                            server_name = server.config.name
+                            registered_server = FlockRegistry.get_server(
+                                server_name
+                            )
+                        else:
+                            # servers must be registered.
+                            registered_server = FlockRegistry.get_server(
+                                name=server
+                            )
+                        if registered_server:
+                            server_tools = await registered_server.get_tools(
+                                agent_id=self.agent_id,
+                                run_id=self.context.run_id,
+                            )
+                        else:
+                            logger.warning(
+                                f"No Server with name '{server}' registered! Skipping."
+                            )
+                        mcp_tools = mcp_tools + server_tools
+
                 # --------------------------------------------------
                 # Optional DI middleware pipeline
                 # --------------------------------------------------
@@ -391,7 +438,10 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                 else:
                     # No DI container – standard execution
                     result = await self.evaluator.evaluate(
-                        self, current_inputs, registered_tools
+                        self,
+                    current_inputs,
+                    registered_tools,
+                    mcp_tools=mcp_tools,
                     )
             except Exception as eval_error:
                 logger.error(
@@ -662,7 +712,14 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
 
         FlockRegistry = get_registry()
 
-        exclude = ["context", "evaluator", "modules", "handoff_router", "tools"]
+        exclude = [
+            "context",
+            "evaluator",
+            "modules",
+            "handoff_router",
+            "tools",
+            "servers",
+        ]
 
         is_descrition_callable = False
         is_input_callable = False
@@ -755,6 +812,25 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                 f"Added {len(serialized_modules)} modules to agent '{self.name}'"
             )
 
+        # --- Serialize Servers ---
+        if self.servers:
+            logger.debug(
+                f"Serializing {len(self.servers)} servers for agent '{self.name}'"
+            )
+            serialized_servers = []
+            for server in self.servers:
+                if isinstance(server, FlockMCPServerBase):
+                    serialized_servers.append(server.config.name)
+                else:
+                    # Write it down as a list of server names.
+                    serialized_servers.append(server)
+
+            if serialized_servers:
+                data["mcp_servers"] = serialized_servers
+                logger.debug(
+                    f"Added {len(serialized_servers)} servers to agent '{self.name}'"
+                )
+
         # --- Serialize Tools (Callables) ---
         if self.tools:
             logger.debug(
@@ -848,6 +924,7 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
         component_configs = {}
         callable_configs = {}
         tool_config = []
+        servers_config = []
         agent_data = {}
 
         component_keys = [
@@ -863,6 +940,8 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
         ]
         tool_key = "tools"
 
+        servers_key = "mcp_servers"
+
         for key, value in data.items():
             if key in component_keys and value is not None:
                 component_configs[key] = value
@@ -870,8 +949,11 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                 callable_configs[key] = value
             elif key == tool_key and value is not None:
                 tool_config = value  # Expecting a list of names
+            elif key == servers_key and value is not None:
+                servers_config = value  # Expecting a list of names
             elif key not in component_keys + callable_keys + [
-                tool_key
+                tool_key,
+                servers_key,
             ]:  # Avoid double adding
                 agent_data[key] = value
             # else: ignore keys that are None or already handled
@@ -987,6 +1069,37 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                         exc_info=True,
                     )
 
+        # --- Deserialize Servers ---
+        agent.servers = []  # Initialize Servers list.
+        if servers_config:
+            logger.debug(
+                f"Deserializing {len(servers_config)} servers for '{agent.name}'"
+            )
+            # Agents keep track of server by getting a list of server names.
+            # The server instances will be retrieved during runtime from the registry. (default behavior)
+
+            for server_name in servers_config:
+                if isinstance(server_name, str):
+                    # Case 1 (default behavior): A server name is passe.
+                    agent.servers.append(server_name)
+                elif isinstance(server_name, FlockMCPServerBase):
+                    # Case 2 (highly unlikely): If someone somehow manages to pass
+                    # an instance of a server during the deserialization step (however that might be achieved)
+                    # check the registry, if the server is already registered, if not, register it
+                    # and store the name in the servers list
+                    FlockRegistry = get_registry()
+                    server_exists = (
+                        FlockRegistry.get_server(server_name.config.name)
+                        is not None
+                    )
+                    if server_exists:
+                        agent.servers.append(server_name.config.name)
+                    else:
+                        FlockRegistry.register_server(
+                            server=server_name
+                        )  # register it.
+                        agent.servers.append(server_name.config.name)
+
         # --- Deserialize Callables ---
         logger.debug(f"Deserializing callable fields for '{agent.name}'")
         # available_callables = registry.get_all_callables() # Incorrect

flock/core/flock_evaluator.py

@@ -47,7 +47,7 @@ class FlockEvaluator(ABC, BaseModel):
 
     @abstractmethod
     async def evaluate(
-        self, agent: Any, inputs: dict[str, Any], tools: list[Any]
+        self, agent: Any, inputs: dict[str, Any], tools: list[Any], mcp_tools: list[Any] | None = None
     ) -> dict[str, Any]:
         """Evaluate inputs to produce outputs."""
         pass