flock-core 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

flock/agent.py

@@ -988,10 +988,23 @@ class AgentBuilder:
     def _normalize_join(self, value: dict | JoinSpec | None) -> JoinSpec | None:
         if value is None or isinstance(value, JoinSpec):
             return value
+        # Phase 2: New JoinSpec API with 'by' and 'within' (time OR count)
+        from datetime import timedelta
+
+        within_value = value.get("within")
+        if isinstance(within_value, (int, float)):
+            # Count window or seconds as float - keep as is
+            within = (
+                int(within_value)
+                if isinstance(within_value, int)
+                else timedelta(seconds=within_value)
+            )
+        else:
+            # Default to 1 minute time window
+            within = timedelta(minutes=1)
         return JoinSpec(
-            kind=value.get("kind", "all_of"),
-            window=float(value.get("window", 0.0)),
-            by=value.get("by"),
+            by=value["by"],  # Required
+            within=within,
         )
 
     def _normalize_batch(self, value: dict | BatchSpec | None) -> BatchSpec | None:

flock/artifact_collector.py

@@ -0,0 +1,158 @@
+"""Artifact collection and waiting pool management for AND gate logic.
+
+This module implements the waiting pool mechanism that enables `.consumes(A, B)`
+to wait for BOTH types before triggering an agent (AND gate logic).
+
+Architecture:
+- Each subscription gets a unique waiting pool identified by (agent_name, subscription_index)
+- Artifacts are collected per type until all required types are present
+- When complete, all collected artifacts are returned for agent execution
+- After triggering, the waiting pool is cleared for the next cycle
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from flock.agent import Agent
+    from flock.artifacts import Artifact
+    from flock.subscription import Subscription
+
+
+class ArtifactCollector:
+    """Manages waiting pools for multi-type subscriptions (AND gate logic).
+
+    Each subscription with multiple types gets a waiting pool that collects
+    artifacts until all required types are present. Single-type subscriptions
+    bypass the waiting pool for immediate triggering.
+
+    Example:
+        agent.consumes(TypeA, TypeB)  # Creates waiting pool for 2 types
+
+        # TypeA published → added to pool (not complete yet)
+        # TypeB published → added to pool (NOW complete!)
+        # → Agent triggered with [TypeA_artifact, TypeB_artifact]
+        # → Waiting pool cleared for next cycle
+    """
+
+    def __init__(self) -> None:
+        """Initialize empty waiting pools."""
+        # Structure: {(agent_name, subscription_index): {type_name: [artifact1, artifact2, ...]}}
+        # Example: {("diagnostician", 0): {"XRay": [artifact1], "LabResult": [artifact2]}}
+        # For count-based AND gates: {"TypeA": [artifact1, artifact2, artifact3]} (3 As collected)
+        self._waiting_pools: dict[tuple[str, int], dict[str, list[Artifact]]] = defaultdict(
+            lambda: defaultdict(list)
+        )
+
+    def add_artifact(
+        self,
+        agent: Agent,
+        subscription: Subscription,
+        artifact: Artifact,
+    ) -> tuple[bool, list[Artifact]]:
+        """Add artifact to waiting pool and check for completeness.
+
+        Args:
+            agent: Agent that will process the artifacts
+            subscription: Subscription that matched the artifact
+            artifact: Artifact to add to the waiting pool
+
+        Returns:
+            Tuple of (is_complete, artifacts):
+                - is_complete: True if all required types are now present
+                - artifacts: List of collected artifacts (empty if incomplete, all artifacts if complete)
+
+        Design Notes:
+            - Single-type subscriptions with count=1 bypass the pool and return immediately complete
+            - Multi-type or count-based subscriptions collect artifacts until all required counts met
+            - Latest artifacts win (keeps most recent N artifacts per type)
+            - After returning complete=True, the pool is automatically cleared
+        """
+        # Single-type subscription with count=1: No waiting needed (immediate trigger)
+        if len(subscription.type_names) == 1 and subscription.type_counts[artifact.type] == 1:
+            return (True, [artifact])
+
+        # Multi-type or count-based subscription: Use waiting pool (AND gate logic)
+
+        # Find subscription index (agents can have multiple subscriptions)
+        try:
+            subscription_index = agent.subscriptions.index(subscription)
+        except ValueError:
+            # Should never happen, but defensive programming
+            raise RuntimeError(
+                f"Subscription not found in agent {agent.name}. "
+                "This indicates an internal orchestrator error."
+            )
+
+        pool_key = (agent.name, subscription_index)
+
+        # Add artifact to pool (collect in list for count-based logic)
+        self._waiting_pools[pool_key][artifact.type].append(artifact)
+
+        # Check if all required counts are met
+        is_complete = True
+        for type_name, required_count in subscription.type_counts.items():
+            collected_count = len(self._waiting_pools[pool_key][type_name])
+            if collected_count < required_count:
+                is_complete = False
+                break
+
+        if is_complete:
+            # Complete! Collect all artifacts (flatten lists) and clear the pool
+            artifacts = []
+            for type_name, required_count in subscription.type_counts.items():
+                # Take exactly the required count (latest artifacts)
+                type_artifacts = self._waiting_pools[pool_key][type_name]
+                artifacts.extend(type_artifacts[:required_count])
+
+            del self._waiting_pools[pool_key]  # Clear for next cycle
+            return (True, artifacts)
+        # Incomplete - still waiting for more artifacts
+        return (False, [])
+
+    def get_waiting_status(
+        self, agent: Agent, subscription_index: int
+    ) -> dict[str, list[Artifact]]:
+        """Get current waiting pool contents for debugging/inspection.
+
+        Args:
+            agent: Agent to inspect
+            subscription_index: Index of the subscription
+
+        Returns:
+            Dictionary mapping type names to lists of collected artifacts (empty if none)
+        """
+        pool_key = (agent.name, subscription_index)
+        # Return a copy to prevent external mutation
+        pool = self._waiting_pools.get(pool_key, {})
+        return {type_name: list(artifacts) for type_name, artifacts in pool.items()}
+
+    def clear_waiting_pool(self, agent: Agent, subscription_index: int) -> None:
+        """Manually clear a waiting pool.
+
+        Useful for cleanup or resetting agent state.
+
+        Args:
+            agent: Agent whose pool to clear
+            subscription_index: Index of the subscription
+        """
+        pool_key = (agent.name, subscription_index)
+        if pool_key in self._waiting_pools:
+            del self._waiting_pools[pool_key]
+
+    def clear_all_pools(self) -> None:
+        """Clear all waiting pools.
+
+        Useful for orchestrator shutdown or test cleanup.
+        """
+        self._waiting_pools.clear()
+
+    def get_pool_count(self) -> int:
+        """Get total number of active waiting pools (for metrics/debugging)."""
+        return len(self._waiting_pools)
+
+
+__all__ = ["ArtifactCollector"]

flock/batch_accumulator.py

@@ -0,0 +1,252 @@
+"""
+BatchAccumulator: Manages batch collection with size/timeout triggers.
+
+Supports BatchSpec-based batching:
+- Accumulates artifacts in batches per subscription
+- Flushes on size threshold (e.g., batch of 25)
+- Flushes on timeout (e.g., every 30 seconds)
+- Whichever comes first wins
+- Ensures zero data loss on shutdown
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from flock.artifacts import Artifact
+    from flock.subscription import BatchSpec, Subscription
+
+
+class BatchAccumulator:
+    """
+    Tracks artifact batches waiting for size/timeout triggers.
+
+    Example: For orders, accumulate 25 at a time to batch process payments.
+    When 25th order arrives OR 30 seconds elapse, flush the batch.
+    """
+
+    def __init__(
+        self,
+        *,
+        batch_spec: BatchSpec,
+        created_at: datetime,
+    ):
+        self.batch_spec = batch_spec
+        self.created_at = created_at  # When first artifact arrived
+        self.artifacts: list[Artifact] = []
+
+    def add_artifact(self, artifact: Artifact) -> bool:
+        """
+        Add artifact to batch.
+
+        Returns:
+            True if batch should flush (size threshold reached), False otherwise
+        """
+        self.artifacts.append(artifact)
+
+        # Check size threshold
+        if self.batch_spec.size is not None:
+            if len(self.artifacts) >= self.batch_spec.size:
+                return True  # Flush now (size threshold reached)
+
+        return False  # Not ready to flush yet
+
+    def is_timeout_expired(self) -> bool:
+        """Check if timeout has expired since batch started."""
+        if self.batch_spec.timeout is None:
+            return False
+
+        elapsed = datetime.now() - self.created_at
+        return elapsed >= self.batch_spec.timeout
+
+    def get_artifacts(self) -> list[Artifact]:
+        """Get all artifacts in batch."""
+        return self.artifacts.copy()
+
+    def clear(self) -> None:
+        """Clear the batch after flush."""
+        self.artifacts.clear()
+
+
+class BatchEngine:
+    """
+    Manages batch state for BatchSpec subscriptions.
+
+    Responsibilities:
+    1. Accumulate artifacts per (agent, subscription_index)
+    2. Track batch size and timeout per batch
+    3. Return complete batches when size or timeout threshold met
+    4. Provide shutdown flush for partial batches
+
+    Example usage:
+        engine = BatchEngine()
+
+        # Add artifact to batch
+        should_flush = engine.add_artifact(
+            artifact=order_artifact,
+            subscription=subscription,  # Has BatchSpec
+            subscription_index=0,
+        )
+
+        if should_flush:
+            # Size threshold reached! Flush batch
+            artifacts = engine.flush_batch("agent_name", 0)
+            # Trigger agent with batch
+    """
+
+    def __init__(self):
+        # Batch state per (agent_name, subscription_index)
+        # Key: (agent_name, subscription_index)
+        # Value: BatchAccumulator
+        self.batches: dict[tuple[str, int], BatchAccumulator] = {}
+
+    def add_artifact(
+        self,
+        *,
+        artifact: Artifact,
+        subscription: Subscription,
+        subscription_index: int,
+    ) -> bool:
+        """
+        Add artifact to batch accumulator.
+
+        Returns:
+            True if batch should flush (size threshold reached), False otherwise
+        """
+        if subscription.batch is None:
+            raise ValueError("Subscription must have BatchSpec for batching")
+
+        batch_key = (subscription.agent_name, subscription_index)
+
+        # Get or create batch accumulator
+        if batch_key not in self.batches:
+            self.batches[batch_key] = BatchAccumulator(
+                batch_spec=subscription.batch,
+                created_at=datetime.now(),
+            )
+
+        accumulator = self.batches[batch_key]
+
+        # Add artifact to batch
+        should_flush = accumulator.add_artifact(artifact)
+
+        return should_flush
+
+    def add_artifact_group(
+        self,
+        *,
+        artifacts: list[Artifact],
+        subscription: Subscription,
+        subscription_index: int,
+    ) -> bool:
+        """
+        Add a GROUP of artifacts (e.g., correlated pair) as a SINGLE batch item.
+
+        This is used for JoinSpec + BatchSpec combinations where we want to batch
+        correlated groups, not individual artifacts.
+
+        Example: JoinSpec + BatchSpec(size=2) means "batch 2 correlated pairs",
+                 not "batch 2 individual artifacts".
+
+        Returns:
+            True if batch should flush (size threshold reached), False otherwise
+        """
+        if subscription.batch is None:
+            raise ValueError("Subscription must have BatchSpec for batching")
+
+        batch_key = (subscription.agent_name, subscription_index)
+
+        # Get or create batch accumulator
+        if batch_key not in self.batches:
+            self.batches[batch_key] = BatchAccumulator(
+                batch_spec=subscription.batch,
+                created_at=datetime.now(),
+            )
+
+        accumulator = self.batches[batch_key]
+
+        # Add ALL artifacts from the group
+        for artifact in artifacts:
+            accumulator.artifacts.append(artifact)
+
+        # Check size threshold - count GROUPS, not artifacts
+        # We track how many groups have been added by checking batch_spec metadata
+        if subscription.batch.size is not None:
+            # For group batching, we need to track group count separately
+            # For now, we'll use a simple heuristic: count groups by dividing by expected group size
+            # But this is NOT perfect - we need better tracking
+
+            # BETTER APPROACH: Count how many times we've called add_artifact_group
+            # For now, let's use artifact count as a proxy and check if we've hit the threshold
+            # This will work correctly if all groups are the same size
+
+            # Actually, let's track group count properly:
+            if not hasattr(accumulator, "_group_count"):
+                accumulator._group_count = 0
+
+            accumulator._group_count += 1
+
+            if accumulator._group_count >= subscription.batch.size:
+                return True  # Flush now
+
+        return False  # Not ready to flush yet
+
+    def flush_batch(self, agent_name: str, subscription_index: int) -> list[Artifact] | None:
+        """
+        Flush a batch and return its artifacts.
+
+        Returns:
+            List of artifacts in batch, or None if no batch exists
+        """
+        batch_key = (agent_name, subscription_index)
+
+        accumulator = self.batches.get(batch_key)
+        if accumulator is None or not accumulator.artifacts:
+            return None
+
+        # Get artifacts and clear batch
+        artifacts = accumulator.get_artifacts()
+        del self.batches[batch_key]
+
+        return artifacts
+
+    def check_timeouts(self) -> list[tuple[str, int]]:
+        """
+        Check all batches for timeout expiry.
+
+        Returns:
+            List of (agent_name, subscription_index) tuples that should flush
+        """
+        expired = []
+
+        for batch_key, accumulator in list(self.batches.items()):
+            if accumulator.is_timeout_expired():
+                expired.append(batch_key)
+
+        return expired
+
+    def flush_all(self) -> list[tuple[str, int, list[Artifact]]]:
+        """
+        Flush ALL partial batches (for shutdown).
+
+        Returns:
+            List of (agent_name, subscription_index, artifacts) tuples
+        """
+        results = []
+
+        for batch_key, accumulator in list(self.batches.items()):
+            if accumulator.artifacts:
+                artifacts = accumulator.get_artifacts()
+                agent_name, subscription_index = batch_key
+                results.append((agent_name, subscription_index, artifacts))
+
+        # Clear all batches after flush
+        self.batches.clear()
+
+        return results
+
+
+__all__ = ["BatchAccumulator", "BatchEngine"]

