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:
    """One sampled assignment, with everything the benchmark layer needs."""
    task_id: str
    agent_id: str
    agent_role: str
    tau: float
    eta: float
    p_norm: float    # normalized selection probability over this candidate set
    p_score: float   # standard §3.3 P_decision = τ^α·η^β·100 (what the metric reports)
    explored: bool   # True when the ε-greedy branch picked uniformly

    def telemetry(self) -> dict:
        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,
        }


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

    def __init__(self, rng: Optional[random.Random] = None):
        seed = os.getenv("ACO_SEED")
        self.rng = rng or (random.Random(int(seed)) if seed 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 = []
        for task in candidates:
            tau = await self.get_tau(task.agent_role, agent_id)
            eta = self.compute_eta(
                task, agent_capabilities, free_slots=free_slots,
                dependents_count=dependents_counts.get(task.task_id, 0),
            )
            scored.append((task, tau, eta, (tau ** alpha) * (eta ** beta)))

        total = sum(w for *_xs, w in scored)
        explored = self.rng.random() < epsilon
        if explored or total <= 0:
            task, tau, eta, weight = scored[self.rng.randrange(len(scored))]
        else:
            pick = self.rng.random() * total
            cum = 0.0
            task, tau, eta, weight = scored[-1]
            for cand in scored:
                cum += cand[3]
                if pick <= cum:
                    task, tau, eta, weight = cand
                    break
        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,
        )


decision_engine = DecisionEngine()