caudate-cli 0.1.0__py3-none-any.whl

nn/policy.py

@@ -0,0 +1,214 @@
+"""Graduation policy — when does Caudate get to influence the agent?
+
+Five trust levels. Promotion needs both *enough samples* and
+*good-enough accuracy*; demotion is symmetric so a regression after a
+bad retrain doesn't keep her in a level she no longer earns.
+
+    SILENT        no checkpoint loaded                  no predictions
+    OBSERVER      checkpoint loaded                     predicts but ignored
+    WHISPER       accuracy ≥ 0.50, ≥ 50 predictions     prediction → prompt hint
+    ADVISOR       accuracy ≥ 0.70, ≥ 200 predictions    prediction → constrained tool list  (ADR 0007)
+    CONTROLLER    accuracy ≥ 0.80, ≥ 500 predictions    prediction → constrained list + thinking gate
+
+The accuracy used is `composite()` from the scorer — a weighted blend
+of tool / tier / think accuracy, matching the training loss weights.
+
+Persistence: the *current* level is saved alongside the scorecard so a
+restart doesn't reset her to OBSERVER. Re-promotion still requires the
+sample/accuracy gates so a forced demotion (e.g. via /caudate demote)
+holds until earned back.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+from enum import IntEnum
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class TrustLevel(IntEnum):
+    SILENT = 0
+    OBSERVER = 1
+    WHISPER = 2
+    ADVISOR = 3
+    CONTROLLER = 4
+
+    @property
+    def label(self) -> str:
+        return self.name.lower()
+
+
+@dataclass
+class _Gate:
+    level: TrustLevel
+    min_samples: int
+    min_accuracy: float
+
+
+# Promotion thresholds. Demotion uses the previous gate's threshold
+# minus a small hysteresis to avoid flapping.
+_GATES: list[_Gate] = [
+    _Gate(TrustLevel.OBSERVER,   min_samples=0,   min_accuracy=0.0),
+    _Gate(TrustLevel.WHISPER,    min_samples=50,  min_accuracy=0.50),
+    _Gate(TrustLevel.ADVISOR,    min_samples=200, min_accuracy=0.70),
+    _Gate(TrustLevel.CONTROLLER, min_samples=500, min_accuracy=0.80),
+]
+_HYSTERESIS = 0.05  # accuracy must drop this much below a gate to demote
+
+
+class GraduationPolicy:
+    """Owns the current trust level. Reads scorer + advisor state to update."""
+
+    def __init__(self, state_path: Path | str = "data/nn/policy.json"):
+        self.state_path = Path(state_path)
+        self.state_path.parent.mkdir(parents=True, exist_ok=True)
+        self._level: TrustLevel = TrustLevel.SILENT
+        self._frozen: bool = False  # /caudate freeze pins the level
+        self._load()
+
+    # ------------------------------------------------------------------
+
+    @property
+    def level(self) -> TrustLevel:
+        return self._level
+
+    @property
+    def frozen(self) -> bool:
+        return self._frozen
+
+    def can_predict(self) -> bool:
+        return self._level >= TrustLevel.OBSERVER
+
+    def can_whisper(self) -> bool:
+        return self._level >= TrustLevel.WHISPER
+
+    def can_advise(self) -> bool:
+        return self._level >= TrustLevel.ADVISOR
+
+    def can_control(self) -> bool:
+        return self._level >= TrustLevel.CONTROLLER
+
+    # ------------------------------------------------------------------
+    # Updates
+    # ------------------------------------------------------------------
+
+    def update(self, samples: int, composite_acc: float, advisor_loaded: bool) -> TrustLevel:
+        """Recompute the trust level from observed performance.
+
+        Called by the observer after every turn. No-op if frozen.
+        """
+        if self._frozen:
+            return self._level
+
+        if not advisor_loaded:
+            new_level = TrustLevel.SILENT
+        else:
+            new_level = self._level_for(samples, composite_acc, current=self._level)
+
+        if new_level != self._level:
+            old = self._level
+            self._level = new_level
+            self._save()
+            arrow = "↑" if new_level > old else "↓"
+            logger.info(
+                f"Caudate trust level {arrow} {old.label} → {new_level.label} "
+                f"(samples={samples}, composite={composite_acc:.3f})"
+            )
+        return self._level
+
+    def freeze(self, level: TrustLevel | None = None) -> None:
+        """Pin the level (useful for testing or rolling back a bad retrain)."""
+        if level is not None:
+            self._level = level
+        self._frozen = True
+        self._save()
+
+    def thaw(self) -> None:
+        self._frozen = False
+        self._save()
+
+    def force(self, level: TrustLevel) -> None:
+        """Manually set the level. Doesn't freeze — next update can change it."""
+        self._level = level
+        self._save()
+
+    # ------------------------------------------------------------------
+    # Inspection
+    # ------------------------------------------------------------------
+
+    def next_gate(self) -> _Gate | None:
+        """The next promotion gate (or None if already at top)."""
+        for gate in _GATES:
+            if gate.level > self._level:
+                return gate
+        return None
+
+    def progress_to_next(self, samples: int, composite_acc: float) -> dict[str, Any]:
+        """How close is she to the next promotion?"""
+        gate = self.next_gate()
+        if gate is None:
+            return {"at_top": True}
+        return {
+            "next_level": gate.level.label,
+            "samples_needed": max(0, gate.min_samples - samples),
+            "current_samples": samples,
+            "samples_progress": min(1.0, samples / max(1, gate.min_samples)),
+            "accuracy_needed": gate.min_accuracy,
+            "current_accuracy": composite_acc,
+            "accuracy_progress": min(1.0, composite_acc / max(1e-8, gate.min_accuracy)),
+        }
+
+    def report(self, samples: int, composite_acc: float) -> dict[str, Any]:
+        return {
+            "level": self._level.label,
+            "level_int": int(self._level),
+            "frozen": self._frozen,
+            "next": self.progress_to_next(samples, composite_acc),
+        }
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _level_for(
+        self, samples: int, composite_acc: float, current: TrustLevel,
+    ) -> TrustLevel:
+        # Promote: highest gate she clears
+        target = TrustLevel.OBSERVER
+        for gate in _GATES:
+            if samples >= gate.min_samples and composite_acc >= gate.min_accuracy:
+                target = gate.level
+
+        # Demote with hysteresis: she only loses the current level if
+        # accuracy has fallen *below the entry threshold minus hysteresis*.
+        if target < current:
+            current_gate = next((g for g in _GATES if g.level == current), None)
+            if current_gate is not None:
+                if (composite_acc >= current_gate.min_accuracy - _HYSTERESIS
+                        and samples >= current_gate.min_samples):
+                    target = current
+        return target
+
+    def _load(self) -> None:
+        if not self.state_path.exists():
+            return
+        try:
+            data = json.loads(self.state_path.read_text())
+            self._level = TrustLevel(int(data.get("level", 0)))
+            self._frozen = bool(data.get("frozen", False))
+        except Exception as e:
+            logger.debug(f"policy load failed: {e}")
+
+    def _save(self) -> None:
+        try:
+            self.state_path.write_text(json.dumps({
+                "level": int(self._level),
+                "frozen": self._frozen,
+            }, indent=2))
+        except Exception as e:
+            logger.debug(f"policy save failed: {e}")