flock/correlation_engine.py

@@ -0,0 +1,223 @@
+"""
+CorrelationEngine: Manages correlated AND gates with time/count windows.
+
+Supports JoinSpec-based correlation:
+- Extracts correlation keys from artifacts
+- Groups artifacts by correlation key
+- Enforces time windows (timedelta) or count windows (int)
+- Triggers agents when all required types arrive within window
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from flock.artifacts import Artifact
+    from flock.subscription import JoinSpec, Subscription
+
+
+class CorrelationGroup:
+    """
+    Tracks artifacts waiting for correlation within a specific key group.
+
+    Example: For patient-123, track X-ray (TypeA) and Lab results (TypeB).
+    When both arrive within time/count window, trigger the agent.
+    """
+
+    def __init__(
+        self,
+        *,
+        correlation_key: Any,
+        required_types: set[str],
+        type_counts: dict[str, int],
+        window_spec: timedelta | int,
+        created_at_sequence: int,
+    ):
+        self.correlation_key = correlation_key
+        self.required_types = required_types  # e.g., {"TypeA", "TypeB"}
+        self.type_counts = type_counts  # e.g., {"TypeA": 1, "TypeB": 1}
+        self.window_spec = window_spec  # timedelta or int
+        self.created_at_sequence = (
+            created_at_sequence  # Global sequence when first artifact arrived
+        )
+        self.created_at_time: datetime | None = None  # Timestamp when first artifact arrived
+
+        # Waiting pool: type -> list of artifacts
+        self.waiting_artifacts: dict[str, list[Artifact]] = defaultdict(list)
+
+    def add_artifact(self, artifact: Artifact, current_sequence: int) -> None:
+        """Add artifact to this correlation group's waiting pool."""
+        if self.created_at_time is None:
+            from datetime import timezone
+
+            self.created_at_time = datetime.now(timezone.utc)
+
+        self.waiting_artifacts[artifact.type].append(artifact)
+
+    def is_complete(self) -> bool:
+        """Check if all required types have arrived with correct counts."""
+        for type_name, required_count in self.type_counts.items():
+            if len(self.waiting_artifacts.get(type_name, [])) < required_count:
+                return False
+        return True
+
+    def is_expired(self, current_sequence: int) -> bool:
+        """Check if this correlation group has expired based on window."""
+        if isinstance(self.window_spec, int):
+            # Count window: expired if current sequence exceeds created + window
+            return (current_sequence - self.created_at_sequence) > self.window_spec
+        if isinstance(self.window_spec, timedelta):
+            # Time window: expired if current time exceeds created + window
+            if self.created_at_time is None:
+                return False
+            from datetime import timezone
+
+            elapsed = datetime.now(timezone.utc) - self.created_at_time
+            return elapsed > self.window_spec
+        return False
+
+    def get_artifacts(self) -> list[Artifact]:
+        """Get all artifacts in the order they should be passed to the agent."""
+        result = []
+        for type_name in self.required_types:
+            # Get the required number of artifacts for this type
+            required_count = self.type_counts[type_name]
+            artifacts_for_type = self.waiting_artifacts[type_name][:required_count]
+            result.extend(artifacts_for_type)
+        return result
+
+
+class CorrelationEngine:
+    """
+    Manages correlation state for JoinSpec subscriptions.
+
+    Responsibilities:
+    1. Extract correlation keys from artifacts using JoinSpec.by lambda
+    2. Group artifacts by correlation key
+    3. Track time/count windows per correlation group
+    4. Return complete correlation groups when all types arrive within window
+    5. Clean up expired correlations
+
+    Example usage:
+        engine = CorrelationEngine()
+
+        # Add artifact to correlation tracking
+        completed = engine.add_artifact(
+            artifact=xray_artifact,
+            subscription=subscription,  # Has JoinSpec with by + within
+            agent_name="diagnostician"
+        )
+
+        if completed:
+            # All types arrived! Trigger agent with correlated artifacts
+            artifacts = completed.get_artifacts()
+    """
+
+    def __init__(self):
+        # Global artifact sequence (for count windows)
+        self.global_sequence = 0
+
+        # Correlation state per (agent, subscription_index)
+        # Key: (agent_name, subscription_index)
+        # Value: dict[correlation_key, CorrelationGroup]
+        self.correlation_groups: dict[tuple[str, int], dict[Any, CorrelationGroup]] = defaultdict(
+            dict
+        )
+
+    def add_artifact(
+        self,
+        *,
+        artifact: Artifact,
+        subscription: Subscription,
+        subscription_index: int,
+    ) -> CorrelationGroup | None:
+        """
+        Add artifact to correlation tracking.
+
+        Returns:
+            CorrelationGroup if correlation is complete, None otherwise
+        """
+        # Increment global sequence (for count windows)
+        self.global_sequence += 1
+        current_sequence = self.global_sequence
+
+        # Extract correlation key using JoinSpec.by lambda
+        if subscription.join is None:
+            raise ValueError("Subscription must have JoinSpec for correlation")
+
+        join_spec: JoinSpec = subscription.join
+
+        # Parse artifact payload to extract correlation key
+        from flock.registry import type_registry
+
+        model_cls = type_registry.resolve(artifact.type)
+        payload_instance = model_cls(**artifact.payload)
+
+        try:
+            correlation_key = join_spec.by(payload_instance)
+        except Exception:
+            # Key extraction failed - skip this artifact
+            # TODO: Log warning?
+            return None
+
+        # Get or create correlation group for this key
+        pool_key = (subscription.agent_name, subscription_index)
+        groups = self.correlation_groups[pool_key]
+
+        if correlation_key not in groups:
+            # Create new correlation group
+            groups[correlation_key] = CorrelationGroup(
+                correlation_key=correlation_key,
+                required_types=subscription.type_names,
+                type_counts=subscription.type_counts,
+                window_spec=join_spec.within,
+                created_at_sequence=current_sequence,
+            )
+
+        group = groups[correlation_key]
+
+        # Check if group expired (for count windows, check BEFORE adding)
+        if group.is_expired(current_sequence):
+            # Group expired - remove it and start fresh
+            del groups[correlation_key]
+            # Create new group
+            groups[correlation_key] = CorrelationGroup(
+                correlation_key=correlation_key,
+                required_types=subscription.type_names,
+                type_counts=subscription.type_counts,
+                window_spec=join_spec.within,
+                created_at_sequence=current_sequence,
+            )
+            group = groups[correlation_key]
+
+        # Add artifact to group
+        group.add_artifact(artifact, current_sequence)
+
+        # Check if correlation is complete
+        if group.is_complete():
+            # Complete! Remove from tracking and return
+            completed_group = groups.pop(correlation_key)
+            return completed_group
+
+        # Not complete yet
+        return None
+
+    def cleanup_expired(self, agent_name: str, subscription_index: int) -> None:
+        """Clean up expired correlation groups for a specific subscription."""
+        pool_key = (agent_name, subscription_index)
+        groups = self.correlation_groups.get(pool_key, {})
+
+        # Remove expired groups
+        expired_keys = [
+            key for key, group in groups.items() if group.is_expired(self.global_sequence)
+        ]
+
+        for key in expired_keys:
+            del groups[key]
+
+
+__all__ = ["CorrelationEngine", "CorrelationGroup"]

flock/dashboard/collector.py

@@ -80,6 +80,9 @@ class AgentSnapshot:
     first_seen: datetime
     last_seen: datetime
     signature: str
+    logic_operations: list[dict] = field(
+        default_factory=list
+    )  # Phase 1.2: JoinSpec/BatchSpec config
 
 
 class DashboardEventCollector(AgentComponent):
@@ -514,6 +517,7 @@ class DashboardEventCollector(AgentComponent):
             first_seen=snapshot.first_seen,
             last_seen=snapshot.last_seen,
             signature=snapshot.signature,
+            logic_operations=[dict(op) for op in snapshot.logic_operations],  # Phase 1.2
         )
 
     def _snapshot_to_record(self, snapshot: AgentSnapshot) -> AgentSnapshotRecord: