flock-core 0.4.527__py3-none-any.whl → 0.5.0b0__py3-none-any.whl

flock/core/serialization/flock_serializer.py

@@ -1,7 +1,6 @@
 # src/flock/core/serialization/flock_serializer.py
 """Handles serialization and deserialization logic for Flock instances."""
 
-import builtins
 import importlib
 import importlib.util
 import inspect
@@ -13,9 +12,10 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import BaseModel, create_model
 
-# Need registry access
-from flock.core.flock_registry import get_registry
 from flock.core.logging.logging import get_logger
+
+# Need registry access
+from flock.core.registry import get_registry
 from flock.core.serialization.serialization_utils import (
     # Assuming this handles basic serialization needs
     extract_pydantic_models_from_type_string,
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
 
 
 logger = get_logger("serialization.flock")
-FlockRegistry = get_registry()
+registry = get_registry()
 
 
 class FlockSerializer:
@@ -186,7 +186,7 @@ class FlockSerializer:
                         "description_callable"
                     ]
                     description_callable = agent_instance.description
-                    path_str = FlockRegistry.get_callable_path_string(
+                    path_str = registry.get_callable_path_string(
                         description_callable
                     )
                     if path_str:
@@ -205,7 +205,7 @@ class FlockSerializer:
                     )
                     input_callable_name = agent_data["input_callable"]
                     input_callable = agent_instance.input
-                    path_str = FlockRegistry.get_callable_path_string(
+                    path_str = registry.get_callable_path_string(
                         input_callable
                     )
                     if path_str:
@@ -224,7 +224,7 @@ class FlockSerializer:
                     )
                     output_callable_name = agent_data["output_callable"]
                     output_callable = agent_instance.output
