flock-core 0.5.10__py3-none-any.whl → 0.5.20__py3-none-any.whl

flock/orchestrator/event_emitter.py

@@ -0,0 +1,167 @@
+"""Event emission for real-time dashboard updates.
+
+Phase 5A: Extracted from orchestrator.py to reduce coupling to dashboard.
+
+This module handles WebSocket event emission for dashboard visualization of
+batch and correlation logic operations. Separating this code reduces the
+orchestrator's dependency on dashboard-specific components.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.core.artifacts import Artifact
+    from flock.core.subscription import Subscription
+    from flock.orchestrator.batch_accumulator import BatchEngine
+    from flock.orchestrator.correlation_engine import CorrelationEngine
+
+
+class EventEmitter:
+    """Manages WebSocket event emission for dashboard updates.
+
+    This module is responsible for broadcasting real-time events about
+    batch accumulation and correlation group status to connected dashboard
+    clients via WebSocket.
+
+    Phase 5A: Extracted to reduce orchestrator coupling to dashboard.
+    """
+
+    def __init__(self, websocket_manager: Any | None = None):
+        """Initialize EventEmitter with WebSocket manager.
+
+        Args:
+            websocket_manager: WebSocket manager instance for broadcasting events.
+                If None, event emission is disabled (dashboard not active).
+        """
+        self._websocket_manager = websocket_manager
+
+    def set_websocket_manager(self, websocket_manager: Any | None) -> None:
+        """Update the WebSocket manager (called when dashboard is enabled).
+
+        Args:
+            websocket_manager: WebSocket manager instance for broadcasting
+        """
+        self._websocket_manager = websocket_manager
+
+    async def emit_correlation_updated(
+        self,
+        *,
+        correlation_engine: CorrelationEngine,
+        agent_name: str,
+        subscription_index: int,
+        artifact: Artifact,
+    ) -> None:
+        """Emit CorrelationGroupUpdatedEvent for real-time dashboard updates.
+
+        Called when an artifact is added to a correlation group that is not yet complete.
+
+        Args:
+            correlation_engine: CorrelationEngine instance with current state
+            agent_name: Name of the agent with the JoinSpec subscription
+            subscription_index: Index of the subscription in the agent's subscriptions list
+            artifact: The artifact that triggered this update
+        """
+        # Only emit if dashboard is enabled
+        if self._websocket_manager is None:
+            return
+
+        # Import _get_correlation_groups helper from dashboard service
+        from flock.dashboard.routes.helpers import _get_correlation_groups
+
+        # Get current correlation groups state from engine
+        groups = _get_correlation_groups(
+            correlation_engine, agent_name, subscription_index
+        )
+
+        if not groups:
+            return  # No groups to report (shouldn't happen, but defensive)
+
+        # Find the group that was just updated (match by last updated time or artifact ID)
+        # For now, we'll emit an event for the FIRST group that's still waiting
+        # In practice, the artifact we just added should be in one of these groups
+        for group_state in groups:
+            if not group_state["is_complete"]:
+                # Import CorrelationGroupUpdatedEvent
+                from flock.dashboard.events import CorrelationGroupUpdatedEvent
+
+                # Build and emit event
+                event = CorrelationGroupUpdatedEvent(
+                    agent_name=agent_name,
+                    subscription_index=subscription_index,
+                    correlation_key=group_state["correlation_key"],
+                    collected_types=group_state["collected_types"],
+                    required_types=group_state["required_types"],
+                    waiting_for=group_state["waiting_for"],
+                    elapsed_seconds=group_state["elapsed_seconds"],
+                    expires_in_seconds=group_state["expires_in_seconds"],
+                    expires_in_artifacts=group_state["expires_in_artifacts"],
+                    artifact_id=str(artifact.id),
+                    artifact_type=artifact.type,
+                    is_complete=group_state["is_complete"],
+                )
+
+                # Broadcast via WebSocket
+                await self._websocket_manager.broadcast(event)
+                break  # Only emit one event per artifact addition
+
+    async def emit_batch_item_added(
+        self,
+        *,
+        batch_engine: BatchEngine,
+        agent_name: str,
+        subscription_index: int,
+        subscription: Subscription,
+        artifact: Artifact,
+    ) -> None:
+        """Emit BatchItemAddedEvent for real-time dashboard updates.
+
+        Called when an artifact is added to a batch that hasn't reached flush threshold.
+
+        Args:
+            batch_engine: BatchEngine instance with current state
+            agent_name: Name of the agent with the BatchSpec subscription
+            subscription_index: Index of the subscription in the agent's subscriptions list
+            subscription: The subscription with BatchSpec configuration
+            artifact: The artifact that triggered this update
+        """
+        # Only emit if dashboard is enabled
+        if self._websocket_manager is None:
+            return
+
+        # Import _get_batch_state helper from dashboard service
+        from flock.dashboard.routes.helpers import _get_batch_state
+
+        # Get current batch state from engine
+        batch_state = _get_batch_state(
+            batch_engine, agent_name, subscription_index, subscription.batch
+        )
+
+        if not batch_state:
+            return  # No batch to report (shouldn't happen, but defensive)
+
+        # Import BatchItemAddedEvent
+        from flock.dashboard.events import BatchItemAddedEvent
+
+        # Build and emit event
+        event = BatchItemAddedEvent(
+            agent_name=agent_name,
+            subscription_index=subscription_index,
+            items_collected=batch_state["items_collected"],
+            items_target=batch_state.get("items_target"),
+            items_remaining=batch_state.get("items_remaining"),
+            elapsed_seconds=batch_state["elapsed_seconds"],
+            timeout_seconds=batch_state.get("timeout_seconds"),
+            timeout_remaining_seconds=batch_state.get("timeout_remaining_seconds"),
+            will_flush=batch_state["will_flush"],
+            artifact_id=str(artifact.id),
+            artifact_type=artifact.type,
+        )
+
+        # Broadcast via WebSocket
+        await self._websocket_manager.broadcast(event)
+
+
+__all__ = ["EventEmitter"]

flock/orchestrator/initialization.py

@@ -0,0 +1,184 @@
+"""Orchestrator initialization helper.
+
+Handles component setup and state initialization.
+Extracted from orchestrator.py to reduce __init__ complexity.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from typing import TYPE_CHECKING, Any
+
+from flock.core.store import InMemoryBlackboardStore
+from flock.orchestrator.artifact_collector import ArtifactCollector
+from flock.orchestrator.batch_accumulator import BatchEngine
+from flock.orchestrator.component_runner import ComponentRunner
+from flock.orchestrator.context_builder import ContextBuilder
+from flock.orchestrator.correlation_engine import CorrelationEngine
+from flock.orchestrator.event_emitter import EventEmitter
+from flock.orchestrator.lifecycle_manager import LifecycleManager
+from flock.orchestrator.mcp_manager import MCPManager
+from flock.orchestrator.tracing import TracingManager
+from flock.utils.cli_helper import init_console
+
+
+if TYPE_CHECKING:
+    from flock.components.orchestrator import OrchestratorComponent
+    from flock.core.store import BlackboardStore
+
+
+class OrchestratorInitializer:
+    """Handles orchestrator component initialization.
+
+    Centralizes the complex setup logic from Flock.__init__ into
+    a focused helper class.
+    """
+
+    @staticmethod
+    def initialize_components(
+        store: BlackboardStore | None,
+        context_provider: Any,
+        max_agent_iterations: int,
+        logger: logging.Logger,
+        model: str | None,
+    ) -> dict[str, Any]:
+        """Initialize all orchestrator components and state.
+
+        Args:
+            store: Blackboard storage backend (or None for default)
+            context_provider: Global context provider for agents
+            max_agent_iterations: Circuit breaker limit
+            logger: Logger instance
+            model: Default LLM model
+
+        Returns:
+            Dictionary of initialized components and state
+
+        Examples:
+            >>> logger = logging.getLogger(__name__)
+            >>> components = OrchestratorInitializer.initialize_components(
+            ...     store=None,
+            ...     context_provider=None,
+            ...     max_agent_iterations=1000,
+            ...     logger=logger,
+            ...     model="openai/gpt-4.1",
+            ... )
+            >>> orchestrator.store = components["store"]
+            >>> orchestrator._scheduler = components["scheduler"]
+        """
+        # Initialize console (with error handling for encoding issues)
+        try:
+            init_console(clear_screen=True, show_banner=True, model=model)
+        except (UnicodeEncodeError, UnicodeDecodeError):
+            # Skip banner on Windows consoles with encoding issues (e.g., tests, CI)
+            pass
+
+        # Basic state
+        resolved_store = store or InMemoryBlackboardStore()
+        agents: dict[str, Any] = {}
+        lock = asyncio.Lock()
+        metrics: dict[str, float] = {"artifacts_published": 0, "agent_runs": 0}
+        agent_iteration_count: dict[str, int] = {}
+
+        # Engines
+        artifact_collector = ArtifactCollector()
+        correlation_engine = CorrelationEngine()
+        batch_engine = BatchEngine()
+
+        # Phase 5A modules
+        context_builder = ContextBuilder(
+            store=resolved_store,
+            default_context_provider=context_provider,
+        )
+        event_emitter = EventEmitter(websocket_manager=None)
+        lifecycle_manager = LifecycleManager(
+            correlation_engine=correlation_engine,
+            batch_engine=batch_engine,
+            cleanup_interval=0.1,
+        )
+
+        # Phase 3 modules
+        mcp_manager_instance = MCPManager()
+        tracing_manager = TracingManager()
+
+        # Auto-workflow tracing feature flag
+        auto_workflow_enabled = os.getenv(
+            "FLOCK_AUTO_WORKFLOW_TRACE", "false"
+        ).lower() in {
+            "true",
+            "1",
+            "yes",
+            "on",
+        }
+
+        return {
+            # Basic state
+            "store": resolved_store,
+            "agents": agents,
+            "lock": lock,
+            "metrics": metrics,
+            "agent_iteration_count": agent_iteration_count,
+            # Engines
+            "artifact_collector": artifact_collector,
+            "correlation_engine": correlation_engine,
+            "batch_engine": batch_engine,
+            # Phase 5A modules
+            "context_builder": context_builder,
+            "event_emitter": event_emitter,
+            "lifecycle_manager": lifecycle_manager,
+            # Phase 3 modules
+            "mcp_manager_instance": mcp_manager_instance,
+            "tracing_manager": tracing_manager,
+            # Feature flags
+            "auto_workflow_enabled": auto_workflow_enabled,
+            # Placeholders
+            "websocket_manager": None,
+        }
+
+    @staticmethod
+    def initialize_components_and_runner(
+        components_list: list[OrchestratorComponent],
+        max_agent_iterations: int,
+        logger: logging.Logger,
+    ) -> dict[str, Any]:
+        """Initialize built-in components and create ComponentRunner.
+
+        Args:
+            components_list: List to populate with components
+            max_agent_iterations: Circuit breaker limit
+            logger: Logger instance
+
+        Returns:
+            Dictionary with component_runner and updated components list
+
+        Examples:
+            >>> components = []
+            >>> result = OrchestratorInitializer.initialize_components_and_runner(
+            ...     components, max_agent_iterations=1000, logger=logger
+            ... )
+            >>> component_runner = result["component_runner"]
+        """
+        from flock.components.orchestrator import (
+            BuiltinCollectionComponent,
+            CircuitBreakerComponent,
+            DeduplicationComponent,
+        )
+
+        # Add built-in components
+        components_list.append(
+            CircuitBreakerComponent(max_iterations=max_agent_iterations)
+        )
+        components_list.append(DeduplicationComponent())
+        components_list.append(BuiltinCollectionComponent())
+
+        # Sort by priority
+        components_list.sort(key=lambda c: c.priority)
+
+        # Create ComponentRunner
+        component_runner = ComponentRunner(components_list, logger)
+
+        return {
+            "component_runner": component_runner,
+        }

