flock-core 0.4.0b22__py3-none-any.whl → 0.4.0b24__py3-none-any.whl

flock/core/flock.py

@@ -59,6 +59,10 @@ FlockRegistry = get_registry()  # Get the registry instance
 # Define TypeVar for generic class methods like from_dict
 T = TypeVar("T", bound="Flock")
 
+from rich.traceback import install
+
+install(show_locals=True)
+
 
 class Flock(BaseModel, Serializable):
     """Orchestrator for managing and executing agent systems.
@@ -350,9 +354,13 @@ class Flock(BaseModel, Serializable):
             try:
                 resolved_start_agent = self._agents.get(start_agent_name)
                 if not resolved_start_agent:
-                    resolved_start_agent = FlockRegistry.get_agent(start_agent_name)
+                    resolved_start_agent = FlockRegistry.get_agent(
+                        start_agent_name
+                    )
                     if not resolved_start_agent:
-                        raise ValueError(f"Start agent '{start_agent_name}' not found.")
+                        raise ValueError(
+                            f"Start agent '{start_agent_name}' not found."
+                        )
                     self.add_agent(resolved_start_agent)
 
                 run_context = context if context else FlockContext()
@@ -384,15 +392,20 @@ class Flock(BaseModel, Serializable):
 
                 # Execute workflow
                 if not self.enable_temporal:
-                    result = await run_local_workflow(run_context, box_result=False)
+                    result = await run_local_workflow(
+                        run_context, box_result=False
+                    )
                 else:
-                    result = await run_temporal_workflow(run_context, box_result=False)
+                    result = await run_temporal_workflow(
+                        run_context, box_result=False
+                    )
 
                 span.set_attribute("result.type", str(type(result)))
                 result_str = str(result)
                 span.set_attribute(
                     "result.preview",
-                    result_str[:1000] + ("..." if len(result_str) > 1000 else ""),
+                    result_str[:1000]
+                    + ("..." if len(result_str) > 1000 else ""),
                 )
 
                 if box_result:
@@ -400,13 +413,17 @@ class Flock(BaseModel, Serializable):
                         logger.debug("Boxing final result.")
                         return Box(result)
                     except ImportError:
-                        logger.warning("Box library not installed, returning raw dict.")
+                        logger.warning(
+                            "Box library not installed, returning raw dict."
+                        )
                         return result
                 else:
                     return result
 
             except Exception as e:
-                logger.error(f"Flock run '{self.name}' failed: {e}", exc_info=True)
+                logger.error(
+                    f"Flock run '{self.name}' failed: {e}", exc_info=True
+                )
                 span.record_exception(e)
                 span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
                 return {

flock/core/flock_agent.py

@@ -23,9 +23,9 @@ from rich.console import Console
 
 # Core Flock components (ensure these are importable)
 from flock.core.context.context import FlockContext
-from flock.core.flock_evaluator import FlockEvaluator
-from flock.core.flock_module import FlockModule
-from flock.core.flock_router import FlockRouter
+from flock.core.flock_evaluator import FlockEvaluator, FlockEvaluatorConfig
+from flock.core.flock_module import FlockModule, FlockModuleConfig
+from flock.core.flock_router import FlockRouter, FlockRouterConfig
 from flock.core.logging.logging import get_logger
 
 # Mixins and Serialization components
@@ -45,6 +45,15 @@ tracer = trace.get_tracer(__name__)
 T = TypeVar("T", bound="FlockAgent")
 
 
+SignatureType = (
+    str
+    | Callable[..., str]
+    | type[BaseModel]
+    | Callable[..., type[BaseModel]]
+    | None
+)
+
+
 # Make FlockAgent inherit from Serializable
 class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
     """Core, declarative base class for Flock agents, enabling serialization,
@@ -61,14 +70,14 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
         "",
         description="A human-readable description or a callable returning one.",
     )
