flock-core 0.5.0b2__py3-none-any.whl → 0.5.0b5__py3-none-any.whl

flock/core/mcp/mcp_config.py

@@ -36,14 +36,14 @@ LoggingLevel = Literal[
 ]
 
 
-A = TypeVar("A", bound="FlockMCPCallbackConfigurationBase")
-B = TypeVar("B", bound="FlockMCPConnectionConfigurationBase")
-C = TypeVar("C", bound="FlockMCPConfigurationBase")
-D = TypeVar("D", bound="FlockMCPCachingConfigurationBase")
-E = TypeVar("E", bound="FlockMCPFeatureConfigurationBase")
+A = TypeVar("A", bound="FlockMCPCallbackConfiguration")
+B = TypeVar("B", bound="FlockMCPConnectionConfiguration")
+C = TypeVar("C", bound="FlockMCPConfiguration")
+D = TypeVar("D", bound="FlockMCPCachingConfiguration")
+E = TypeVar("E", bound="FlockMCPFeatureConfiguration")
 
 
-class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
+class FlockMCPCachingConfiguration(BaseModel, Serializable):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -110,7 +110,7 @@ class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
+class FlockMCPCallbackConfiguration(BaseModel, Serializable):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -188,7 +188,7 @@ class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
+class FlockMCPConnectionConfiguration(BaseModel, Serializable):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -298,7 +298,7 @@ class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
+class FlockMCPFeatureConfiguration(BaseModel, Serializable):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
@@ -346,7 +346,7 @@ class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
         )
 
 
