Agentswarm/orchestrator/decision_engine.py

"""ACO decision engine (benchmark Group A, Option A: score-at-pull).

Implements the standard's §3 decision layer over the EXISTING pull dispatch:
when an idle agent asks for work, score all eligible ready tasks with
P_i = τ_i^α·η_i^β / Σ and SAMPLE one (ε-greedy exploration), instead of taking
the first match. The agent side of the pairing is fixed by arrival order — this
is the documented one-sided limitation of Option A (full two-sided matching =
Option B, deferred until Group C can measure whether it helps).

Two independently-switched halves:
  - LEARNING (pheromone deposits) is ALWAYS ON: every task terminal state updates
    the trail. Pure passive observation — flag-off dispatch behavior is unchanged,
    but history accrues so enabling the flag later acts on earned reputation, not
    the cold-start prior.
  - ACTING (probabilistic selection) is gated by ENABLE_ACO_DISPATCH (default OFF).

Honesty rules baked in (org rule #9 — no fabricated signals):
  - τ deposit inputs we DON'T have contribute 0, never a made-up value:
    acceptance (no acceptance signal), time (no per-task target), risk/rollback
    (no signals). quality reuses success until per-task grading exists (run-level
    fixture grading happens AFTER completion, too late for the deposit).
  - η inputs without a source are a documented neutral constant (confidence=0.5,
    identical for every candidate → no ranking distortion) or 0 (risk,
    budget_pressure when unknown).

Determinism: the RNG is seedable via ACO_SEED so CI can assert exact behavior;
production leaves it unseeded.
"""
from __future__ import annotations

import json
import logging
import os
import random
import time
from dataclasses import dataclass
from typing import Dict, List, Optional

from benchmark.metrics import THETA_DEFAULTS, heuristic, p_decision, pheromone

from .redis_client import redis_client

logger = logging.getLogger(__name__)

PHEROMONE_KEY_PREFIX = "pheromone:"

# Trail bounds: keep τ strictly positive (τ^α must stay well-defined) and capped
# (the standard's literal update τ←(1−ρ)τ+Δτ saturates; the cap is the saturation).
TAU_INITIAL = 0.5
TAU_MIN = 0.05
TAU_MAX = 1.0
ETA_MIN = 0.05

URGENCY_NORM_SECONDS = 300.0   # task age at which urgency saturates to 1.0
DEPENDENCY_NORM = 3.0          # dependent-count at which criticality saturates to 1.0
RESOURCE_NORM_SLOTS = 2.0      # free slots at which resource fit saturates to 1.0
CONFIDENCE_NEUTRAL = 0.5       # no confidence signal → same neutral for all candidates
MAX_RECORDED_DECISIONS = 1000  # cap on per-run decision telemetry


def aco_dispatch_enabled() -> bool:
    return os.getenv("ENABLE_ACO_DISPATCH", "false").lower() in {"1", "true", "yes"}


def _hyper(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, str(default)))
    except ValueError:
        return default