-                    path_str = FlockRegistry.get_callable_path_string(
+                    path_str = registry.get_callable_path_string(
                         output_callable
                     )
                     if path_str:
@@ -250,7 +250,7 @@ class FlockSerializer:
                             tool = tool_objs[i]
                             if callable(tool) and not isinstance(tool, type):
                                 path_str = (
-                                    FlockRegistry.get_callable_path_string(tool)
+                                    registry.get_callable_path_string(tool)
                                 )
                                 if path_str:
                                     logger.debug(
@@ -293,7 +293,7 @@ class FlockSerializer:
         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,
+            FlockMCPServer as ConcreteFlockMCPServer,
         )
 
         logger.debug(
@@ -415,7 +415,7 @@ class FlockSerializer:
         type_definitions = {}
         for type_name in type_names:
             try:
-                type_obj = FlockRegistry.get_type(
+                type_obj = registry.get_type(
                     type_name
                 )  # Throws KeyError if not found
                 type_def = FlockSerializer._extract_type_definition(
@@ -490,7 +490,7 @@ class FlockSerializer:
             "file_path": None,
         }
         try:
-            component_class = FlockRegistry.get_component(
+            component_class = registry.get_component(
                 component_type_name
             )  # Raises KeyError if not found
             component_def["module_path"] = getattr(
@@ -540,7 +540,7 @@ class FlockSerializer:
             "file_path": None,
         }
         try:
-            func = FlockRegistry.get_callable(
+            func = registry.get_callable(
                 callable_path
             )  # Raises KeyError if not found
             callable_def["module_path"] = getattr(func, "__module__", "unknown")
@@ -596,7 +596,7 @@ class FlockSerializer:
                     module = importlib.import_module(module_path)
                     if hasattr(module, type_name):
                         type_obj = getattr(module, type_name)
-                        FlockRegistry.register_type(type_obj, type_name)
+                        registry.register_type(type_obj, type_name)
                         logger.info(
                             f"Registered type '{type_name}' from module '{module_path}'"
                         )
@@ -629,7 +629,7 @@ class FlockSerializer:
         type_name: str, type_def: dict[str, Any]
     ) -> None:
         """Dynamically create and register a Pydantic model from schema."""
-        # (Logic remains the same, ensure it uses FlockRegistry.register_type)
+        # (Logic remains the same, ensure it uses registry.register_type)
         schema = type_def.get("schema", {})
         try:
             fields = {}
@@ -641,7 +641,7 @@ class FlockSerializer:
                 fields[field_name] = (field_type, default)
 
             DynamicModel = create_model(type_name, **fields)
-            FlockRegistry.register_type(DynamicModel, type_name)
+            registry.register_type(DynamicModel, type_name)
             logger.info(
                 f"Dynamically created and registered Pydantic model: {type_name}"
             )
@@ -672,27 +672,31 @@ class FlockSerializer:
     @staticmethod
     def _create_dataclass(type_name: str, type_def: dict[str, Any]) -> None:
         """Dynamically create and register a dataclass."""
-        # (Logic remains the same, ensure it uses FlockRegistry.register_type)
+        # (Logic remains the same, ensure it uses registry.register_type)
         from dataclasses import make_dataclass
 
         fields_def = type_def.get("fields", {})
         try:
             fields = []
             for field_name, field_props in fields_def.items():
-                # Safely evaluate type string - requires care!
+                # Safely map type strings to actual types
                 field_type_str = field_props.get("type", "str")
-                try:
-                    field_type = eval(
-                        field_type_str,
-                        {"__builtins__": builtins.__dict__},
-                        {"List": list, "Dict": dict},
-                    )  # Allow basic types
-                except Exception:
-                    field_type = Any
+                type_mapping = {
+                    "str": str,
+                    "int": int,
+                    "float": float,
+                    "bool": bool,
+                    "list": list,
+                    "dict": dict,
+                    "List": list,
+                    "Dict": dict,
+                    "Any": Any,
+                }
+                field_type = type_mapping.get(field_type_str, Any)
                 fields.append((field_name, field_type))
 
             DynamicDataclass = make_dataclass(type_name, fields)
-            FlockRegistry.register_type(DynamicDataclass, type_name)
+            registry.register_type(DynamicDataclass, type_name)
             logger.info(
                 f"Dynamically created and registered dataclass: {type_name}"
             )
@@ -705,7 +709,7 @@ class FlockSerializer:
         path_type: Literal["absolute", "relative"],
     ) -> None:
         """Register component/callable definitions from serialized data."""
-        # (Logic remains the same, ensure it uses FlockRegistry.register_component/register_callable)
+        # (Logic remains the same, ensure it uses registry.register_component/register_callable)
         # Key change: Ensure file_path is handled correctly based on path_type from metadata
         for name, comp_def in component_defs.items():
             logger.debug(
@@ -735,13 +739,13 @@ class FlockSerializer:
                     if hasattr(module, name):
                         obj = getattr(module, name)
                         if kind == "flock_callable" and callable(obj):
-                            FlockRegistry.register_callable(
+                            registry.register_callable(
                                 obj, name
                             )  # Register by simple name
                             # Also register by full path if possible
                             full_path = f"{module_path}.{name}"
                             if full_path != name:
-                                FlockRegistry.register_callable(obj, full_path)
+                                registry.register_callable(obj, full_path)
                             logger.info(
                                 f"Registered callable '{name}' from module '{module_path}'"
                             )
@@ -749,7 +753,7 @@ class FlockSerializer:
                         elif kind == "flock_component" and isinstance(
                             obj, type
                         ):
-                            FlockRegistry.register_component(obj, name)
+                            registry.register_component(obj, name)
                             logger.info(
                                 f"Registered component '{name}' from module '{module_path}'"
                             )
@@ -785,14 +789,14 @@ class FlockSerializer:
                         if hasattr(module, name):
                             obj = getattr(module, name)
                             if kind == "flock_callable" and callable(obj):
-                                FlockRegistry.register_callable(obj, name)
+                                registry.register_callable(obj, name)
                                 logger.info(
                                     f"Registered callable '{name}' from file '{file_path}'"
                                 )
                             elif kind == "flock_component" and isinstance(
                                 obj, type
                             ):
-                                FlockRegistry.register_component(obj, name)
+                                registry.register_component(obj, name)
                                 logger.info(
                                     f"Registered component '{name}' from file '{file_path}'"
                                 )

flock/core/serialization/serialization_utils.py

@@ -24,13 +24,13 @@ from pydantic import BaseModel
 if TYPE_CHECKING:
     pass
 
-from flock.core.flock_registry import get_registry
+from flock.core.registry import get_registry
 from flock.core.logging.logging import get_logger
 
 logger = get_logger("serialization.utils")
 
 # Remove this line to avoid circular import at module level
-# FlockRegistry = get_registry()  # Get singleton instance
+# registry = get_registry()  # Get singleton instance
 
 # --- Serialization Helper ---
 
@@ -188,16 +188,16 @@ def collect_pydantic_models(
 
 def serialize_item(item: Any) -> Any:
     """Recursively prepares an item for serialization (e.g., to dict for YAML/JSON).
-    Converts known callables to their path strings using FlockRegistry.
+    Converts known callables to their path strings using registry.
     Converts Pydantic models using model_dump.
     """
     # Import the registry lazily when needed
-    from flock.core.flock_registry import get_registry
+    from flock.core.registry import get_registry
 
-    FlockRegistry = get_registry()
+    registry = get_registry()
 
     if callable(item) and not isinstance(item, type):
-        path_str = FlockRegistry.get_callable_path_string(
+        path_str = registry.get_callable_path_string(
             item
         )  # Use registry helper
         if path_str:
@@ -248,12 +248,12 @@ def serialize_item(item: Any) -> Any:
     elif isinstance(
         item, type
     ):  # Handle type objects themselves (e.g. if stored directly)
-        type_name = FlockRegistry.get_component_type_name(
+        type_name = registry.get_component_type_name(
             item
         )  # Check components first
         if type_name:
             return {"__component_ref__": type_name}
-        type_name = FlockRegistry._get_path_string(
+        type_name = registry._get_path_string(
             item
         )  # Check regular types/classes by path
         if type_name:
@@ -274,13 +274,13 @@ def serialize_item(item: Any) -> Any:
 
 def deserialize_item(item: Any) -> Any:
     """Recursively processes a deserialized item (e.g., from YAML/JSON dict).
-    Converts reference dicts back to actual callables or types using FlockRegistry.
+    Converts reference dicts back to actual callables or types using registry.
     Handles nested lists and dicts.
     """
     # Import the registry lazily when needed
-    from flock.core.flock_registry import get_registry
+    from flock.core.registry import get_registry
 
-    FlockRegistry = get_registry()
+    registry = get_registry()
 
     if isinstance(item, Mapping):
         if "__callable_ref__" in item and len(item) == 1:
@@ -288,7 +288,7 @@ def deserialize_item(item: Any) -> Any:
             try:
                 # The registry's get_callable needs to handle lookup by simple name OR full path
                 # Or we assume get_callable handles finding the right function from the simple name
-                resolved_callable = FlockRegistry.get_callable(ref_name)
+                resolved_callable = registry.get_callable(ref_name)
                 logger.debug(
                     f"Deserialized callable reference '{ref_name}' to {resolved_callable}"
                 )
@@ -307,7 +307,7 @@ def deserialize_item(item: Any) -> Any:
         elif "__component_ref__" in item and len(item) == 1:
             type_name = item["__component_ref__"]
             try:
-                return FlockRegistry.get_component(type_name)
+                return registry.get_component(type_name)
             except KeyError:
                 logger.error(
                     f"Component reference '{type_name}' not found during deserialization."
@@ -318,7 +318,7 @@ def deserialize_item(item: Any) -> Any:
             try:
                 # For general types, use get_type or fallback to dynamic import like get_callable
                 # Using get_type for now, assuming it needs registration
-                return FlockRegistry.get_type(type_name)
+                return registry.get_type(type_name)
             except KeyError:
                 # Attempt dynamic import as fallback if get_type fails (similar to get_callable)
                 try:
@@ -329,7 +329,7 @@ def deserialize_item(item: Any) -> Any:
                         mod = importlib.import_module(module_name)
                     type_obj = getattr(mod, class_name)
                     if isinstance(type_obj, type):
-                        FlockRegistry.register_type(
+                        registry.register_type(
                             type_obj, type_name
                         )  # Cache it
                         return type_obj
@@ -356,12 +356,12 @@ def deserialize_component(
     data: dict | None, expected_base_type: type
 ) -> Any | None:
     """Deserializes a component (Module, Evaluator, Router) from its dict representation.
-    Uses the 'type' field to find the correct class via FlockRegistry.
+    Uses the 'type' field to find the correct class via registry.
     """
     # Import the registry and COMPONENT_BASE_TYPES lazily when needed
-    from flock.core.flock_registry import COMPONENT_BASE_TYPES, get_registry
+    from flock.core.registry import get_registry
 
-    FlockRegistry = get_registry()
+    registry = get_registry()
 
     if data is None:
         return None
@@ -379,14 +379,17 @@ def deserialize_component(
         return None
 
     try:
-        ComponentClass = FlockRegistry.get_component(type_name)  # Use registry
+        ComponentClass = registry.get_component(type_name)  # Use registry
         # Optional: Keep the base type check
-        if COMPONENT_BASE_TYPES and not issubclass(
-            ComponentClass, expected_base_type
-        ):
-            raise TypeError(
-                f"Deserialized class {type_name} is not a subclass of {expected_base_type.__name__}"
-            )
+        try:
+            from flock.core.component.agent_component_base import AgentComponent
+            if not issubclass(ComponentClass, AgentComponent):
+                logger.warning(
+                    f"Deserialized class '{type_name}' is not a subclass of AgentComponent."
+                )
+        except ImportError:
+            # AgentComponent not available during setup
+            pass
 
         # Recursively deserialize the data *before* passing to Pydantic constructor
         # This handles nested callables/types within the component's config/data

flock/core/util/hydrator.py

@@ -259,7 +259,7 @@ async def _run_hydration_agent(
 
         # Run agent
         result = await temp_flock.run_async(
-            start_agent=agent_name,
+            agent=agent_name,
             input=agent_input_data,
             box_result=False,
         )

flock/core/util/input_resolver.py

@@ -45,7 +45,12 @@ def top_level_to_keys(s: str) -> list[str]:
 
 
 def resolve_inputs(
-    input_spec: str, context: FlockContext, previous_agent_name: str
+    input_spec: str,
+    context: FlockContext,
+    previous_agent_name: str,
+    previous_agent_output:str,
+    previous_agent_handoff_strategy:str,
+    previous_agent_handoff_map:dict[str, str]
 ) -> dict:
     """Build a dictionary of inputs based on the input specification string and the provided context.
 
@@ -63,6 +68,28 @@ def resolve_inputs(
         eg. agent name: "idea_agent", variable: "ia_idea" (ia = idea agent)
         - or set hand off mode to strict to avoid conflicts.
         with strict mode, the agent will only accept inputs from the previous agent.
+        
+
+    Strategy for passing data to the next agent.
+
+    Example:
+    ReviewAgent.next_agent = SummaryAgent
+    ReviewAgent(output = "text:str, keywords:list[str], rating:int")
+    SummaryAgent(input = "text:str, title:str")
+
+    'append' means the difference in signature is appended to the next agent's input signature.
+    SummaryAgent(input = "text:str, title:str, keywords:list[str], rating:int")
+
+    'override' means the target agent's signature is getting overriden.
+    SummaryAgent(input = "text:str, keywords:list[str], rating:int")
+
+    'static' means the the target agent's signature is not changed at all.
+    If source agent has no output fields that match the target agent's input,
+    there will be no data passed to the next agent.
+    SummaryAgent(input = "text:str, title:str")
+
+    'map' means the source agent's output is mapped to the target agent's input
+    based on 'handoff_map' configuration.
 
     Args:
         input_spec: Comma-separated input keys (e.g., "query" or "agent_name.property").
@@ -120,7 +147,7 @@ def resolve_inputs(
                 inputs[key] = context.get_agent_definition(property_name)
                 continue
 
-            # Otherwise, attempt to look up a state variable with the key "agent_name.property"
+            # Otherwise, attempt to look up a state variable with the key "entity_name.property_name"
             inputs[key] = context.get_variable(f"{entity_name}.{property_name}")
             continue
 

flock/mcp/servers/sse/flock_sse_server.py

@@ -15,12 +15,12 @@ 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.flock_mcp_server import FlockMCPServer
+from flock.core.mcp.mcp_client import FlockMCPClient
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
 from flock.core.mcp.mcp_config import (
-    FlockMCPConfigurationBase,
-    FlockMCPConnectionConfigurationBase,
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
 )
 from flock.core.mcp.types.types import SseServerParameters
 
@@ -28,7 +28,7 @@ logger = get_logger("mcp.sse.server")
 tracer = trace.get_tracer(__name__)
 
 
-class FlockSSEConnectionConfig(FlockMCPConnectionConfigurationBase):
+class FlockSSEConnectionConfig(FlockMCPConnectionConfiguration):
     """Concrete ConnectionConfig for an SSEClient."""
 
     # Only thing we need to override here is the concrete transport_type
@@ -42,7 +42,7 @@ class FlockSSEConnectionConfig(FlockMCPConnectionConfigurationBase):
     )
 
 
-class FlockSSEConfig(FlockMCPConfigurationBase):
+class FlockSSEConfig(FlockMCPConfiguration):
     """Configuration for SSE Clients."""
 
     # The only thing we need to override here is the concrete
@@ -53,7 +53,7 @@ class FlockSSEConfig(FlockMCPConfigurationBase):
     )
 
 
-class FlockSSEClient(FlockMCPClientBase):
+class FlockSSEClient(FlockMCPClient):
     """Client for SSE Servers."""
 
     config: FlockSSEConfig = Field(..., description="Client configuration.")
@@ -115,7 +115,7 @@ class FlockSSEClient(FlockMCPClientBase):
         )
 
 
-class FlockSSEClientManager(FlockMCPClientManagerBase):
+class FlockSSEClientManager(FlockMCPClientManager):
     """Manager for handling SSE Clients."""
 
     client_config: FlockSSEConfig = Field(
@@ -132,7 +132,7 @@ class FlockSSEClientManager(FlockMCPClientManagerBase):
         return new_client
 
 
-class FlockSSEServer(FlockMCPServerBase):
+class FlockSSEServer(FlockMCPServer):
     """Class which represents a MCP Server using the SSE Transport type."""
 
     config: FlockSSEConfig = Field(..., description="Config for the server.")

flock/mcp/servers/stdio/flock_stdio_server.py

@@ -14,12 +14,12 @@ 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.flock_mcp_server import FlockMCPServer
+from flock.core.mcp.mcp_client import FlockMCPClient
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
 from flock.core.mcp.mcp_config import (
-    FlockMCPConfigurationBase,
-    FlockMCPConnectionConfigurationBase,
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
 )
 from flock.core.mcp.types.types import StdioServerParameters
 
@@ -27,7 +27,7 @@ logger = get_logger("mcp.stdio.server")
 tracer = trace.get_tracer(__name__)
 
 
-class FlockStdioConnectionConfig(FlockMCPConnectionConfigurationBase):
+class FlockStdioConnectionConfig(FlockMCPConnectionConfiguration):
     """Concrete ConnectionConfig for an StdioClient."""
 
     # Only thing we need to override here is the concrete transport_type
@@ -42,7 +42,7 @@ class FlockStdioConnectionConfig(FlockMCPConnectionConfigurationBase):
     )
 
 
-class FlockStdioConfig(FlockMCPConfigurationBase):
+class FlockStdioConfig(FlockMCPConfiguration):
     """Configuration for Stdio Clients."""
 
     # The only thing we need to override here is the
@@ -53,7 +53,7 @@ class FlockStdioConfig(FlockMCPConfigurationBase):
     )
 
 
-class FlockStdioClient(FlockMCPClientBase):
+class FlockStdioClient(FlockMCPClient):
     """Client for Stdio Servers."""
 
     config: FlockStdioConfig = Field(..., description="Client Configuration.")
@@ -104,7 +104,7 @@ class FlockStdioClient(FlockMCPClientBase):
 
 
 # Not really needed but kept here as an example.
-class FlockStdioClientManager(FlockMCPClientManagerBase):
+class FlockStdioClientManager(FlockMCPClientManager):
     """Manager for handling Stdio Clients."""
 
     client_config: FlockStdioConfig = Field(
@@ -122,7 +122,7 @@ class FlockStdioClientManager(FlockMCPClientManagerBase):
         return new_client
 
 
-class FlockMCPStdioServer(FlockMCPServerBase):
+class FlockMCPStdioServer(FlockMCPServer):
     """Class which represents a MCP Server using the Stdio Transport type.
 
     This means (most likely) that the server is a locally

flock/mcp/servers/streamable_http/flock_streamable_http_server.py

@@ -17,12 +17,12 @@ 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.flock_mcp_server import FlockMCPServer
+from flock.core.mcp.mcp_client import FlockMCPClient
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
 from flock.core.mcp.mcp_config import (
-    FlockMCPConfigurationBase,
-    FlockMCPConnectionConfigurationBase,
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
 )
 from flock.core.mcp.types.types import (
     StreamableHttpServerParameters,
@@ -34,7 +34,7 @@ tracer = trace.get_tracer(__name__)
 GetSessionIdCallback = Callable[[], str | None]
 
 
-class FlockStreamableHttpConnectionConfig(FlockMCPConnectionConfigurationBase):
+class FlockStreamableHttpConnectionConfig(FlockMCPConnectionConfiguration):
     """Concrete ConnectionConfig for a StreamableHttpClient."""
 
     # Only thing we need to override here is the concrete transport_type
@@ -49,7 +49,7 @@ class FlockStreamableHttpConnectionConfig(FlockMCPConnectionConfigurationBase):
     )
 
 
-class FlockStreamableHttpConfig(FlockMCPConfigurationBase):
+class FlockStreamableHttpConfig(FlockMCPConfiguration):
     """Configuration for Streamable HTTP Clients."""
 
     # The only thing we need to override here is the
@@ -60,7 +60,7 @@ class FlockStreamableHttpConfig(FlockMCPConfigurationBase):
     )
 
 
-class FlockStreamableHttpClient(FlockMCPClientBase):
+class FlockStreamableHttpClient(FlockMCPClient):
     """Client for StreamableHttpServers."""
 
     config: FlockStreamableHttpConfig = Field(
@@ -135,7 +135,7 @@ class FlockStreamableHttpClient(FlockMCPClientBase):
         )
 
 
-class FlockStreamableHttpClientManager(FlockMCPClientManagerBase):
+class FlockStreamableHttpClientManager(FlockMCPClientManager):
     """Manager for handling StreamableHttpClients."""
 
     client_config: FlockStreamableHttpConfig = Field(
@@ -153,7 +153,7 @@ class FlockStreamableHttpClientManager(FlockMCPClientManagerBase):
         return new_client
 
 
-class FlockStreamableHttpServer(FlockMCPServerBase):
+class FlockStreamableHttpServer(FlockMCPServer):
     """Class which represents a MCP Server using the streamable Http Transport type."""
 
     config: FlockStreamableHttpConfig = Field(

flock/mcp/servers/websockets/flock_websocket_server.py

@@ -14,12 +14,12 @@ 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.flock_mcp_server import FlockMCPServer
+from flock.core.mcp.mcp_client import FlockMCPClient
+from flock.core.mcp.mcp_client_manager import FlockMCPClientManager
 from flock.core.mcp.mcp_config import (
-    FlockMCPConfigurationBase,
-    FlockMCPConnectionConfigurationBase,
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
 )
 from flock.core.mcp.types.types import WebsocketServerParameters
 
@@ -28,7 +28,7 @@ tracer = trace.get_tracer(__name__)
 
 
 # Optional to provide type hints.
-class FlockWSConnectionConfig(FlockMCPConnectionConfigurationBase):
+class FlockWSConnectionConfig(FlockMCPConnectionConfiguration):
     """Concrete ConnectionConfig for a WS Client."""
 
     # Only thing we need to override here is the concrete transport_type
@@ -44,7 +44,7 @@ class FlockWSConnectionConfig(FlockMCPConnectionConfigurationBase):
 
 
 # Optional to provide type hints.
-class FlockWSConfig(FlockMCPConfigurationBase):
+class FlockWSConfig(FlockMCPConfiguration):
     """Configuration for Websocket clients."""
 
     # The only thing we need to override here is the concrete
@@ -56,7 +56,7 @@ class FlockWSConfig(FlockMCPConfigurationBase):
     )
 
 
-class FlockWSClient(FlockMCPClientBase):
+class FlockWSClient(FlockMCPClient):
     """Client for Websocket servers."""
 
     config: FlockWSConfig = Field(..., description="Client Configuration")
@@ -90,7 +90,7 @@ class FlockWSClient(FlockMCPClientBase):
 
 
 # not really needed, but kept for type hints and as an example.
-class FlockWSClientManager(FlockMCPClientManagerBase):
+class FlockWSClientManager(FlockMCPClientManager):
     """Manager for handling websocket clients."""
 
     client_config: FlockWSConfig = Field(
@@ -106,7 +106,7 @@ class FlockWSClientManager(FlockMCPClientManagerBase):
         return new_client
 
 
-class FlockWSServer(FlockMCPServerBase):
+class FlockWSServer(FlockMCPServer):
     """Class which represents an MCP Server using the websocket transport type."""
 
     config: FlockWSConfig = Field(..., description="Config for the server.")