-class FlockMCPConfigurationBase(BaseModel, Serializable):
+class FlockMCPConfiguration(BaseModel, Serializable):
     """Base Configuration Class for MCP Clients.
 
     Each Client should implement their own config
@@ -357,22 +357,22 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
         ..., description="Name of the server the client connects to."
     )
 
-    connection_config: FlockMCPConnectionConfigurationBase = Field(
+    connection_config: FlockMCPConnectionConfiguration = Field(
         ..., description="MCP Connection Configuration for a client."
     )
 
-    caching_config: FlockMCPCachingConfigurationBase = Field(
-        default_factory=FlockMCPCachingConfigurationBase,
+    caching_config: FlockMCPCachingConfiguration = Field(
+        default_factory=FlockMCPCachingConfiguration,
         description="Configuration for the internal caches of the client.",
     )
 
-    callback_config: FlockMCPCallbackConfigurationBase = Field(
-        default_factory=FlockMCPCallbackConfigurationBase,
+    callback_config: FlockMCPCallbackConfiguration = Field(
+        default_factory=FlockMCPCallbackConfiguration,
         description="Callback configuration for the client.",
     )
 
-    feature_config: FlockMCPFeatureConfigurationBase = Field(
-        default_factory=FlockMCPFeatureConfigurationBase,
+    feature_config: FlockMCPFeatureConfiguration = Field(
+        default_factory=FlockMCPFeatureConfiguration,
         description="Feature configuration for the client.",
     )
 
@@ -425,7 +425,7 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPConnectionConfigurationBase
+                config_cls = FlockMCPConnectionConfiguration
             instance_data["connection_config"] = config_cls.from_dict(connection_config)
         else:
             raise ValueError(f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}")
@@ -436,10 +436,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCachingConfigurationBase
+                config_cls = FlockMCPCachingConfiguration
             instance_data["caching_config"] = config_cls.from_dict(caching_config)
         else:
-            instance_data["caching_config"] = FlockMCPCachingConfigurationBase()
+            instance_data["caching_config"] = FlockMCPCachingConfiguration()
 
         if feature_config:
             try:
@@ -447,10 +447,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPFeatureConfigurationBase
+                config_cls = FlockMCPFeatureConfiguration
             instance_data["feature_config"] = config_cls.from_dict(feature_config)
         else:
-            instance_data["feature_config"] = FlockMCPFeatureConfigurationBase()
+            instance_data["feature_config"] = FlockMCPFeatureConfiguration()
 
         if callback_config:
             try:
@@ -458,10 +458,10 @@ class FlockMCPConfigurationBase(BaseModel, Serializable):
                 config_cls = config_field.annotation
             except (AttributeError, KeyError):
                 # fallback
-                config_cls = FlockMCPCallbackConfigurationBase
+                config_cls = FlockMCPCallbackConfiguration
             instance_data["callback_config"] = config_cls.from_dict(callback_config)
         else:
-            instance_data["callback_config"] = FlockMCPCallbackConfigurationBase()
+            instance_data["callback_config"] = FlockMCPCallbackConfiguration()
 
         return cls(**{k: v for k, v in instance_data.items()})
 

flock/core/orchestration/flock_execution.py

@@ -57,7 +57,7 @@ class FlockExecution:
 
     def run(
         self,
-        start_agent: "FlockAgent | str | None" = None,
+        agent: "FlockAgent | str | None" = None,
         input: dict | None = None,
         context: FlockContext | None = None,
         run_id: str = "",
@@ -69,7 +69,7 @@ class FlockExecution:
         """Synchronous execution wrapper."""
         return self._run_sync(
             self.run_async(
-                start_agent=start_agent,
+                agent=agent,
                 input=input,
                 context=context,
                 run_id=run_id,
@@ -82,7 +82,7 @@ class FlockExecution:
 
     async def run_async(
         self,
-        start_agent: "FlockAgent | str | None" = None,
+        agent: "FlockAgent | str | None" = None,
         input: dict | None = None,
         context: FlockContext | None = None,
         run_id: str = "",
@@ -95,7 +95,7 @@ class FlockExecution:
         # 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,
+            FlockMCPServer as ConcreteFlockServer,
         )
 
         with tracer.start_as_current_span("flock.run_async") as span:
@@ -120,7 +120,7 @@ class FlockExecution:
                         )
 
             # Determine starting agent name
-            start_agent_name = self._resolve_start_agent(start_agent)
+            start_agent_name = self._resolve_start_agent(agent)
 
             # Setup execution context and input
             run_input = input if input is not None else self.flock._start_input
@@ -131,7 +131,7 @@ class FlockExecution:
             span.set_attribute("input", str(run_input))
             span.set_attribute("run_id", effective_run_id)
             span.set_attribute("enable_temporal", self.flock.enable_temporal)
-            
+
             logger.info(
                 f"Initiating Flock run '{self.flock.name}'. Start Agent: '{start_agent_name}'. Temporal: {self.flock.enable_temporal}."
             )
@@ -169,7 +169,7 @@ class FlockExecution:
                 logger.error(f"Flock run '{self.flock.name}' failed: {e}", exc_info=True)
                 span.record_exception(e)
                 span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
-                
+
                 # Return a consistent error structure
                 error_output = {
                     "error": str(e),
@@ -179,7 +179,7 @@ class FlockExecution:
                 }
                 return Box(error_output) if box_result else error_output
 
-    def _resolve_start_agent(self, start_agent: "FlockAgent | str | None") -> str:
+    def _resolve_start_agent(self, agent: "FlockAgent | str | None") -> str:
         """Resolve the start agent name from various input types."""
         from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
         from flock.core.registry import get_registry
@@ -188,12 +188,12 @@ class FlockExecution:
 
         # Determine starting agent name
         start_agent_name: str | None = None
-        if isinstance(start_agent, ConcreteFlockAgent):
-            start_agent_name = start_agent.name
+        if isinstance(agent, ConcreteFlockAgent):
+            start_agent_name = agent.name
             if start_agent_name not in self.flock._agents:  # Add if not already present
-                self.flock.add_agent(start_agent)
-        elif isinstance(start_agent, str):
-            start_agent_name = start_agent
+                self.flock.add_agent(agent)
+        elif isinstance(agent, str):
+            start_agent_name = agent
         else:  # start_agent is None
             start_agent_name = self.flock._start_agent_name
 

flock/core/orchestration/flock_initialization.py

@@ -3,18 +3,18 @@
 
 import os
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from opentelemetry.baggage import get_baggage, set_baggage
 
-from flock.core.registry import get_registry
 from flock.core.logging.logging import get_logger
+from flock.core.registry import get_registry
 from flock.core.util.cli_helper import init_console
 
 if TYPE_CHECKING:
     from flock.core.flock import Flock
     from flock.core.flock_agent import FlockAgent
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("flock.initialization")
 
@@ -28,7 +28,7 @@ class FlockInitialization:
     def setup(
         self,
         agents: list["FlockAgent"] | None = None,
-        servers: list["FlockMCPServerBase"] | None = None,
+        servers: list["FlockMCPServer"] | None = None,
     ) -> None:
         """Handle all initialization side effects and setup."""
         # Register passed servers first (agents may depend on them)
@@ -64,10 +64,10 @@ class FlockInitialization:
             enable_temporal=self.flock.enable_temporal,
         )
 
-    def _register_servers(self, servers: list["FlockMCPServerBase"]) -> None:
+    def _register_servers(self, servers: list["FlockMCPServer"]) -> None:
         """Register servers with the Flock instance."""
         from flock.core.mcp.flock_mcp_server import (
-            FlockMCPServerBase as ConcreteFlockMCPServer,
+            FlockMCPServer as ConcreteFlockMCPServer,
         )
 
         for server in servers:

flock/core/orchestration/flock_server_manager.py

@@ -3,12 +3,14 @@
 
 from typing import TYPE_CHECKING
 
-from flock.core.flock_server_manager import FlockServerManager as InternalServerManager
+from flock.core.flock_server_manager import (
+    FlockServerManager as InternalServerManager,
+)
 from flock.core.logging.logging import get_logger
 
 if TYPE_CHECKING:
     from flock.core.flock import Flock
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("flock.server_manager")
 
@@ -21,12 +23,12 @@ class FlockServerManager:
         # Use the existing internal server manager
         self._internal_mgr = InternalServerManager()
 
-    def add_server(self, server: "FlockMCPServerBase") -> "FlockMCPServerBase":
+    def add_server(self, server: "FlockMCPServer") -> "FlockMCPServer":
         """Adds a server instance to this Flock configuration and registry as well as set it up to be managed by internal manager."""
-        from flock.core.registry import get_registry
         from flock.core.mcp.flock_mcp_server import (
-            FlockMCPServerBase as ConcreteFlockMCPServer,
+            FlockMCPServer as ConcreteFlockMCPServer,
         )
+        from flock.core.registry import get_registry
 
         registry = get_registry()
 
@@ -60,6 +62,6 @@ class FlockServerManager:
         return await self._internal_mgr.__aexit__(exc_type, exc_val, exc_tb)
 
     @property
-    def servers(self) -> dict[str, "FlockMCPServerBase"]:
+    def servers(self) -> dict[str, "FlockMCPServer"]:
         """Returns the dictionary of servers managed by this Flock instance."""
         return self.flock._servers

flock/core/registry/__init__.py

@@ -5,32 +5,38 @@ This module provides a complete refactor of the FlockRegistry using
 the proven composition pattern from Flock and FlockAgent refactoring.
 """
 
