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

flock/core/registry/decorators.py

@@ -0,0 +1,137 @@
+# src/flock/core/registry/decorators.py
+"""Registry decorators for component, tool, and type registration."""
+
+import inspect
+from collections.abc import Callable
+from typing import Any, TypeVar, overload
+
+from flock.core.registry.registry_hub import get_registry
+
+ClassType = TypeVar("ClassType", bound=type)
+FuncType = TypeVar("FuncType", bound=Callable)
+ConfigType = TypeVar("ConfigType")
+
+
+# --- Component Registration Decorator ---
+
+@overload
+def flock_component(cls: ClassType) -> ClassType: ...  # Basic registration
+
+
+@overload
+def flock_component(
+    *, name: str | None = None, config_class: type[ConfigType] | None = None
+) -> Callable[[ClassType], ClassType]: ...  # With options
+
+
+def flock_component(
+    cls: ClassType | None = None,
+    *,
+    name: str | None = None,
+    config_class: type[ConfigType] | None = None,
+) -> Any:
+    """Decorator to register a Flock Component class and optionally link its config class."""
+    registry = get_registry()
+
+    def decorator(inner_cls: ClassType) -> ClassType:
+        if not inspect.isclass(inner_cls):
+            raise TypeError("@flock_component can only decorate classes.")
+
+        component_name = name or inner_cls.__name__
+        registry.register_component(inner_cls, name=component_name)
+
+        # If config_class is provided, register the mapping
+        if config_class:
+            registry.register_config_component_pair(config_class, inner_cls)
+
+        return inner_cls
+
+    if cls is None:
+        # Called as @flock_component(name="...", config_class=...)
+        return decorator
+    else:
+        # Called as @flock_component
+        return decorator(cls)
+
+
+# --- Tool/Callable Registration Decorator ---
+
+@overload
+def flock_tool(func: FuncType) -> FuncType: ...
+
+
+@overload
+def flock_tool(
+    *, name: str | None = None
+) -> Callable[[FuncType], FuncType]: ...
+
+
+def flock_tool(func: FuncType | None = None, *, name: str | None = None) -> Any:
+    """Decorator to register a callable function/method as a Tool (or general callable).
+
+    Usage:
+        @flock_tool
+        def my_web_search(query: str): ...
+
+        @flock_tool(name="utils.calculate_pi")
+        def compute_pi(): ...
+    """
+    registry = get_registry()
+
+    def decorator(inner_func: FuncType) -> FuncType:
+        if not callable(inner_func):
+            raise TypeError("@flock_tool can only decorate callables.")
+        # Let registry handle default name generation if None
+        registry.register_callable(inner_func, name=name)
+        return inner_func
+
+    if func is None:
+        # Called as @flock_tool(name="...")
+        return decorator
+    else:
+        # Called as @flock_tool
+        return decorator(func)
+
+
+# Alias for clarity
+flock_callable = flock_tool
+
+
+# --- Type Registration Decorator ---
+
+@overload
+def flock_type(cls: ClassType) -> ClassType: ...
+
+
+@overload
+def flock_type(
+    *, name: str | None = None
+) -> Callable[[ClassType], ClassType]: ...
+
+
+def flock_type(cls: ClassType | None = None, *, name: str | None = None) -> Any:
+    """Decorator to register a Type (Pydantic Model, Dataclass) used in signatures.
+
+    Usage:
+        @flock_type
+        class MyDataModel(BaseModel): ...
+
+        @flock_type(name="UserInput")
+        @dataclass
+        class UserQuery: ...
+    """
+    registry = get_registry()
+
+    def decorator(inner_cls: ClassType) -> ClassType:
+        if not inspect.isclass(inner_cls):
+            raise TypeError("@flock_type can only decorate classes.")
+        type_name = name or inner_cls.__name__
+        registry.register_type(inner_cls, name=type_name)
+        return inner_cls
+
+    if cls is None:
+        # Called as @flock_type(name="...")
+        return decorator
+    else:
+        # Called as @flock_type
+        return decorator(cls)