@dataclass
class Decision:
    """A fully replayable decision trace for one sampled assignment (#10 DoD).

    Carries not just the CHOSEN task's scores but the entire context needed to re-derive and
    explain the choice WITHOUT the live RNG: the full candidate set (each candidate's τ/η/weight/
    p_norm/dependents), the hyperparameters (α/β/ε), the seed, free_slots, the ε-greedy branch,
    and the recorded RNG draws (explore_draw + the exploit pick or explore index).
    `DecisionEngine.replay_decision(trace)` reconstructs the chosen task_id from this alone.
    """
    # chosen assignment
    task_id: str
    agent_id: str
    agent_role: str
    tau: float
    eta: float
    p_norm: float
    p_score: float           # standard §3.3 P_decision = τ^α·η^β·100 (the reported metric)
    # replay context
    explored: bool           # True when the ε-greedy branch picked uniformly
    alpha: float
    beta: float
    epsilon: float
    seed: Optional[int]
    free_slots: int
    total_weight: float
    explore_draw: float          # rng.random() compared against ε to pick the branch
    select_pick: Optional[float]  # exploit: rng.random()*total_weight (absolute); None if explored
    select_index: int             # resolved index into `candidates`
    candidates: List[dict]        # full set in evaluated order: task_id/agent_role/tau/eta/weight/p_norm/dependents

    def telemetry(self) -> dict:
        # Replay-critical fields (weights, select_pick, total_weight) kept full-precision; only
        # human-display duplicates are rounded.
        return {
            "task_id": self.task_id,
            "agent_id": self.agent_id,
            "agent_role": self.agent_role,
            "tau": round(self.tau, 6),
            "eta": round(self.eta, 6),
            "p_norm": round(self.p_norm, 6),
            "p_score": round(self.p_score, 6),
            "explored": self.explored,
            "alpha": self.alpha,
            "beta": self.beta,
            "epsilon": self.epsilon,
            "seed": self.seed,
            "free_slots": self.free_slots,
            "total_weight": self.total_weight,
            "explore_draw": self.explore_draw,
            "select_pick": self.select_pick,
            "select_index": self.select_index,
            "candidates": self.candidates,
        }