-from flock.core.registry.registry_hub import RegistryHub, get_registry
-from flock.core.registry.decorators import flock_component, flock_tool, flock_callable, flock_type
+
+from flock.core.registry.agent_registry import AgentRegistry
+from flock.core.registry.callable_registry import CallableRegistry
+from flock.core.registry.component_discovery import ComponentDiscovery
 
 # Specialized registry components (for advanced usage)
 from flock.core.registry.component_registry import ComponentRegistry
-from flock.core.registry.callable_registry import CallableRegistry
-from flock.core.registry.agent_registry import AgentRegistry
+from flock.core.registry.config_mapping import ConfigMapping
+from flock.core.registry.decorators import (
+    flock_callable,
+    flock_component,
+    flock_tool,
+    flock_type,
+)
+from flock.core.registry.registry_hub import RegistryHub, get_registry
 from flock.core.registry.server_registry import ServerRegistry
 from flock.core.registry.type_registry import TypeRegistry
-from flock.core.registry.config_mapping import ConfigMapping
-from flock.core.registry.component_discovery import ComponentDiscovery
 
 __all__ = [
     # Main API
     "RegistryHub",
     "get_registry",
-    
+
     # Decorators
     "flock_component",
-    "flock_tool", 
+    "flock_tool",
     "flock_callable",
     "flock_type",
-    
+
     # Specialized registries (for advanced usage)
     "ComponentRegistry",
-    "CallableRegistry", 
+    "CallableRegistry",
     "AgentRegistry",
     "ServerRegistry",
     "TypeRegistry",

flock/core/registry/registry_hub.py

@@ -8,8 +8,9 @@ from flock.core.logging.logging import get_logger
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+
     from flock.core.flock_agent import FlockAgent
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("registry.hub")
 
@@ -83,7 +84,9 @@ class RegistryHub:
     def discovery(self):
         """Component discovery helper."""
         if not hasattr(self, '_discovery_helper'):
-            from flock.core.registry.component_discovery import ComponentDiscovery
+            from flock.core.registry.component_discovery import (
+                ComponentDiscovery,
+            )
             self._discovery_helper = ComponentDiscovery(self)
         return self._discovery_helper
 
@@ -101,11 +104,11 @@ class RegistryHub:
         """Get all registered agent names."""
         return self.agents.get_all_agent_names()
 
-    def register_server(self, server: "FlockMCPServerBase") -> None:
+    def register_server(self, server: "FlockMCPServer") -> None:
         """Register a FlockMCPServer instance."""
         self.servers.register_server(server)
 
-    def get_server(self, name: str) -> "FlockMCPServerBase | None":
+    def get_server(self, name: str) -> "FlockMCPServer | None":
         """Get a registered FlockMCPServer instance by name."""
         return self.servers.get_server(name)
 

flock/core/registry/server_registry.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
 from flock.core.logging.logging import get_logger
 
 if TYPE_CHECKING:
-    from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
 
 logger = get_logger("registry.servers")
 
@@ -17,9 +17,9 @@ class ServerRegistry:
 
     def __init__(self, lock: threading.RLock):
         self._lock = lock
-        self._servers: dict[str, "FlockMCPServerBase"] = {}
+        self._servers: dict[str, FlockMCPServer] = {}
 
-    def register_server(self, server: "FlockMCPServerBase") -> None:
+    def register_server(self, server: "FlockMCPServer") -> None:
         """Register a flock mcp server by its name."""
         if not hasattr(server.config, "name") or not server.config.name:
             logger.error("Attempted to register a server without a valid 'name' attribute.")
@@ -28,11 +28,11 @@ class ServerRegistry:
         with self._lock:
             if server.config.name in self._servers and self._servers[server.config.name] != server:
                 logger.warning(f"Server '{server.config.name}' already registered. Overwriting.")
-            
+
             self._servers[server.config.name] = server
             logger.debug(f"Registered server: {server.config.name}")
 
-    def get_server(self, name: str) -> "FlockMCPServerBase | None":
+    def get_server(self, name: str) -> "FlockMCPServer | None":
         """Retrieve a registered FlockMCPServer instance by name."""
         with self._lock:
             server = self._servers.get(name)
@@ -45,7 +45,7 @@ class ServerRegistry:
         with self._lock:
             return list(self._servers.keys())
 
-    def get_all_servers(self) -> dict[str, "FlockMCPServerBase"]:
+    def get_all_servers(self) -> dict[str, "FlockMCPServer"]:
         """Get all registered servers."""
         with self._lock:
             return self._servers.copy()

flock/core/serialization/flock_serializer.py

@@ -12,9 +12,10 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import BaseModel, create_model
 
+from flock.core.logging.logging import get_logger
+
 # Need registry access
 from flock.core.registry import get_registry
-from flock.core.logging.logging import get_logger
 from flock.core.serialization.serialization_utils import (
     # Assuming this handles basic serialization needs
     extract_pydantic_models_from_type_string,
@@ -292,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(

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/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(