flock-core 0.5.11__py3-none-any.whl → 0.5.21__py3-none-any.whl

flock/orchestrator/artifact_manager.py

@@ -0,0 +1,168 @@
+"""Artifact publishing and persistence."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+from flock.core.artifacts import Artifact
+from flock.core.visibility import PublicVisibility, Visibility
+from flock.registry import type_registry
+
+
+if TYPE_CHECKING:
+    from flock.core import Flock
+    from flock.core.store import BlackboardStore
+    from flock.orchestrator import AgentScheduler
+
+
+class ArtifactManager:
+    """Manages artifact publishing and persistence.
+
+    Responsibilities:
+    - Normalize different input types (BaseModel, dict, Artifact)
+    - Persist artifacts to store
+    - Trigger scheduling after publish
+    - Handle batch publishing
+    """
+
+    def __init__(
+        self, orchestrator: Flock, store: BlackboardStore, scheduler: AgentScheduler
+    ):
+        """Initialize artifact manager.
+
+        Args:
+            orchestrator: Flock orchestrator instance
+            store: Blackboard store for persistence
+            scheduler: Scheduler for triggering agent execution
+        """
+        self._orchestrator = orchestrator
+        self._store = store
+        self._scheduler = scheduler
+        self._logger = orchestrator._logger
+
+    async def publish(
+        self,
+        obj: BaseModel | dict | Artifact,
+        *,
+        visibility: Visibility | None = None,
+        correlation_id: str | None = None,
+        partition_key: str | None = None,
+        tags: set[str] | None = None,
+        is_dashboard: bool = False,
+    ) -> Artifact:
+        """Publish an artifact to the blackboard (event-driven).
+
+        All agents with matching subscriptions will be triggered according to
+        their filters (type, predicates, visibility, etc).
+
+        Args:
+            obj: Object to publish (BaseModel instance, dict, or Artifact)
+            visibility: Access control (defaults to PublicVisibility)
+            correlation_id: Optional correlation ID for request tracing
+            partition_key: Optional partition key for sharding
+            tags: Optional tags for channel-based routing
+            is_dashboard: Internal flag for dashboard events
+
+        Returns:
+            The published Artifact
+
+        Examples:
+            >>> # Publish a model instance (recommended)
+            >>> task = Task(name="Deploy", priority=5)
+            >>> await artifact_manager.publish(task)
+
+            >>> # Publish with custom visibility
+            >>> await artifact_manager.publish(
+            ...     task, visibility=PrivateVisibility(agents={"admin"})
+            ... )
+        """
+        # Handle different input types
+        if isinstance(obj, Artifact):
+            # Already an artifact - publish as-is
+            artifact = obj
+        elif isinstance(obj, BaseModel):
+            # BaseModel instance - get type from registry
+            type_name = type_registry.name_for(type(obj))
+            artifact = Artifact(
+                type=type_name,
+                payload=obj.model_dump(),
+                produced_by="external",
+                visibility=visibility or PublicVisibility(),
+                correlation_id=correlation_id or uuid4(),
+                partition_key=partition_key,
+                tags=tags or set(),
+            )
+        elif isinstance(obj, dict):
+            # Dict must have 'type' key
+            if "type" not in obj:
+                raise ValueError(
+                    "Dict input must contain 'type' key. "
+                    "Example: {'type': 'Task', 'name': 'foo', 'priority': 5}"
+                )
+            # Support both {'type': 'X', 'payload': {...}} and {'type': 'X', ...}
+            type_name = obj["type"]
+            if "payload" in obj:
+                payload = obj["payload"]
+            else:
+                payload = {k: v for k, v in obj.items() if k != "type"}
+
+            artifact = Artifact(
+                type=type_name,
+                payload=payload,
+                produced_by="external",
+                visibility=visibility or PublicVisibility(),
+                correlation_id=correlation_id,
+                partition_key=partition_key,
+                tags=tags or set(),
+            )
+        else:
+            raise TypeError(
+                f"Cannot publish object of type {type(obj).__name__}. "
+                "Expected BaseModel, dict, or Artifact."
+            )
+
+        # Persist and schedule matching agents
+        await self.persist_and_schedule(artifact)
+        return artifact
+
+    async def publish_many(
+        self, objects: Iterable[BaseModel | dict | Artifact], **kwargs: Any
+    ) -> list[Artifact]:
+        """Publish multiple artifacts at once (event-driven).
+
+        Args:
+            objects: Iterable of objects to publish
+            **kwargs: Passed to each publish() call (visibility, tags, etc)
+
+        Returns:
+            List of published Artifacts
+
+        Example:
+            >>> tasks = [
+            ...     Task(name="Deploy", priority=5),
+            ...     Task(name="Test", priority=3),
+            ... ]
+            >>> await artifact_manager.publish_many(tasks, tags={"sprint-3"})
+        """
+        artifacts = []
+        for obj in objects:
+            artifact = await self.publish(obj, **kwargs)
+            artifacts.append(artifact)
+        return artifacts
+
+    async def persist_and_schedule(self, artifact: Artifact) -> None:
+        """Persist artifact to store and trigger scheduling.
+
+        Args:
+            artifact: Artifact to publish
+        """
+        await self._store.publish(artifact)
+        self._orchestrator.metrics["artifacts_published"] += 1
+        await self._scheduler.schedule_artifact(artifact)
+
+
+__all__ = ["ArtifactManager"]

flock/{batch_accumulator.py → orchestrator/batch_accumulator.py}

@@ -16,8 +16,8 @@ from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
-    from flock.artifacts import Artifact
-    from flock.subscription import BatchSpec, Subscription
+    from flock.core.artifacts import Artifact
+    from flock.core.subscription import BatchSpec, Subscription
 
 
 class BatchAccumulator:

flock/orchestrator/component_runner.py

@@ -0,0 +1,389 @@
+"""Component lifecycle hook execution for orchestrator.
+
+This module manages the execution of OrchestratorComponent hooks in priority order.
+Components can modify artifacts, control scheduling decisions, and handle collection logic.
+"""
+
+from __future__ import annotations
+
+import logging
+from asyncio import Task
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.components.orchestrator import OrchestratorComponent
+    from flock.core import Agent
+    from flock.core.artifacts import Artifact
+    from flock.core.subscription import Subscription
+
+
+class ComponentRunner:
+    """Executes orchestrator component hooks in priority order.
+
+    This class manages the component lifecycle including initialization,
+    artifact processing, scheduling decisions, and shutdown. All hooks
+    execute in priority order (lower priority number = earlier execution).
+
+    Attributes:
+        _components: List of orchestrator components (sorted by priority)
+        _logger: Logger instance for component execution tracking
+        _initialized: Flag to prevent multiple initializations
+    """
+
+    def __init__(
+        self,
+        components: list[OrchestratorComponent],
+        logger: logging.Logger | None = None,
+    ) -> None:
+        """Initialize the component runner.
+
+        Args:
+            components: List of orchestrator components (should be pre-sorted by priority)
+            logger: Logger instance (defaults to module logger if not provided)
+        """
+        self._components = components
+        self._logger = logger or logging.getLogger(__name__)
+        self._initialized = False
+
+    @property
+    def components(self) -> list[OrchestratorComponent]:
+        """Get the list of registered components."""
+        return self._components
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if components have been initialized."""
+        return self._initialized
+
+    async def run_initialize(self, orchestrator: Any) -> None:
+        """Initialize all components in priority order (called once).
+
+        Executes on_initialize hook for each component. Sets _initialized
+        flag to prevent multiple initializations.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+        """
+        if self._initialized:
+            return
+
+        self._logger.info(
+            f"Initializing {len(self._components)} orchestrator components"
+        )
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+            self._logger.debug(
+                f"Initializing component: name={comp_name}, priority={component.priority}"
+            )
+
+            try:
+                await component.on_initialize(orchestrator)
+            except Exception as e:
+                self._logger.exception(
+                    f"Component initialization failed: name={comp_name}, error={e!s}"
+                )
+                raise
+
+        self._initialized = True
+        self._logger.info(f"All components initialized: count={len(self._components)}")
+
+    async def run_artifact_published(
+        self, orchestrator: Any, artifact: Artifact
+    ) -> Artifact | None:
+        """Run on_artifact_published hooks (returns modified artifact or None to block).
+
+        Components execute in priority order, each receiving the artifact from the
+        previous component (chaining). If any component returns None, the artifact
+        is blocked and scheduling stops.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            artifact: The artifact being published
+
+        Returns:
+            Modified artifact or None if blocked by a component
+        """
+        current_artifact = artifact
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+            self._logger.debug(
+                f"Running on_artifact_published: component={comp_name}, "
+                f"artifact_type={current_artifact.type}, artifact_id={current_artifact.id}"
+            )
+
+            try:
+                result = await component.on_artifact_published(
+                    orchestrator, current_artifact
+                )
+
+                if result is None:
+                    self._logger.info(
+                        f"Artifact blocked by component: component={comp_name}, "
+                        f"artifact_type={current_artifact.type}, artifact_id={current_artifact.id}"
+                    )
+                    return None
+
+                current_artifact = result
+            except Exception as e:
+                self._logger.exception(
+                    f"Component hook failed: component={comp_name}, "
+                    f"hook=on_artifact_published, error={e!s}"
+                )
+                raise
+
+        return current_artifact
+
+    async def run_before_schedule(
+        self,
+        orchestrator: Any,
+        artifact: Artifact,
+        agent: Agent,
+        subscription: Subscription,
+    ) -> Any:
+        """Run on_before_schedule hooks (returns CONTINUE, SKIP, or DEFER).
+
+        Components execute in priority order. First component to return SKIP or
+        DEFER stops execution and returns that decision.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            artifact: The artifact being scheduled
+            agent: The target agent
+            subscription: The subscription being evaluated
+
+        Returns:
+            ScheduleDecision (CONTINUE, SKIP, or DEFER)
+        """
+        from flock.components.orchestrator import ScheduleDecision
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+
+            self._logger.debug(
+                f"Running on_before_schedule: component={comp_name}, "
+                f"agent={agent.name}, artifact_type={artifact.type}"
+            )
+
+            try:
+                decision = await component.on_before_schedule(
+                    orchestrator, artifact, agent, subscription
+                )
+
+                if decision == ScheduleDecision.SKIP:
+                    self._logger.info(
+                        f"Scheduling skipped by component: component={comp_name}, "
+                        f"agent={agent.name}, artifact_type={artifact.type}, decision=SKIP"
+                    )
+                    return ScheduleDecision.SKIP
+
+                if decision == ScheduleDecision.DEFER:
+                    self._logger.debug(
+                        f"Scheduling deferred by component: component={comp_name}, "
+                        f"agent={agent.name}, decision=DEFER"
+                    )
+                    return ScheduleDecision.DEFER
+
+            except Exception as e:
+                self._logger.exception(
+                    f"Component hook failed: component={comp_name}, "
+                    f"hook=on_before_schedule, error={e!s}"
+                )
+                raise
+
+        return ScheduleDecision.CONTINUE
+
+    async def run_collect_artifacts(
+        self,
+        orchestrator: Any,
+        artifact: Artifact,
+        agent: Agent,
+        subscription: Subscription,
+    ) -> Any:
+        """Run on_collect_artifacts hooks (returns first non-None result).
+
+        Components execute in priority order. First component to return non-None
+        wins (short-circuit). If all return None, default is immediate scheduling.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            artifact: The artifact being collected
+            agent: The target agent
+            subscription: The subscription being evaluated
+
+        Returns:
+            CollectionResult (complete=True/False, artifacts=[...])
+        """
+        from flock.components.orchestrator import CollectionResult
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+
+            self._logger.debug(
+                f"Running on_collect_artifacts: component={comp_name}, "
+                f"agent={agent.name}, artifact_type={artifact.type}"
+            )
+
+            try:
+                result = await component.on_collect_artifacts(
+                    orchestrator, artifact, agent, subscription
+                )
+
+                if result is not None:
+                    self._logger.debug(
+                        f"Collection handled by component: component={comp_name}, "
+                        f"complete={result.complete}, artifact_count={len(result.artifacts)}"
+                    )
+                    return result
+            except Exception as e:
+                self._logger.exception(
+                    f"Component hook failed: component={comp_name}, "
+                    f"hook=on_collect_artifacts, error={e!s}"
+                )
+                raise
+
+        # Default: immediate scheduling with single artifact
+        self._logger.debug(
+            f"No component handled collection, using default: "
+            f"agent={agent.name}, artifact_type={artifact.type}"
+        )
+        return CollectionResult.immediate([artifact])
+
+    async def run_before_agent_schedule(
+        self, orchestrator: Any, agent: Agent, artifacts: list[Artifact]
+    ) -> list[Artifact] | None:
+        """Run on_before_agent_schedule hooks (returns modified artifacts or None to block).
+
+        Components execute in priority order, each receiving artifacts from the
+        previous component (chaining). If any component returns None, scheduling
+        is blocked.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            agent: The target agent
+            artifacts: List of artifacts to schedule
+
+        Returns:
+            Modified artifacts list or None if scheduling blocked
+        """
+        current_artifacts = artifacts
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+
+            self._logger.debug(
+                f"Running on_before_agent_schedule: component={comp_name}, "
+                f"agent={agent.name}, artifact_count={len(current_artifacts)}"
+            )
+
+            try:
+                result = await component.on_before_agent_schedule(
+                    orchestrator, agent, current_artifacts
+                )
+
+                if result is None:
+                    self._logger.info(
+                        f"Agent scheduling blocked by component: component={comp_name}, "
+                        f"agent={agent.name}"
+                    )
+                    return None
+
+                current_artifacts = result
+            except Exception as e:
+                self._logger.exception(
+                    f"Component hook failed: component={comp_name}, "
+                    f"hook=on_before_agent_schedule, error={e!s}"
+                )
+                raise
+
+        return current_artifacts
+
+    async def run_agent_scheduled(
+        self,
+        orchestrator: Any,
+        agent: Agent,
+        artifacts: list[Artifact],
+        task: Task[Any],
+    ) -> None:
+        """Run on_agent_scheduled hooks (notification only, non-blocking).
+
+        Components execute in priority order. Exceptions are logged but don't
+        prevent other components from executing or block scheduling.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+            agent: The scheduled agent
+            artifacts: List of artifacts for the agent
+            task: The asyncio task for agent execution
+        """
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+
+            self._logger.debug(
+                f"Running on_agent_scheduled: component={comp_name}, "
+                f"agent={agent.name}, artifact_count={len(artifacts)}"
+            )
+
+            try:
+                await component.on_agent_scheduled(orchestrator, agent, artifacts, task)
+            except Exception as e:
+                self._logger.warning(
+                    f"Component notification hook failed (non-critical): "
+                    f"component={comp_name}, hook=on_agent_scheduled, error={e!s}"
+                )
+                # Don't propagate - this is a notification hook
+
+    async def run_idle(self, orchestrator: Any) -> None:
+        """Run on_orchestrator_idle hooks when orchestrator becomes idle.
+
+        Components execute in priority order. Exceptions are logged but don't
+        prevent other components from executing.
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+        """
+        self._logger.debug(
+            f"Running on_orchestrator_idle hooks: component_count={len(self._components)}"
+        )
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+
+            try:
+                await component.on_orchestrator_idle(orchestrator)
+            except Exception as e:
+                self._logger.warning(
+                    f"Component idle hook failed (non-critical): "
+                    f"component={comp_name}, hook=on_orchestrator_idle, error={e!s}"
+                )
+
+    async def run_shutdown(self, orchestrator: Any) -> None:
+        """Run on_shutdown hooks when orchestrator shuts down.
+
+        Components execute in priority order. Exceptions are logged but don't
+        prevent shutdown of other components (best-effort cleanup).
+
+        Args:
+            orchestrator: The Flock orchestrator instance
+        """
+        self._logger.info(
+            f"Shutting down {len(self._components)} orchestrator components"
+        )
+
+        for component in self._components:
+            comp_name = component.name or component.__class__.__name__
+            self._logger.debug(f"Shutting down component: name={comp_name}")
+
+            try:
+                await component.on_shutdown(orchestrator)
+            except Exception as e:
+                self._logger.exception(
+                    f"Component shutdown failed: component={comp_name}, "
+                    f"hook=on_shutdown, error={e!s}"
+                )
+                # Continue shutting down other components
+
+
+__all__ = ["ComponentRunner"]