flock-core 0.3.23__py3-none-any.whl → 0.3.31__py3-none-any.whl

flock/core/flock_registry.py

@@ -0,0 +1,532 @@
+# src/flock/core/flock_registry.py
+"""Centralized registry for managing Agents, Callables, Types, and Component Classes
+within the Flock framework to support dynamic lookup and serialization.
+"""
+
+from __future__ import annotations  # Add this at the very top
+
+import importlib
+import inspect
+import sys
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import is_dataclass
+from typing import (  # Add TYPE_CHECKING
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from flock.core.flock_agent import (
+        FlockAgent,  # Import only for type checking
+    )
+    from flock.core.flock_evaluator import FlockEvaluator
+    from flock.core.flock_module import FlockModule
+    from flock.core.flock_router import FlockRouter
+
+    COMPONENT_BASE_TYPES = (FlockModule, FlockEvaluator, FlockRouter)
+    IS_COMPONENT_CHECK_ENABLED = True
+else:
+    # Define dummy types or skip check if not type checking
+    FlockAgent = Any  # Or define a dummy class
+    COMPONENT_BASE_TYPES = ()
+    IS_COMPONENT_CHECK_ENABLED = False
+
+# Fallback if core types aren't available during setup
+
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("registry")
+T = TypeVar("T")
+ClassType = TypeVar("ClassType", bound=type)
+FuncType = TypeVar("FuncType", bound=Callable)
+
+
+class FlockRegistry:
+    """Singleton registry for Agents, Callables (functions/methods).
+
+    Types (Pydantic/Dataclasses used in signatures), and Component Classes
+    (Modules, Evaluators, Routers).
+    """
+
+    _instance = None
+
+    _agents: dict[str, FlockAgent]
+    _callables: dict[str, Callable]
+    _types: dict[str, type]
+    _components: dict[str, type]  # For Module, Evaluator, Router classes
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialize()
+            logger.info("FlockRegistry instance created.")
+        return cls._instance
+
+    def _initialize(self):
+        """Initialize the internal dictionaries."""
+        self._agents = {}
+        self._callables = {}
+        self._types = {}
+        self._components = {}
+        logger.debug("FlockRegistry initialized internal stores.")
+        # Auto-register core Python types
+        self._register_core_types()
+
+    def _register_core_types(self):
+        """Registers 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}")
+
+    # --- Path String Generation ---
+    @staticmethod
+    def _get_path_string(obj: Callable | type) -> str | None:
+        """Generates a unique path string 'module.ClassName' or 'module.function_name'."""
+        try:
+            module = obj.__module__
+            name = obj.__name__
+            if module == "builtins":
+                return name
+            # Check if it's nested (basic check, might not cover all edge cases)
+            if "." in name and hasattr(sys.modules[module], name.split(".")[0]):
+                # Likely a nested class/method - serialization might need custom handling or pickle
+                logger.warning(
+                    f"Object {name} appears nested in {module}. Path string might be ambiguous."
+                )
+            return f"{module}.{name}"
+        except AttributeError:
+            logger.warning(f"Could not determine module/name for object: {obj}")
+            return None
+
+    # --- Agent Registration ---
+    def register_agent(self, agent: FlockAgent) -> None:
+        """Registers a FlockAgent instance by its name."""
+        if not hasattr(agent, "name") or not agent.name:
+            logger.error(
+                "Attempted to register an agent without a valid 'name' attribute."
+            )
+            return
+        if agent.name in self._agents and self._agents[agent.name] != agent:
+            logger.warning(
+                f"Agent '{agent.name}' already registered. Overwriting."
+            )
+        self._agents[agent.name] = agent
+        logger.debug(f"Registered agent: {agent.name}")
+
+    def get_agent(self, name: str) -> FlockAgent | None:
+        """Retrieves a registered FlockAgent instance by name."""
+        agent = self._agents.get(name)
+        if not agent:
+            logger.warning(f"Agent '{name}' not found in registry.")
+        return agent
+
+    def get_all_agent_names(self) -> list[str]:
+        """Returns a list of names of all registered agents."""
+        return list(self._agents.keys())
+
+    # --- Callable Registration ---
+    def register_callable(
+        self, func: Callable, name: str | None = None
+    ) -> str | None:
+        """Registers a callable (function/method). Returns its path string identifier."""
+        path_str = name or self._get_path_string(func)
+        if path_str:
+            if (
+                path_str in self._callables
+                and self._callables[path_str] != func
+            ):
+                logger.warning(
+                    f"Callable '{path_str}' already registered. Overwriting."
+                )
+            self._callables[path_str] = func
+            logger.debug(f"Registered callable: {path_str}")
+            return path_str
+        return None
+
+    def get_callable(self, path_str: str) -> Callable:
+        """Retrieves a callable by its path string, attempting dynamic import if not found."""
+        if path_str in self._callables:
+            return self._callables[path_str]
+
+        logger.debug(
+            f"Callable '{path_str}' not in registry, attempting dynamic import."
+        )
+        try:
+            if "." not in path_str:  # Built-ins
+                builtins_module = importlib.import_module("builtins")
+                if hasattr(builtins_module, path_str):
+                    func = getattr(builtins_module, path_str)
+                    if callable(func):
+                        self.register_callable(func, path_str)  # Cache it
+                        return func
+                raise KeyError(f"Built-in callable '{path_str}' not found.")
+
+            module_name, func_name = path_str.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            func = getattr(module, func_name)
+            if callable(func):
+                self.register_callable(
+                    func, path_str
+                )  # Cache dynamically imported
+                return func
+            else:
+                raise TypeError(
+                    f"Dynamically imported object '{path_str}' is not callable."
+                )
+        except (ImportError, AttributeError, KeyError, TypeError) as e:
+            logger.error(
+                f"Failed to dynamically load/find callable '{path_str}': {e}"
+            )
+            raise KeyError(
+                f"Callable '{path_str}' not found or failed to load: {e}"
+            ) from e
+
+    def get_callable_path_string(self, func: Callable) -> str | None:
+        """Gets the path string for a callable, registering it if necessary."""
+        for path_str, registered_func in self._callables.items():
+            if func == registered_func:
+                return path_str
+        # If not found by identity, generate path, register, and return
+        return self.register_callable(func)
+
+    # --- Type Registration ---
+    def register_type(
+        self, type_obj: type, name: str | None = None
+    ) -> str | None:
+        """Registers a class/type (Pydantic, Dataclass, etc.) used in signatures."""
+        type_name = name or type_obj.__name__
+        if type_name:
+            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
+        return None
+
+    def get_type(self, type_name: str) -> type:
+        """Retrieves a registered type by its name."""
+        if type_name in self._types:
+            return self._types[type_name]
+        else:
+            # Consider adding dynamic import attempts for types if needed,
+            # but explicit registration is generally safer for types.
+            logger.error(f"Type '{type_name}' not found in registry.")
+            raise KeyError(
+                f"Type '{type_name}' not found. Ensure it is registered."
+            )
+
+    # --- Component Class Registration ---
+    def register_component(
+        self, component_class: type, name: str | None = None
+    ) -> str | None:
+        """Registers a component class (Module, Evaluator, Router)."""
+        type_name = name or component_class.__name__
+        if type_name:
+            # Optional: Add check if it's a subclass of expected bases
+            # if COMPONENT_BASE_TYPES and not issubclass(component_class, COMPONENT_BASE_TYPES):
+            #     logger.warning(f"Registering class '{type_name}' which is not a standard Flock component type.")
+            if (
+                type_name in self._components
+                and self._components[type_name] != component_class
+            ):
+                logger.warning(
+                    f"Component class '{type_name}' already registered. Overwriting."
+                )
+            self._components[type_name] = component_class
+            logger.debug(f"Registered component class: {type_name}")
+            return type_name
+        return None
+
+    def get_component(self, type_name: str) -> type:
+        """Retrieves a component class by its type name."""
+        if type_name in self._components:
+            return self._components[type_name]
+        else:
+            # Dynamic import attempts similar to get_callable could be added here if desired,
+            # targeting likely module locations based on type_name conventions.
+            logger.error(
+                f"Component class '{type_name}' not found in registry."
+            )
+            raise KeyError(
+                f"Component class '{type_name}' not found. Ensure it is registered."
+            )
+
+    def get_component_type_name(self, component_class: type) -> str | None:
+        """Gets the type name for a component class, registering it if necessary."""
+        for type_name, registered_class in self._components.items():
+            if component_class == registered_class:
+                return type_name
+        # If not found, register using class name and return
+        return self.register_component(component_class)
+
+    # --- Auto-Registration ---
+    def register_module_components(self, module_or_path: Any) -> None:
+        """Scans a module (object or path string) and automatically registers.
+
+        - Functions as callables.
+        - Pydantic Models and Dataclasses as types.
+        - Subclasses of FlockModule, FlockEvaluator, FlockRouter as components.
+        """
+        try:
+            if isinstance(module_or_path, str):
+                module = importlib.import_module(module_or_path)
+            elif inspect.ismodule(module_or_path):
+                module = module_or_path
+            else:
+                logger.error(
+                    f"Invalid input for auto-registration: {module_or_path}. Must be module object or path string."
+                )
+                return
+
+            logger.info(
+                f"Auto-registering components from module: {module.__name__}"
+            )
+            registered_count = {"callable": 0, "type": 0, "component": 0}
+
+            for name, obj in inspect.getmembers(module):
+                if name.startswith("_"):
+                    continue  # Skip private/internal
+
+                # Register Functions as Callables
+                if (
+                    inspect.isfunction(obj)
+                    and obj.__module__ == module.__name__
+                ):
+                    if self.register_callable(obj):
+                        registered_count["callable"] += 1
+
+                # Register Classes (Types and Components)
+                elif inspect.isclass(obj) and obj.__module__ == module.__name__:
+                    is_component = False
+                    # Register as Component if subclass of base types
+                    if (
+                        COMPONENT_BASE_TYPES
+                        and issubclass(obj, COMPONENT_BASE_TYPES)
+                        and self.register_component(obj)
+                    ):
+                        registered_count["component"] += 1
+                        is_component = True  # Mark as component
+
+                    # Register as Type if Pydantic Model or Dataclass
+                    # A component can also be a type used in signatures
+                    base_model_or_dataclass = isinstance(obj, type) and (
+                        issubclass(obj, BaseModel) or is_dataclass(obj)
+                    )
+                    if (
+                        base_model_or_dataclass
+                        and self.register_type(obj)
+                        and not is_component
+                    ):
+                        # Only increment type count if it wasn't already counted as component
+                        registered_count["type"] += 1
+
+            logger.info(
+                f"Auto-registration summary for {module.__name__}: "
+                f"{registered_count['callable']} callables, "
+                f"{registered_count['type']} types, "
+                f"{registered_count['component']} components."
+            )
+
+        except Exception as e:
+            logger.error(
+                f"Error during auto-registration for {module_or_path}: {e}",
+                exc_info=True,
+            )
+
+
+# --- Initialize Singleton ---
+_registry_instance = FlockRegistry()
+
+
+# --- Convenience Access ---
+# Provide a function to easily get the singleton instance
+def get_registry() -> FlockRegistry:
+    """Returns the singleton FlockRegistry instance."""
+    return _registry_instance
+
+
+# Type hinting for decorators to preserve signature
+@overload
+def flock_component(cls: ClassType) -> ClassType: ...
+@overload
+def flock_component(
+    *, name: str | None = None
+) -> Callable[[ClassType], ClassType]: ...
+
+
+def flock_component(
+    cls: ClassType | None = None, *, name: str | None = None
+) -> Any:
+    """Decorator to register a Flock Component class (Module, Evaluator, Router).
+
+    Usage:
+        @flock_component
+        class MyModule(FlockModule): ...
+
+        @flock_component(name="CustomRouterAlias")
+        class MyRouter(FlockRouter): ...
+    """
+    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)
+        return inner_cls
+
+    if cls is None:
+        # Called as @flock_component(name="...")
+        return decorator
+    else:
+        # Called as @flock_component
+        return decorator(cls)
+
+
+# Type hinting for decorators
+@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 if desired
+# flock_callable = flock_tool
+
+
+@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)
+
+
+# --- Auto-register known core components and tools ---
+def _auto_register_by_path():
+    components_to_register = [
+        (
+            "flock.evaluators.declarative.declarative_evaluator",
+            "DeclarativeEvaluator",
+        ),
+        ("flock.evaluators.memory.memory_evaluator", "MemoryEvaluator"),
+        ("flock.modules.output.output_module", "OutputModule"),
+        ("flock.modules.performance.metrics_module", "MetricsModule"),
+        ("flock.modules.memory.memory_module", "MemoryModule"),
+        # ("flock.modules.hierarchical.module", "HierarchicalMemoryModule"), # Uncomment if exists
+        ("flock.routers.default.default_router", "DefaultRouter"),
+        ("flock.routers.llm.llm_router", "LLMRouter"),
+        ("flock.routers.agent.agent_router", "AgentRouter"),
+    ]
+    for module_path, class_name in components_to_register:
+        try:
+            module = importlib.import_module(module_path)
+            component_class = getattr(module, class_name)
+            _registry_instance.register_component(component_class)
+        except (ImportError, AttributeError) as e:
+            logger.warning(f"{class_name} not found for auto-registration: {e}")
+
+    # Auto-register standard tools by scanning modules
+    tool_modules = [
+        "flock.core.tools.basic_tools",
+        "flock.core.tools.azure_tools",
+        "flock.core.tools.dev_tools.github",
+        "flock.core.tools.llm_tools",
+        "flock.core.tools.markdown_tools",
+    ]
+    for module_path in tool_modules:
+        try:
+            _registry_instance.register_module_components(module_path)
+        except ImportError as e:
+            logger.warning(
+                f"Could not auto-register tools from {module_path}: {e}"
+            )
+
+
+# Bootstrapping the registry
+_auto_register_by_path()

flock/core/logging/logging.py

@@ -46,50 +46,110 @@ def get_current_trace_id() -> str:
     return "no-trace"
 
 
-# ---------------------------------------------------------------------
-# 2. A color map for different logger names
-#    You can add or change entries as you like.
-# ---------------------------------------------------------------------
 COLOR_MAP = {
-    "flock": "magenta",
-    "interpreter": "cyan",
-    "memory": "yellow",
-    "activities": "blue",
+    # Core & Orchestration
+    "flock": "magenta",  # Color only
+    "agent": "blue",  # Color only
+    "workflow": "cyan",  # Color only
+    "activities": "cyan",
     "context": "green",
-    "registry": "white",
+    # Components & Mechanisms
+    "registry": "yellow",  # Color only
+    "serialization": "yellow",
+    "serialization.utils": "light-yellow",
+    "evaluator": "light-blue",
+    "module": "light-green",
+    "router": "light-magenta",
+    "mixin.dspy": "yellow",
+    # Specific Modules (Examples)
+    "memory": "yellow",
+    "module.output": "green",
+    "module.metrics": "blue",
+    "module.zep": "red",
+    "module.hierarchical": "light-green",
+    # Tools & Execution
     "tools": "light-black",
-    "agent": "light-magenta",
+    "interpreter": "light-yellow",
+    # API Components
+    "api": "white",  # Color only
+    "api.main": "white",
+    "api.endpoints": "light-black",
+    "api.run_store": "light-black",
+    "api.ui": "light-blue",  # Color only
+    "api.ui.routes": "light-blue",
+    "api.ui.utils": "cyan",
+    # Default/Unknown
+    "unknown": "light-black",
 }
 
 LOGGERS = [
+    "flock",  # Core Flock orchestration
+    "agent",  # General agent operations
+    "context",  # Context management
+    "registry",  # Unified registry operations (new)
+    "serialization",  # General serialization (new - can be base for others)
+    "serialization.utils",  # Serialization helpers (new, more specific)
+    "evaluator",  # Base evaluator category (new/optional)
+    "module",  # Base module category (new/optional)
+    "router",  # Base router category (new/optional)
+    "mixin.dspy",  # DSPy integration specifics (new)
+    "memory",  # Memory module specifics
+    "module.output",  # Output module specifics (example specific module)
+    "module.metrics",  # Metrics module specifics (example specific module)
+    "module.zep",  # Zep module specifics (example specific module)
+    "module.hierarchical",  # Hierarchical memory specifics (example specific module)
+    "interpreter",  # Code interpreter (if still used)
+    "activities",  # Temporal activities
+    "workflow",  # Temporal workflow logic
+    "tools",  # Tool execution/registration
+    "api",  # General API server (new)
+    "api.main",  # API main setup (new)
+    "api.endpoints",  # API endpoints (new)
+    "api.run_store",  # API run state management (new)
+    "api.ui",  # UI general (new)
+    "api.ui.routes",  # UI routes (new)
+    "api.ui.utils",  # UI utils (new)
+]
+
+BOLD_CATEGORIES = [
     "flock",
-    "interpreter",
-    "memory",
-    "activities",
-    "context",
-    "registry",
-    "tools",
     "agent",
+    "workflow",
+    "registry",
+    "api",
+    "api.ui",
 ]
 
 
 def color_for_category(category: str) -> str:
-    """Return the ANSI color code name for the given category."""
-    return COLOR_MAP.get(category, "magenta")  # fallback color
+    """Return the Rich markup color code name for the given category."""
+    # Handle potentially nested names like 'serialization.utils'
+    # Try exact match first, then go up the hierarchy
+    if category in COLOR_MAP:
+        return COLOR_MAP[category]
+    parts = category.split(".")
+    for i in range(len(parts) - 1, 0, -1):
+        parent_category = ".".join(parts[:i])
+        if parent_category in COLOR_MAP:
+            return COLOR_MAP[parent_category]
+    # Fallback to default 'unknown' color
+    return COLOR_MAP.get("unknown", "light-black")  # Final fallback
 
 
 def custom_format(record):
-    """A formatter that applies truncation to the entire formatted message."""
+    """A formatter that applies truncation and sequential styling tags."""
     t = record["time"].strftime("%Y-%m-%d %H:%M:%S")
     level_name = record["level"].name
     category = record["extra"].get("category", "unknown")
     trace_id = record["extra"].get("trace_id", "no-trace")
-    color = color_for_category(category)
+    color_tag = color_for_category(
+        category
+    )  # Get the color tag name (e.g., "yellow")
 
-    # Get the formatted message (already includes args)
     message = record["message"]
+    message = message.replace("{", "{{").replace("}", "}}")
 
-    # Apply truncation to the full formatted message
+    # MAX_LENGTH = 500 # Example value
     if len(message) > MAX_LENGTH:
         truncated_chars = len(message) - MAX_LENGTH
         message = (
@@ -97,10 +157,24 @@ def custom_format(record):
             + f"<yellow>...+({truncated_chars} chars)</yellow>"
         )
 
+    # Determine if category needs bolding (can refine this logic)
+    needs_bold = category in BOLD_CATEGORIES
+
+    # Apply tags sequentially
+    category_styled = f"[{category}]"  # Start with the plain category name
+    category_styled = (
+        f"<{color_tag}>{category_styled}</{color_tag}>"  # Wrap with color
+    )
+    if needs_bold:
+        category_styled = (
+            f"<bold>{category_styled}</bold>"  # Wrap with bold if needed
+        )
+
+    # Final format string using sequential tags for category
     return (
         f"<green>{t}</green> | <level>{level_name: <8}</level> | "
         f"<cyan>[trace_id: {trace_id}]</cyan> | "
-        f"<{color}>[{category}]</{color}> | {message}\n"
+        f"{category_styled} | {message}\n"  # Apply the sequentially styled category
     )