flock/core/registry/registry_hub.py

@@ -0,0 +1,205 @@
+# src/flock/core/registry/registry_hub.py
+"""Main registry hub using composition pattern with thread safety."""
+
+import threading
+from typing import TYPE_CHECKING, Any, TypeVar
+
+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 FlockMCPServer
+
+logger = get_logger("registry.hub")
+
+T = TypeVar("T")
+ConfigType = TypeVar("ConfigType")
+ClassType = TypeVar("ClassType")
+
+
+class RegistryHub:
+    """Thread-safe registry hub using composition pattern.
+    
+    Main coordinator for all registry types following the successful
+    pattern from Flock and FlockAgent refactoring.
+    """
+
+    def __init__(self):
+        self._lock = threading.RLock()
+        logger.debug("RegistryHub initialized with thread safety")
+
+    # --- Lazy-loaded composition helpers (following Flock pattern) ---
+
+    @property
+    def components(self):
+        """Component class registry helper."""
+        if not hasattr(self, '_components_helper'):
+            from flock.core.registry.component_registry import ComponentRegistry
+            self._components_helper = ComponentRegistry(self._lock)
+        return self._components_helper
+
+    @property
+    def callables(self):
+        """Callable registry helper."""
+        if not hasattr(self, '_callables_helper'):
+            from flock.core.registry.callable_registry import CallableRegistry
+            self._callables_helper = CallableRegistry(self._lock)
+        return self._callables_helper
+
+    @property
+    def agents(self):
+        """Agent registry helper."""
+        if not hasattr(self, '_agents_helper'):
+            from flock.core.registry.agent_registry import AgentRegistry
+            self._agents_helper = AgentRegistry(self._lock)
+        return self._agents_helper
+
+    @property
+    def servers(self):
+        """Server registry helper."""
+        if not hasattr(self, '_servers_helper'):
+            from flock.core.registry.server_registry import ServerRegistry
+            self._servers_helper = ServerRegistry(self._lock)
+        return self._servers_helper
+
+    @property
+    def types(self):
+        """Type registry helper."""
+        if not hasattr(self, '_types_helper'):
+            from flock.core.registry.type_registry import TypeRegistry
+            self._types_helper = TypeRegistry(self._lock)
+        return self._types_helper
+
+    @property
+    def config_mapping(self):
+        """Config mapping helper."""
+        if not hasattr(self, '_config_mapping_helper'):
+            from flock.core.registry.config_mapping import ConfigMapping
+            self._config_mapping_helper = ConfigMapping(self._lock)
+        return self._config_mapping_helper
+
+    @property
+    def discovery(self):
+        """Component discovery helper."""
+        if not hasattr(self, '_discovery_helper'):
+            from flock.core.registry.component_discovery import (
+                ComponentDiscovery,
+            )
+            self._discovery_helper = ComponentDiscovery(self)
+        return self._discovery_helper
+
+    # --- High-level registry operations (delegate to helpers) ---
+
+    def register_agent(self, agent: "FlockAgent", *, force: bool = False) -> None:
+        """Register a FlockAgent instance."""
+        self.agents.register_agent(agent, force=force)
+
+    def get_agent(self, name: str) -> "FlockAgent | None":
+        """Get a registered FlockAgent instance by name."""
+        return self.agents.get_agent(name)
+
+    def get_all_agent_names(self) -> list[str]:
+        """Get all registered agent names."""
+        return self.agents.get_all_agent_names()
+
+    def register_server(self, server: "FlockMCPServer") -> None:
+        """Register a FlockMCPServer instance."""
+        self.servers.register_server(server)
+
+    def get_server(self, name: str) -> "FlockMCPServer | None":
+        """Get a registered FlockMCPServer instance by name."""
+        return self.servers.get_server(name)
+
+    def get_all_server_names(self) -> list[str]:
+        """Get all registered server names."""
+        return self.servers.get_all_server_names()
+
+    def register_callable(self, func: "Callable", name: str | None = None) -> str | None:
+        """Register a callable function/method."""
+        return self.callables.register_callable(func, name)
+
+    def get_callable(self, name_or_path: str) -> "Callable":
+        """Get a registered callable by name or path."""
+        return self.callables.get_callable(name_or_path)
+
+    def get_callable_path_string(self, func: "Callable") -> str | None:
+        """Get the path string for a callable."""
+        return self.callables.get_callable_path_string(func)
+
+    def register_type(self, type_obj: type, name: str | None = None) -> str | None:
+        """Register a type (Pydantic Model, Dataclass, etc.)."""
+        return self.types.register_type(type_obj, name)
+
+    def get_type(self, type_name: str) -> type:
+        """Get a registered type by name."""
+        return self.types.get_type(type_name)
+
+    def register_component(self, component_class: type, name: str | None = None) -> str | None:
+        """Register a component class."""
+        return self.components.register_component(component_class, name)
+
+    def get_component(self, type_name: str) -> type:
+        """Get a registered component class by name."""
+        return self.components.get_component(type_name)
+
+    def get_component_type_name(self, component_class: type) -> str | None:
+        """Get the type name for a component class."""
+        return self.components.get_component_type_name(component_class)
+
+    def register_config_component_pair(self, config_cls: type[ConfigType], component_cls: type[ClassType]) -> None:
+        """Register a config-to-component mapping."""
+        self.config_mapping.register_config_component_pair(config_cls, component_cls)
+
+    def get_component_class_for_config(self, config_cls: type[ConfigType]) -> type[ClassType] | None:
+        """Get the component class for a config class."""
+        return self.config_mapping.get_component_class_for_config(config_cls)
+
+    def discover_and_register_components(self) -> None:
+        """Auto-discover and register components from known packages."""
+        self.discovery.discover_and_register_components()
+
+    def register_module_components(self, module_or_path: Any) -> None:
+        """Register components from a specific module."""
+        self.discovery.register_module_components(module_or_path)
+
+    # --- Utility methods ---
+
+    def clear_all(self) -> None:
+        """Clear all registries (useful for testing)."""
+        with self._lock:
+            if hasattr(self, '_agents_helper'):
+                self.agents.clear()
+            if hasattr(self, '_servers_helper'):
+                self.servers.clear()
+            if hasattr(self, '_callables_helper'):
+                self.callables.clear()
+            if hasattr(self, '_types_helper'):
+                self.types.clear()
+            if hasattr(self, '_components_helper'):
+                self.components.clear()
+            if hasattr(self, '_config_mapping_helper'):
+                self.config_mapping.clear_config_mappings()
+            logger.debug("Cleared all registries")
+
+    def get_registry_summary(self) -> dict[str, int]:
+        """Get a summary of all registries."""
+        with self._lock:
+            return {
+                "agents": len(self.agents.get_all_agents()),
+                "servers": len(self.servers.get_all_servers()),
+                "callables": len(self.callables.get_all_callables()),
+                "types": len(self.types.get_all_types()),
+                "components": len(self.components.get_all_components()),
+                "config_mappings": len(self.config_mapping.get_all_config_mappings()),
+            }
+
+
+# --- Global singleton instance ---
+_default_registry_hub = RegistryHub()
+
+
+def get_registry() -> RegistryHub:
+    """Get the default thread-safe registry hub instance."""
+    return _default_registry_hub