flock/orchestrator/lifecycle_manager.py

@@ -0,0 +1,226 @@
+"""Lifecycle management for background tasks and cleanup.
+
+Phase 5A: Extracted from orchestrator.py to isolate background task coordination.
+
+This module handles background tasks for batch timeouts and correlation cleanup,
+reducing orchestrator complexity and centralizing async task management.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from asyncio import Task
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.orchestrator.batch_accumulator import BatchEngine
+    from flock.orchestrator.correlation_engine import CorrelationEngine
+
+
+class LifecycleManager:
+    """Manages background tasks for batch and correlation lifecycle.
+
+    This module centralizes all background task management for:
+    - Correlation group cleanup (time-based expiry)
+    - Batch timeout checking (timeout-based flushing)
+
+    Phase 5A: Extracted to reduce orchestrator complexity and improve testability.
+    """
+
+    def __init__(
+        self,
+        *,
+        correlation_engine: CorrelationEngine,
+        batch_engine: BatchEngine,
+        cleanup_interval: float = 0.1,
+    ):
+        """Initialize LifecycleManager with engines and intervals.
+
+        Args:
+            correlation_engine: Engine managing correlation groups
+            batch_engine: Engine managing batch accumulation
+            cleanup_interval: How often to check for expiry (seconds, default: 0.1)
+        """
+        self._correlation_engine = correlation_engine
+        self._batch_engine = batch_engine
+        self._cleanup_interval = cleanup_interval
+
+        # Background tasks
+        self._correlation_cleanup_task: Task[Any] | None = None
+        self._batch_timeout_task: Task[Any] | None = None
+
+        # Callback for batch timeout flushing (set by orchestrator)
+        self._batch_timeout_callback: Any | None = None
+
+        self._logger = logging.getLogger(__name__)
+
+    async def start_correlation_cleanup(self) -> None:
+        """Start background correlation cleanup loop if not already running.
+
+        This ensures expired correlation groups are periodically discarded.
+        Called when there are pending correlations during run_until_idle.
+        """
+        if (
+            self._correlation_cleanup_task is None
+            or self._correlation_cleanup_task.done()
+        ):
+            self._correlation_cleanup_task = asyncio.create_task(
+                self._correlation_cleanup_loop()
+            )
+
+    def set_batch_timeout_callback(self, callback: Any) -> None:
+        """Set the callback to invoke when batches timeout.
+
+        Args:
+            callback: Async function to call when timeout checking. Should handle
+                flushing expired batches and scheduling agent tasks.
+        """
+        self._batch_timeout_callback = callback
+
+    async def start_batch_timeout_checker(self) -> None:
+        """Start background batch timeout checker loop if not already running.
+
+        This ensures timeout-expired batches are periodically flushed.
+        Called when there are pending batches during run_until_idle.
+        """
+        if self._batch_timeout_task is None or self._batch_timeout_task.done():
+            self._batch_timeout_task = asyncio.create_task(
+                self._batch_timeout_checker_loop()
+            )
+
+    async def shutdown(self) -> None:
+        """Cancel and cleanup all background tasks.
+
+        Called during orchestrator shutdown to ensure clean resource cleanup.
+        """
+        # Cancel correlation cleanup task if running
+        if self._correlation_cleanup_task and not self._correlation_cleanup_task.done():
+            self._correlation_cleanup_task.cancel()
+            try:
+                await self._correlation_cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+        # Cancel batch timeout checker if running
+        if self._batch_timeout_task and not self._batch_timeout_task.done():
+            self._batch_timeout_task.cancel()
+            try:
+                await self._batch_timeout_task
+            except asyncio.CancelledError:
+                pass
+
+    # Background Loops ─────────────────────────────────────────────────────
+
+    async def _correlation_cleanup_loop(self) -> None:
+        """Background task that periodically cleans up expired correlation groups.
+
+        Runs continuously until all correlation groups are cleared or orchestrator shuts down.
+        Checks every 100ms for time-based expired correlations and discards them.
+        """
+        try:
+            while True:
+                await asyncio.sleep(self._cleanup_interval)
+                self._cleanup_expired_correlations()
+
+                # Stop if no correlation groups remain
+                if not self._correlation_engine.correlation_groups:
+                    self._correlation_cleanup_task = None
+                    break
+        except asyncio.CancelledError:
+            # Clean shutdown
+            self._correlation_cleanup_task = None
+            raise
+
+    def _cleanup_expired_correlations(self) -> None:
+        """Clean up all expired correlation groups across all subscriptions.
+
+        Called periodically by background task to enforce time-based correlation windows.
+        Discards incomplete correlations that have exceeded their time window.
+        """
+        # Get all active subscription keys
+        for agent_name, subscription_index in list(
+            self._correlation_engine.correlation_groups.keys()
+        ):
+            self._correlation_engine.cleanup_expired(agent_name, subscription_index)
+
+    async def _batch_timeout_checker_loop(self) -> None:
+        """Background task that periodically checks for batch timeouts.
+
+        Runs continuously until all batches are cleared or orchestrator shuts down.
+        Checks every 100ms for expired batches and flushes them via callback.
+        """
+        try:
+            while True:
+                await asyncio.sleep(self._cleanup_interval)
+
+                # Call the timeout callback to check and flush expired batches
+                if self._batch_timeout_callback:
+                    await self._batch_timeout_callback()
+
+                # Stop if no batches remain
+                if not self._batch_engine.batches:
+                    self._batch_timeout_task = None
+                    break
+        except asyncio.CancelledError:
+            # Clean shutdown
+            self._batch_timeout_task = None
+            raise
+
+    # Helper Methods ───────────────────────────────────────────────────────
+
+    async def check_batch_timeouts(self, orchestrator_callback: Any) -> None:
+        """Check all batches for timeout expiry and invoke callback for expired batches.
+
+        This method is called periodically by the background timeout checker
+        or manually (in tests) to enforce timeout-based batching.
+
+        Args:
+            orchestrator_callback: Async function to call for each expired batch.
+                Signature: async def callback(agent_name: str, subscription_index: int,
+                                              artifacts: list[Artifact]) -> None
+        """
+        expired_batches = self._batch_engine.check_timeouts()
+
+        for agent_name, subscription_index in expired_batches:
+            # Flush the expired batch
+            artifacts = self._batch_engine.flush_batch(agent_name, subscription_index)
+
+            if artifacts is not None:
+                # Invoke orchestrator callback to schedule task
+                await orchestrator_callback(agent_name, subscription_index, artifacts)
+
+    async def flush_all_batches(self, orchestrator_callback: Any) -> None:
+        """Flush all partial batches (for shutdown - ensures zero data loss).
+
+        Args:
+            orchestrator_callback: Async function to call for each flushed batch.
+                Signature: async def callback(agent_name: str, subscription_index: int,
+                                              artifacts: list[Artifact]) -> None
+        """
+        all_batches = self._batch_engine.flush_all()
+
+        for agent_name, subscription_index, artifacts in all_batches:
+            # Invoke orchestrator callback to schedule task
+            await orchestrator_callback(agent_name, subscription_index, artifacts)
+
+    # Properties ───────────────────────────────────────────────────────────
+
+    @property
+    def has_pending_correlations(self) -> bool:
+        """Check if there are any pending correlation groups."""
+        return any(
+            groups and any(group.waiting_artifacts for group in groups.values())
+            for groups in self._correlation_engine.correlation_groups.values()
+        )
+
+    @property
+    def has_pending_batches(self) -> bool:
+        """Check if there are any pending batches."""
+        return any(
+            accumulator.artifacts for accumulator in self._batch_engine.batches.values()
+        )
+
+
+__all__ = ["LifecycleManager"]