nn/runtime.py

@@ -0,0 +1,343 @@
+"""Runtime advisor — load Caudate's weights, expose `predict()` to the agent.
+
+The advisor is **non-destructive**. The agent doesn't *have* to follow
+Caudate's predictions. Phase 1 just logs every prediction for
+measurement. Future phases can:
+
+  - bias the LLM's tool prompt with predicted-tool hints
+  - skip the heuristic router and use predicted tier directly
+  - flip thinking on when predicted helpful
+
+Each prediction lands in `data/nn/predictions.jsonl` so we can later
+correlate (predicted, actual, observed reward) and decide whether
+Caudate is good enough to graduate from advisor to controller.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from nn.caudate import Caudate
+from nn.config import NNConfig
+from nn.data import (
+    ConversationSample, ReplayBuffer, SourceVocab, ToolVocab, collate,
+)
+from nn.format import ChatMessage, ToolCall, ToolDef
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Prediction:
+    tool: str
+    tool_confidence: float
+    tier: str                # "fast" | "slow"
+    tier_confidence: float
+    think: float             # in [0, 1]
+    value: float             # in [0, 1]
+    # Top-K candidates (name, probability) sorted high → low. Includes
+    # the predicted tool itself at position 0. Drives ADR-0007 constrained
+    # routing: at ADVISOR+, the LLM's tool list is pruned to this set.
+    # The synthetic "<no_tool>" slot is excluded — pruning only acts on
+    # real-tool predictions per ADR-0007's "soft suggest" decision.
+    tool_top_k: list[tuple[str, float]] = field(default_factory=list)
+
+
+class CaudateAdvisor:
+    """Inference-only wrapper around a trained Caudate."""
+
+    def __init__(
+        self,
+        cfg: NNConfig,
+        model: Caudate,
+        vocab: ToolVocab,
+        source_vocab: SourceVocab | None = None,
+    ):
+        self.cfg = cfg
+        self.model = model.eval()
+        self.vocab = vocab
+        self.source_vocab = source_vocab or SourceVocab(
+            cap=getattr(cfg, "source_vocab_size", SourceVocab.DEFAULT_CAP)
+        )
+        self.device = next(model.parameters()).device
+        # Reverse lookup over the bounded n_tools_out range. Multiple
+        # vocab ids may collide on the same head index because
+        # n_tools_out < vocab size in general; the reverse map keeps the
+        # first one seen, which is fine for explanations.
+        self._head_idx_to_name: dict[int, str] = {}
+        for name, idx in vocab.to_dict().items():
+            slot = idx % cfg.n_tools_out
+            self._head_idx_to_name.setdefault(slot, name)
+
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(
+        self,
+        messages: list[str],
+        tool_history: list[str],
+        mood: list[float],
+        image_paths: list[str] | None = None,
+        model_source: str = "<unknown>",
+        available_tools: list[ToolDef] | None = None,
+    ) -> Prediction:
+        """Run inference. Optional ``available_tools`` enables the
+        contrastive tool head: pass the list of ToolDef the assistant
+        can call right now and the head scores each one. When None or
+        empty, the tool head sits out (no candidates → no prediction)
+        and the returned Prediction has ``tool="<no_tool>"`` with
+        zero confidence.
+
+        ``messages`` are role-prefixed strings ("user: hi"); they're
+        parsed back into ChatMessages internally. ``tool_history`` is
+        the rolling list of tool names called in recent turns —
+        becomes a synthetic assistant message at the end of the
+        prefix so the encoder picks it up via the tool-history
+        channel."""
+        # Parse role-prefixed strings → ChatMessages
+        conv: list[ChatMessage] = []
+        for m in messages:
+            if not m:
+                continue
+            role, sep, content = m.partition(": ")
+            if not sep:
+                role, content = "user", m
+            conv.append(ChatMessage(role=role, content=content))
+        if tool_history:
+            conv.append(ChatMessage(
+                role="assistant",
+                content="",
+                tool_calls=[ToolCall(name=t) for t in tool_history if t],
+            ))
+
+        sample = ConversationSample(
+            conversation=conv,
+            tools=list(available_tools or []),
+            mood=mood,
+            image_paths=image_paths or [],
+            target_tool="<unk>", target_tier=0, target_think=0.0,
+            target_value=0.0,
+            model_source=model_source,
+        )
+        batch = collate(
+            [sample], self.cfg, self.vocab, source_vocab=self.source_vocab,
+        )
+        out = self.model(
+            messages=batch["messages"],
+            tool_ids=batch["tool_ids"].to(self.device),
+            mood=batch["mood"].to(self.device),
+            image_paths=batch.get("image_paths"),
+            source_id=batch["source_id"].to(self.device),
+            tool_specs=batch.get("tool_specs"),
+        )
+
+        # Contrastive tool head — only meaningful when candidates were
+        # supplied. The first candidate slot is always <no_tool>, so a
+        # prediction of "answer directly" is a normal class.
+        tool_top_k: list[tuple[str, float]] = []
+        if "tool_logits" in out:
+            tool_probs = torch.softmax(out["tool_logits"], dim=-1)[0]
+            tool_idx = int(tool_probs.argmax().item())
+            tool_conf = float(tool_probs[tool_idx].item())
+            # Slot 0 is the synthetic <no_tool>; the rest are the
+            # caller's available_tools in supplied order.
+            if tool_idx == 0:
+                tool_name = "<no_tool>"
+            elif available_tools and (tool_idx - 1) < len(available_tools):
+                tool_name = available_tools[tool_idx - 1].name
+            else:
+                tool_name = "<unk>"
+
+            # Top-K real-tool candidates for ADR-0007 constrained routing.
+            # Skip slot 0 (<no_tool>) — pruning acts on real tools only.
+            k = min(self.cfg.advisor_top_k, tool_probs.shape[0] - 1)
+            if available_tools and k > 0:
+                # Indices 1..N map to available_tools[0..N-1]. Score those
+                # slots only, then sort desc.
+                n_real = min(len(available_tools), tool_probs.shape[0] - 1)
+                scored = [
+                    (available_tools[i].name, float(tool_probs[i + 1].item()))
+                    for i in range(n_real)
+                ]
+                scored.sort(key=lambda x: -x[1])
+                tool_top_k = scored[:k]
+        else:
+            tool_name = "<no_tool>"
+            tool_conf = 0.0
+
+        tier_probs = torch.softmax(out["tier_logits"], dim=-1)[0]
+        tier_idx = int(tier_probs.argmax().item())
+        tier_conf = float(tier_probs[tier_idx].item())
+        tier_name = ("fast", "slow")[tier_idx]
+
+        think = float(torch.sigmoid(out["think_logit"])[0].item())
+        value = float(torch.sigmoid(out["value"])[0].item())
+
+        pred = Prediction(
+            tool=tool_name, tool_confidence=tool_conf,
+            tier=tier_name, tier_confidence=tier_conf,
+            think=think, value=value,
+            tool_top_k=tool_top_k,
+        )
+        self._log(pred)
+        return pred
+
+    # ------------------------------------------------------------------
+    # ADR-0006 Phase A — feature-success prediction surface
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict_feature_success(
+        self,
+        feature_text: str,
+        model_id: str | None = None,
+    ) -> dict[str, Any] | None:
+        """Return P(feature succeeds | text, model_source) using the
+        feature_success head if present, else None.
+
+        Returns None — not a fake prediction — when the loaded checkpoint
+        was trained before the head existed (so its weights weren't
+        learned and the output would be noise). The forge_advisor falls
+        back to the heuristic baseline in that case."""
+        # The head was added in nn/heads.py:EXTENDED_HEADS. A checkpoint
+        # trained before this addition will have the module (heads are
+        # constructed from the spec list, not the state-dict) but its
+        # weights will be at Xavier init, not learned. Distinguish by
+        # checking whether the head's weights look trained — a very
+        # cheap proxy is whether the bias on the last linear is non-zero
+        # (Xavier init zeros the bias; even a few training steps move
+        # it). When indistinguishable, the caller treats None as
+        # "use baseline".
+        head_mod = getattr(self.model.heads, "feature_success", None)
+        if head_mod is None:
+            return None
+        last_linear = head_mod[-1] if hasattr(head_mod, "__len__") else None
+        if last_linear is not None and getattr(last_linear, "bias", None) is not None:
+            if torch.allclose(last_linear.bias, torch.zeros_like(last_linear.bias)):
+                # Looks like un-trained init. Return None so the caller
+                # falls back to the heuristic.
+                return None
+
+        sample = ConversationSample(
+            conversation=[ChatMessage(role="user", content=feature_text or "")],
+            tools=[],
+            mood=[0.5, 0.5, 0.5, 0.5],
+            image_paths=[],
+            target_tool="<unk>", target_tier=0, target_think=0.0,
+            target_value=0.0,
+            model_source=model_id or "<unknown>",
+        )
+        batch = collate(
+            [sample], self.cfg, self.vocab, source_vocab=self.source_vocab,
+        )
+        out = self.model(
+            messages=batch["messages"],
+            tool_ids=batch["tool_ids"].to(self.device),
+            mood=batch["mood"].to(self.device),
+            image_paths=batch.get("image_paths"),
+            source_id=batch["source_id"].to(self.device),
+        )
+        logit = out.get("feature_success_logit")
+        if logit is None:
+            return None
+        # (B,) tensor → scalar
+        p = float(torch.sigmoid(logit.detach())[0].item())
+        return {
+            "success_prob": max(0.0, min(1.0, p)),
+            "reason": f"caudate head · model={model_id or '<unknown>'}",
+        }
+
+    # ------------------------------------------------------------------
+
+    def _log(self, pred: Prediction) -> None:
+        try:
+            path = Path(self.cfg.advisor_log_path)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            with path.open("a") as f:
+                f.write(json.dumps({
+                    "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                    "tool": pred.tool, "tool_conf": round(pred.tool_confidence, 4),
+                    "tier": pred.tier, "tier_conf": round(pred.tier_confidence, 4),
+                    "think": round(pred.think, 4),
+                    "value": round(pred.value, 4),
+                }) + "\n")
+        except Exception as e:
+            logger.debug(f"advisor log write failed: {e}")
+
+
+def load_advisor(cfg: NNConfig | None = None) -> CaudateAdvisor | None:
+    """Load Caudate's checkpoint; return None if absent or mismatched.
+
+    Safe to call at every agent startup — failure is silent so the
+    agent works fine even before any training has happened.
+
+    Architecture detection: if a `caudate.meta.json` sidecar exists
+    next to the checkpoint, its `config` block overrides the default
+    architecture knobs (d_model, n_layers, n_heads, d_ff). This lets a
+    retrained checkpoint at a different scale (e.g. Phase-5 20M / 95M)
+    load cleanly without the loader hard-coding the old shape.
+    """
+    cfg = cfg or NNConfig()
+    ckpt_path = Path(cfg.checkpoint_path)
+    if not ckpt_path.exists():
+        return None
+    try:
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    except Exception as e:
+        logger.warning(f"Caudate checkpoint load failed: {e}")
+        return None
+
+    # Prefer meta-file config over default — handles scaled retrains.
+    meta_path = Path(getattr(cfg, "metadata_path", "data/nn/caudate.meta.json"))
+    if meta_path.exists():
+        try:
+            import json
+            meta = json.loads(meta_path.read_text())
+            arch = meta.get("config") or {}
+            arch_keys = ("d_model", "n_heads", "n_layers", "d_ff", "dropout",
+                         "n_tools_out", "source_vocab_size", "tool_vocab_size",
+                         "tool_embed_dim", "mood_dim", "history_window",
+                         "msg_window", "image_window", "use_vision",
+                         "vision_backend", "vision_embed_dim", "vision_dtype",
+                         "vision_encoder_name", "text_embed_dim")
+            overrides = {k: arch[k] for k in arch_keys if k in arch}
+            if overrides:
+                cfg = cfg.model_copy(update=overrides)
+                logger.info(
+                    f"Caudate loading with meta-overridden arch: "
+                    f"d_model={cfg.d_model} n_layers={cfg.n_layers}"
+                )
+        except Exception as e:
+            logger.warning(f"meta override failed, using defaults: {e}")
+
+    try:
+        vocab = ToolVocab.from_dict(ckpt.get("vocab", {}))
+        source_vocab = SourceVocab.from_dict(
+            ckpt.get("source_vocab", {}),
+            cap=getattr(cfg, "source_vocab_size", SourceVocab.DEFAULT_CAP),
+        )
+        model = Caudate(cfg)
+        # strict=False — pre-Phase-2 checkpoints have no source_embed
+        # weights; the Caudate constructor zero-inits them so the
+        # baseline behaviour is preserved until the next retrain.
+        missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
+        if missing or unexpected:
+            logger.info(
+                f"advisor checkpoint loaded non-strict  "
+                f"missing={list(missing)[:5]}  unexpected={list(unexpected)[:5]}"
+            )
+    except Exception as e:
+        logger.warning(f"Caudate model rebuild failed: {e}")
+        return None
+    return CaudateAdvisor(cfg=cfg, model=model, vocab=vocab, source_vocab=source_vocab)
+
+
+# Backward-compat alias — old code that imported NNAdvisor still works.
+NNAdvisor = CaudateAdvisor