flock/core/registry/server_registry.py

@@ -0,0 +1,57 @@
+# src/flock/core/registry/server_registry.py
+"""MCP Server registration and lookup functionality."""
+
+import threading
+from typing import TYPE_CHECKING
+
+from flock.core.logging.logging import get_logger
+
+if TYPE_CHECKING:
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
+
+logger = get_logger("registry.servers")
+
+
+class ServerRegistry:
+    """Manages FlockMCPServerBase registration and lookup with thread safety."""
+
+    def __init__(self, lock: threading.RLock):
+        self._lock = lock
+        self._servers: dict[str, FlockMCPServer] = {}
+
+    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.")
+            return
+
+        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) -> "FlockMCPServer | None":
+        """Retrieve a registered FlockMCPServer instance by name."""
+        with self._lock:
+            server = self._servers.get(name)
+            if not server:
+                logger.warning(f"Server '{name}' not found in registry.")
+            return server
+
+    def get_all_server_names(self) -> list[str]:
+        """Return a list of names for all registered servers."""
+        with self._lock:
+            return list(self._servers.keys())
+
+    def get_all_servers(self) -> dict[str, "FlockMCPServer"]:
+        """Get all registered servers."""
+        with self._lock:
+            return self._servers.copy()
+
+    def clear(self) -> None:
+        """Clear all registered servers."""
+        with self._lock:
+            self._servers.clear()
+            logger.debug("Cleared all registered servers")