class DecisionEngine:
    """Pheromone store + η scorer + ε-greedy probabilistic selection."""

    def __init__(self, rng: Optional[random.Random] = None):
        seed = os.getenv("ACO_SEED")
        self.seed = int(seed) if seed else None
        self.rng = rng or (random.Random(self.seed) if self.seed is not None else random.Random())

    # --- pheromone trail (τ) -------------------------------------------------
    @staticmethod
    def _trail_key(agent_role: str, agent_id: str) -> str:
        return f"{PHEROMONE_KEY_PREFIX}{agent_role}:{agent_id}"

    async def get_tau(self, agent_role: str, agent_id: str) -> float:
        raw = await redis_client.get(self._trail_key(agent_role, agent_id))
        if not raw:
            return TAU_INITIAL
        try:
            return float(json.loads(raw).get("tau", TAU_INITIAL))
        except Exception:
            return TAU_INITIAL

    async def deposit(
        self,
        *,
        agent_role: str,
        agent_id: str,
        success: bool,
        cost_ratio: float = 0.0,
    ) -> float:
        """Update the trail after a task reaches a terminal state (always on).

        Standard §3.1: τ(t+1) = (1−ρ)·τ(t) + Δτ, clamped to [TAU_MIN, TAU_MAX].
        Δτ comes from metrics.pheromone() with only the signals we actually have:
        success (1/0), quality (=success — see module docstring), cost (task cost
        as a fraction of the run budget, 0 when unknown). Absent signals are 0.
        """
        rho = _hyper("ACO_RHO", THETA_DEFAULTS["rho"])
        s = 1.0 if success else 0.0
        delta = pheromone(
            success=s, quality=s, acceptance=0.0,
            cost=max(0.0, min(1.0, cost_ratio)), time=0.0, risk=0.0, rollback=0.0,
        )
        old = await self.get_tau(agent_role, agent_id)
        new = max(TAU_MIN, min(TAU_MAX, (1.0 - rho) * old + delta))
        await redis_client.set(
            self._trail_key(agent_role, agent_id),
            json.dumps({"tau": new, "updated_at": time.time()}),
        )
        logger.debug("pheromone %s:%s %.3f -> %.3f (Δ=%.3f)", agent_role, agent_id, old, new, delta)
        return new

    # --- heuristic (η) -------------------------------------------------------
    @staticmethod
    def compute_eta(task, agent_capabilities: List[str], *, free_slots: int,
                    dependents_count: int, now: Optional[float] = None) -> float:
        """A-priori desirability of (agent, task) from signals we actually hold."""
        now = now or time.time()
        required = set(task.required_capabilities or [])
        caps = set(agent_capabilities or [])
        union = required | caps
        match = (len(required & caps) / len(union)) if union else 1.0  # Jaccard: rewards focus
        urgency = min(1.0, max(0.0, now - task.created_at) / URGENCY_NORM_SECONDS)
        dependency = min(1.0, dependents_count / DEPENDENCY_NORM)
        resource = min(1.0, max(0, free_slots) / RESOURCE_NORM_SLOTS)
        eta = heuristic(
            match=match, urgency=urgency, dependency=dependency, resource=resource,
            confidence=CONFIDENCE_NEUTRAL, risk=0.0, budget_pressure=0.0,
        )
        return max(ETA_MIN, eta)

    # --- selection (P) -------------------------------------------------------
    async def select(self, agent_id: str, agent_capabilities: List[str],
                     candidates: List, *, free_slots: int,
                     dependents_counts: Dict[str, int]) -> Optional[Decision]:
        """Sample one task from the eligible candidates with P_i = τ^α·η^β / Σ.

        ε-greedy: with probability ε pick uniformly (exploration keeps low-τ
        agents able to rebuild their trail and prevents starvation).
        """
        if not candidates:
            return None
        alpha = _hyper("ACO_ALPHA", THETA_DEFAULTS["alpha"])
        beta = _hyper("ACO_BETA", THETA_DEFAULTS["beta"])
        epsilon = _hyper("ACO_EPSILON", THETA_DEFAULTS["epsilon"])

        scored = []  # (task, tau, eta, weight, dependents)
        for task in candidates:
            tau = await self.get_tau(task.agent_role, agent_id)
            dependents = dependents_counts.get(task.task_id, 0)
            eta = self.compute_eta(
                task, agent_capabilities, free_slots=free_slots, dependents_count=dependents,
            )
            scored.append((task, tau, eta, (tau ** alpha) * (eta ** beta), dependents))

        total = sum(s[3] for s in scored)
        explore_draw = self.rng.random()
        explored = explore_draw < epsilon
        select_pick = None
        if explored or total <= 0:
            select_index = self.rng.randrange(len(scored))
        else:
            select_pick = self.rng.random() * total
            cum = 0.0
            select_index = len(scored) - 1
            for i, s in enumerate(scored):
                cum += s[3]
                if select_pick <= cum:
                    select_index = i
                    break

        task, tau, eta, weight, _dep = scored[select_index]
        candidate_traces = [
            {
                "task_id": s[0].task_id,
                "agent_role": s[0].agent_role,
                "tau": s[1],
                "eta": s[2],
                "weight": s[3],
                "p_norm": (s[3] / total) if total > 0 else 1.0 / len(scored),
                "dependents": s[4],
            }
            for s in scored
        ]
        return Decision(
            task_id=task.task_id,
            agent_id=agent_id,
            agent_role=task.agent_role,
            tau=tau,
            eta=eta,
            p_norm=(weight / total) if total > 0 else 1.0 / len(scored),
            p_score=p_decision(tau, eta, alpha, beta),
            explored=explored,
            alpha=alpha,
            beta=beta,
            epsilon=epsilon,
            seed=self.seed,
            free_slots=free_slots,
            total_weight=total,
            explore_draw=explore_draw,
            select_pick=select_pick,
            select_index=select_index,
            candidates=candidate_traces,
        )

    @staticmethod
    def replay_decision(trace: dict) -> str:
        """Reconstruct the chosen task_id from a stored DecisionTrace — no RNG, no live state.

        Proves the trace is sufficient to explain/replay the selection: same candidates + branch +
        recorded draw → same choice. Mirrors select()'s logic exactly.
        """
        cands = trace["candidates"]
        if trace.get("explored") or trace.get("total_weight", 0.0) <= 0:
            return cands[trace["select_index"]]["task_id"]
        pick = trace["select_pick"]
        cum = 0.0
        for c in cands:
            cum += c["weight"]
            if pick <= cum:
                return c["task_id"]
        return cands[-1]["task_id"]


decision_engine = DecisionEngine()