-    input: str | Callable[..., str] | None = Field(
+    input: SignatureType = Field(
         None,
         description=(
             "Signature for input keys. Supports type hints (:) and descriptions (|). "
             "E.g., 'query: str | Search query, context: dict | Conversation context'. Can be a callable."
         ),
     )
-    output: str | Callable[..., str] | None = Field(
+    output: SignatureType = Field(
         None,
         description=(
             "Signature for output keys. Supports type hints (:) and descriptions (|). "
@@ -111,6 +120,43 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
         description="Runtime context associated with the flock execution.",
     )
 
+    def __init__(
+        self,
+        name: str,
+        model: str | None = None,
+        description: str | Callable[..., str] | None = "",
+        input: SignatureType = None,
+        output: SignatureType = None,
+        tools: list[Callable[..., Any]] | None = None,
+        evaluator: "FlockEvaluator | None" = None,
+        handoff_router: "FlockRouter | None" = None,
+        modules: dict[str, "FlockModule"] | None = None,  # Use dict for modules
+        write_to_file: bool = False,
+        wait_for_input: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            model=model,
+            description=description,
+            input=input,  # Store the raw input spec
+            output=output,  # Store the raw output spec
+            tools=tools,
+            write_to_file=write_to_file,
+            wait_for_input=wait_for_input,
+            evaluator=evaluator,
+            handoff_router=handoff_router,
+            modules=modules
+            if modules is not None
+            else {},  # Ensure modules is a dict
+            **kwargs,
+        )
+
+        if isinstance(self.input, type) and issubclass(self.input, BaseModel):
+            self._input_model = self.input
+        if isinstance(self.output, type) and issubclass(self.output, BaseModel):
+            self._output_model = self.output
+
     # --- Existing Methods (add_module, remove_module, etc.) ---
     # (Keep these methods as they were, adding type hints where useful)
     def add_module(self, module: FlockModule) -> None:
@@ -359,6 +405,114 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                 span.record_exception(temporal_error)
                 raise
 
+    def add_component(
+        self,
+        config_instance: FlockModuleConfig
+        | FlockRouterConfig
+        | FlockEvaluatorConfig,
+        component_name: str | None = None,
+    ) -> "FlockAgent":
+        """Adds or replaces a component (Evaluator, Router, Module) based on its configuration object.
+
+        Args:
+            config_instance: An instance of a config class inheriting from
+                             FlockModuleConfig, FlockRouterConfig, or FlockEvaluatorConfig.
+            component_name: Explicit name for the component (required for Modules if not in config).
+
+        Returns:
+            self for potential chaining.
+        """
+        from flock.core.flock_registry import get_registry
+
+        config_type = type(config_instance)
+        registry = get_registry()  # Get registry instance
+        logger.debug(
+            f"Attempting to add component via config: {config_type.__name__}"
+        )
+
+        # --- 1. Find Component Class using Registry Map ---
+        ComponentClass = registry.get_component_class_for_config(config_type)
+
+        if not ComponentClass:
+            logger.error(
+                f"No component class registered for config type {config_type.__name__}. Use @flock_component(config_class=...) on the component."
+            )
+            raise TypeError(
+                f"Cannot find component class for config {config_type.__name__}"
+            )
+
+        component_class_name = ComponentClass.__name__
+        logger.debug(
+            f"Found component class '{component_class_name}' mapped to config '{config_type.__name__}'"
+        )
+
+        # --- 2. Determine Assignment Target and Name (Same as before) ---
+        instance_name = component_name
+        attribute_name: str = ""
+
+        if issubclass(ComponentClass, FlockEvaluator):
+            attribute_name = "evaluator"
+            if not instance_name:
+                instance_name = getattr(
+                    config_instance, "name", component_class_name.lower()
+                )
+
+        elif issubclass(ComponentClass, FlockRouter):
+            attribute_name = "handoff_router"
+            if not instance_name:
+                instance_name = getattr(
+                    config_instance, "name", component_class_name.lower()
+                )
+
+        elif issubclass(ComponentClass, FlockModule):
+            attribute_name = "modules"
+            if not instance_name:
+                instance_name = getattr(
+                    config_instance, "name", component_class_name.lower()
+                )
+            if not instance_name:
+                raise ValueError(
+                    "Module name must be provided either in config or as component_name argument."
+                )
+            # Ensure config has name if module expects it
+            if hasattr(config_instance, "name") and not getattr(
+                config_instance, "name", None
+            ):
+                setattr(config_instance, "name", instance_name)
+
+        else:  # Should be caught by registry map logic ideally
+            raise TypeError(
+                f"Class '{component_class_name}' mapped from config is not a valid Flock component."
+            )
+
+        # --- 3. Instantiate the Component (Same as before) ---
+        try:
+            init_args = {"config": config_instance, "name": instance_name}
+
+            component_instance = ComponentClass(**init_args)
+        except Exception as e:
+            logger.error(
+                f"Failed to instantiate {ComponentClass.__name__} with config {config_type.__name__}: {e}",
+                exc_info=True,
+            )
+            raise RuntimeError(f"Component instantiation failed: {e}") from e
+
+        # --- 4. Assign to the Agent (Same as before) ---
+        if attribute_name == "modules":
+            if not isinstance(self.modules, dict):
+                self.modules = {}
+            self.modules[instance_name] = component_instance
+            logger.info(
+                f"Added/Updated module '{instance_name}' (type: {ComponentClass.__name__}) to agent '{self.name}'"
+            )
+        else:
+            setattr(self, attribute_name, component_instance)
+            logger.info(
+                f"Set {attribute_name} to {ComponentClass.__name__} (instance name: '{instance_name}') for agent '{self.name}'"
+            )
+
+        return self
+
     # resolve_callables remains useful for dynamic definitions
     def resolve_callables(self, context: FlockContext | None = None) -> None:
         """Resolves callable fields (description, input, output) using context."""
@@ -428,81 +582,71 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
             exclude_none=True,  # Exclude None values for cleaner output
         )
         logger.debug(f"Base agent data for '{self.name}': {list(data.keys())}")
+        serialized_modules = {}
+
+        def add_serialized_component(component: Any, field_name: str):
+            if component:
+                comp_type = type(component)
+                type_name = FlockRegistry.get_component_type_name(
+                    comp_type
+                )  # Get registered name
+                if type_name:
+                    try:
+                        serialized_component_data = serialize_item(component)
+
+                        if not isinstance(serialized_component_data, dict):
+                            logger.error(
+                                f"Serialization of component {type_name} for field '{field_name}' did not result in a dictionary. Got: {type(serialized_component_data)}"
+                            )
+                            serialized_modules[field_name] = {
+                                "type": type_name,
+                                "name": getattr(component, "name", "unknown"),
+                                "error": "serialization_failed_non_dict",
+                            }
+                        else:
+                            serialized_component_data["type"] = type_name
+                            serialized_modules[field_name] = (
+                                serialized_component_data
+                            )
+                            logger.debug(
+                                f"Successfully serialized component for field '{field_name}' (type: {type_name})"
+                            )
 
-        # --- Serialize Components using Registry Type Names ---
-        # Evaluator
-        if self.evaluator:
-            logger.debug(f"Serializing evaluator for agent '{self.name}'")
-            evaluator_type_name = FlockRegistry.get_component_type_name(
-                type(self.evaluator)
-            )
-            if evaluator_type_name:
-                # Recursively serialize the evaluator's dict representation
-                evaluator_dict = serialize_item(
-                    self.evaluator.model_dump(mode="json", exclude_none=True)
-                )
-                evaluator_dict["type"] = evaluator_type_name  # Add type marker
-                data["evaluator"] = evaluator_dict
-                logger.debug(
-                    f"Added evaluator of type '{evaluator_type_name}' to agent '{self.name}'"
-                )
-            else:
-                logger.warning(
-                    f"Could not get registered type name for evaluator {type(self.evaluator).__name__} in agent '{self.name}'. Skipping serialization."
-                )
-
-        # Router
-        if self.handoff_router:
-            logger.debug(f"Serializing router for agent '{self.name}'")
-            router_type_name = FlockRegistry.get_component_type_name(
-                type(self.handoff_router)
-            )
-            if router_type_name:
-                router_dict = serialize_item(
-                    self.handoff_router.model_dump(
-                        mode="json", exclude_none=True
-                    )
-                )
-                router_dict["type"] = router_type_name
-                data["handoff_router"] = router_dict
-                logger.debug(
-                    f"Added router of type '{router_type_name}' to agent '{self.name}'"
-                )
-            else:
-                logger.warning(
-                    f"Could not get registered type name for router {type(self.handoff_router).__name__} in agent '{self.name}'. Skipping serialization."
-                )
-
-        # Modules
-        if self.modules:
-            logger.debug(
-                f"Serializing {len(self.modules)} modules for agent '{self.name}'"
-            )
-            serialized_modules = {}
-            for name, module_instance in self.modules.items():
-                module_type_name = FlockRegistry.get_component_type_name(
-                    type(module_instance)
-                )
-                if module_type_name:
-                    module_dict = serialize_item(
-                        module_instance.model_dump(
-                            mode="json", exclude_none=True
+                    except Exception as e:
+                        logger.error(
+                            f"Failed to serialize component {type_name} for field '{field_name}': {e}",
+                            exc_info=True,
                         )
-                    )
-                    module_dict["type"] = module_type_name
-                    serialized_modules[name] = module_dict
-                    logger.debug(
-                        f"Added module '{name}' of type '{module_type_name}' to agent '{self.name}'"
-                    )
+                        serialized_modules[field_name] = {
+                            "type": type_name,
+                            "name": getattr(component, "name", "unknown"),
+                            "error": "serialization_failed",
+                        }
                 else:
                     logger.warning(
-                        f"Could not get registered type name for module {type(module_instance).__name__} ('{name}') in agent '{self.name}'. Skipping."
+                        f"Cannot serialize unregistered component {comp_type.__name__} for field '{field_name}'"
                     )
-            if serialized_modules:
-                data["modules"] = serialized_modules
-                logger.debug(
-                    f"Added {len(serialized_modules)} modules to agent '{self.name}'"
-                )
+
+        add_serialized_component(self.evaluator, "evaluator")
+        if serialized_modules:
+            data["evaluator"] = serialized_modules["evaluator"]
+            logger.debug(f"Added evaluator to agent '{self.name}'")
+
+        serialized_modules = {}
+        add_serialized_component(self.handoff_router, "handoff_router")
+        if serialized_modules:
+            data["handoff_router"] = serialized_modules["handoff_router"]
+            logger.debug(f"Added handoff_router to agent '{self.name}'")
+
+        serialized_modules = {}
+        for module in self.modules.values():
+            add_serialized_component(module, module.name)
+
+        if serialized_modules:
+            data["modules"] = serialized_modules
+            logger.debug(
+                f"Added {len(serialized_modules)} modules to agent '{self.name}'"
+            )
 
         # --- Serialize Tools (Callables) ---
         if self.tools:

flock/core/flock_factory.py

@@ -3,7 +3,7 @@
 from collections.abc import Callable
 from typing import Any
 
-from flock.core.flock_agent import FlockAgent
+from flock.core.flock_agent import FlockAgent, SignatureType
 from flock.core.logging.formatters.themes import OutputTheme
 from flock.evaluators.declarative.declarative_evaluator import (
     DeclarativeEvaluator,
@@ -24,15 +24,15 @@ class FlockFactory:
         name: str,
         description: str | Callable[..., str] | None = None,
         model: str | Callable[..., str] | None = None,
-        input: str | Callable[..., str] | None = None,
-        output: str | Callable[..., str] | None = None,
+        input: SignatureType = None,
+        output: SignatureType = None,
         tools: list[Callable[..., Any] | Any] | None = None,
         use_cache: bool = True,
         enable_rich_tables: bool = False,
         output_theme: OutputTheme = OutputTheme.abernathy,
         wait_for_input: bool = False,
         temperature: float = 0.0,
-        max_tokens: int = 4096,
+        max_tokens: int = 8192,
         alert_latency_threshold_ms: int = 30000,
         no_output: bool = False,
         print_context: bool = False,
@@ -42,10 +42,12 @@ class FlockFactory:
     ) -> FlockAgent:
         """Creates a default FlockAgent.
 
-        The default agent includes a declarative evaluator with the following modules:
+        The default agent includes the following modules:
+        - DeclarativeEvaluator
         - OutputModule
+        - MetricsModule
 
-        It also includes often needed configurations like cache usage, rich tables, and output theme.
+        It also includes direct acces to the most important configurations.
         """
         eval_config = DeclarativeEvaluatorConfig(
             model=model,