flock/core/registry/type_registry.py

@@ -0,0 +1,86 @@
+# src/flock/core/registry/type_registry.py
+"""Type registration and lookup functionality."""
+
+import threading
+from typing import Any, Literal, Mapping, Optional, Sequence, TypeVar, Union
+
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("registry.types")
+
+
+class TypeRegistry:
+    """Manages type registration and lookup with thread safety."""
+
+    def __init__(self, lock: threading.RLock):
+        self._lock = lock
+        self._types: dict[str, type] = {}
+        self._register_core_types()
+
+    def register_type(self, type_obj: type, name: str | None = None) -> str | None:
+        """Register a class/type (Pydantic, Dataclass, etc.) used in signatures."""
+        type_name = name or type_obj.__name__
+        if not type_name:
+            logger.error(f"Could not determine name for type: {type_obj}")
+            return None
+
+        with self._lock:
+            if type_name in self._types and self._types[type_name] != type_obj:
+                logger.warning(f"Type '{type_name}' already registered. Overwriting.")
+            
+            self._types[type_name] = type_obj
+            logger.debug(f"Registered type: {type_name}")
+            return type_name
+
+    def get_type(self, type_name: str) -> type:
+        """Retrieve a registered type by its name."""
+        with self._lock:
+            if type_name in self._types:
+                return self._types[type_name]
+            
+            # Consider adding dynamic import attempts for types if needed,
+            # but explicit registration is generally safer for types.
+            logger.warning(f"Type '{type_name}' not found in registry. Will attempt to build it from builtins.")
+            raise KeyError(f"Type '{type_name}' not found. Ensure it is registered.")
+
+    def get_all_types(self) -> dict[str, type]:
+        """Get all registered types."""
+        with self._lock:
+            return self._types.copy()
+
+    def clear(self) -> None:
+        """Clear all registered types (except core types)."""
+        with self._lock:
+            # Save core types
+            core_types = {
+                name: type_obj for name, type_obj in self._types.items()
+                if type_obj in [str, int, float, bool, list, dict, tuple, set, Any, Mapping, Sequence, TypeVar, Literal, Optional, Union]
+            }
+            self._types.clear()
+            self._types.update(core_types)
+            logger.debug("Cleared all registered types (keeping core types)")
+
+    def _register_core_types(self):
+        """Register common built-in and typing types."""
+        core_types = [
+            str,
+            int, 
+            float,
+            bool,
+            list,
+            dict,
+            tuple,
+            set,
+            Any,
+            Mapping,
+            Sequence,
+            TypeVar,
+            Literal,
+            Optional,
+            Union,  # Common typing generics
+        ]
+        for t in core_types:
+            try:
+                self.register_type(t)
+            except Exception as e:
+                logger.error(f"Failed to auto-register core type {t}: {e}")