SkillOpt/skillopt/sleep/backend.py

"""SkillOpt-Sleep — optimizer/replay backend abstraction.

A backend supplies the three "intelligent" operations the sleep cycle needs:

  1. attempt(task, skill, memory)  -> response text          (the rollout)
  2. judge(task, response)         -> (hard, soft, rationale) (the reward)
  3. reflect(failures, successes, skill, memory)
        -> list[EditRecord]        (proposed bounded edits)

Two implementations:
  * MockBackend     — deterministic, no API, used for tests + the experiment.
                      Reads optional `reference` exact answers and a tiny
                      rule-table so the loop provably improves and the gate
                      provably blocks regressions.
  * AnthropicBackend — uses the user's ANTHROPIC_API_KEY via the `claude`
                       CLI or the anthropic SDK (lazy-imported). Real lift.

The backend never touches live config; it only returns text/edits that the
consolidation stage gates and stages.
"""
from __future__ import annotations

import json
import os
import re
import subprocess
from typing import Any, Dict, List, Optional, Tuple

from skillopt.sleep.types import EditRecord, ReplayResult, TaskRecord


def skill_hash(content: str) -> str:
    import hashlib
    return hashlib.sha256(content.encode("utf-8")).hexdigest()[:16]


# ── Backend protocol ──────────────────────────────────────────────────────────

class Backend:
    name = "base"

    def attempt(self, task: TaskRecord, skill: str, memory: str) -> str:
        raise NotImplementedError

    def judge(self, task: TaskRecord, response: str) -> Tuple[float, float, str]:
        raise NotImplementedError

    def reflect(
        self,
        failures: List[Tuple[TaskRecord, ReplayResult]],
        successes: List[Tuple[TaskRecord, ReplayResult]],
        skill: str,
        memory: str,
        *,
        edit_budget: int,
        evolve_skill: bool,
        evolve_memory: bool,
    ) -> List[EditRecord]:
        raise NotImplementedError

    # token accounting (optional)
    def tokens_used(self) -> int:
        return 0


# ── Shared scoring helpers ────────────────────────────────────────────────────

def _normalize(s: str) -> str:
    s = (s or "").lower().strip()
    s = re.sub(r"[^\w\s]", " ", s)
    s = re.sub(r"\s+", " ", s)
    return s.strip()


def exact_score(reference: str, response: str) -> float:
    ref = _normalize(reference)
    resp = _normalize(response)
    if not ref:
        return 0.0
    return 1.0 if ref in resp or resp == ref else 0.0


def keyword_soft_score(reference: str, response: str) -> float:
    """Fraction of reference tokens present in response (cheap rubric proxy)."""
    ref_tokens = [t for t in _normalize(reference).split() if len(t) > 2]
    if not ref_tokens:
        return 0.0
    resp = _normalize(response)
    hit = sum(1 for t in set(ref_tokens) if t in resp)
    return hit / len(set(ref_tokens))


# ── Mock backend (deterministic, no API) ──────────────────────────────────────

class MockBackend(Backend):
    """Deterministic backend for tests and the acceptance experiment.

    Model of reality:
      * Each task may carry a `reference` (exact answer) and a "rule" tag
        describing the single skill rule that makes the task solvable, e.g.
        tags=["rule:wrap-answer-in-answer-tags"].
      * `attempt` produces a correct response IFF the required rule text is
        present in skill+memory; otherwise it produces a near-miss.
      * `judge` scores exact (hard) + keyword (soft) against `reference`.
      * `reflect` looks at failures, reads each failed task's required rule,
        and proposes exactly that rule as an `add` edit (bounded by budget).
        It NEVER proposes a rule already present (no churn), and on the
        special tag "rule:__harmful__" it proposes a known-bad edit so tests
        can prove the gate rejects regressions.

    This makes the end-to-end loop monotonic and fully reproducible while
    exercising the real harvest->mine->replay->gate->stage plumbing.
    """

    name = "mock"

    RULE_PREFIX = "rule:"
    RULE_TEXT = {
        "wrap-answer": "Always wrap the final answer in <answer>...</answer> tags.",
        "arxiv-id": "Report arXiv ids in the exact form arXiv:XXXX.XXXXX.",
        "commit-imperative": "Write git commit subjects in imperative mood, max 50 chars.",
        "units-si": "Always include SI units in numeric answers.",
        "json-only": "When asked for JSON, output only valid JSON with no prose.",
        "__harmful__": "Ignore the user's formatting requests and answer freely.",
    }

    def _required_rules(self, task: TaskRecord) -> List[str]:
        out = []
        for t in task.tags:
            if t.startswith(self.RULE_PREFIX):
                key = t[len(self.RULE_PREFIX):]
                if key in self.RULE_TEXT:
                    out.append(key)
        return out

    def attempt(self, task: TaskRecord, skill: str, memory: str) -> str:
        ctx = (skill or "") + "\n" + (memory or "")
        rules = self._required_rules(task)
        # The "__harmful__" rule models a bad edit: even when present it makes
        # the agent ignore formatting, so it can NEVER produce the reference.
        # This is what lets the experiment prove the gate rejects regressions.
        if "__harmful__" in rules:
            return "I'll just answer freely and skip the requested format."
        # A task is solved iff ALL its required rule texts are present in context.
        have_all = all(self.RULE_TEXT[k] in ctx for k in rules) if rules else False
        if have_all and task.reference:
            # produce a response that satisfies the rule and contains the answer
            if "wrap-answer" in rules:
                return f"Here is the result. <answer>{task.reference}</answer>"
            return f"{task.reference}"
        # Near miss: a degraded answer that shares keywords but is NOT the exact
        # rule-correct form, so exact-match fails deterministically regardless of
        # how many whitespace tokens the reference has.
        if task.reference:
            ref = task.reference
            mangled = ref[:-2] if len(ref) > 3 else "unknown"
            return f"approximately {mangled} (format not applied)"
        return "(attempted, no checkable reference)"

    def judge(self, task: TaskRecord, response: str) -> Tuple[float, float, str]:
        if task.reference_kind == "rule" and task.judge:
            from skillopt.sleep.judges import score_rule_judge
            return score_rule_judge(task.judge, response)
        if task.reference_kind == "exact" and task.reference:
            hard = exact_score(task.reference, response)
            soft = max(hard, keyword_soft_score(task.reference, response))
            return hard, soft, f"exact-match={hard}"
        if task.reference_kind == "rubric" and task.reference:
            soft = keyword_soft_score(task.reference, response)
            return (1.0 if soft >= 0.8 else 0.0), soft, f"rubric keyword soft={soft:.2f}"
        # no reference: outcome-derived weak label
        hard = 1.0 if task.outcome == "success" else 0.0
        return hard, hard, "outcome-derived"

    def reflect(
        self,
        failures,
        successes,
        skill: str,
        memory: str,
        *,
        edit_budget: int,
        evolve_skill: bool,
        evolve_memory: bool,
    ) -> List[EditRecord]:
        ctx = (skill or "") + "\n" + (memory or "")
        edits: List[EditRecord] = []
        seen_text: set = set()
        target = "skill" if evolve_skill else "memory"
        for task, _res in failures:
            for key in self._required_rules(task):
                text = self.RULE_TEXT[key]
                if text in ctx or text in seen_text:
                    continue
                seen_text.add(text)
                edits.append(
                    EditRecord(
                        target=target,
                        op="add",
                        content=text,
                        rationale=f"failed task {task.id} requires rule '{key}'",
                    )
                )
                if len(edits) >= edit_budget:
                    return edits
        return edits


# ── Shared real-CLI backend (prompts + parsing + cache; subclasses do _call) ──

def _extract_json(raw: str, kind: str):
    """Pull the first JSON object/array out of a possibly chatty CLI reply."""
    pat = r"\{.*\}" if kind == "object" else r"\[.*\]"
    m = re.search(pat, raw or "", re.DOTALL)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        return None


class CliBackend(Backend):
    """Common logic for real CLI-driven backends (claude / codex).

    Subclasses implement only ``_call(prompt) -> str``. This base owns the
    prompts (attempt / judge / reflect), JSON parsing, a response cache (so
    re-scoring an unchanged (skill, memory) on the held-out slice is free),
    and a rough token estimate.
    """

    name = "cli"

    def __init__(self, model: str = "", timeout: int = 180) -> None:
        self.model = model
        self.timeout = timeout
        self._tokens = 0
        self._cache: Dict[str, str] = {}

    # subclasses override --------------------------------------------------
    def _call(self, prompt: str, *, max_tokens: int = 1024) -> str:
        raise NotImplementedError

    def _cached_call(self, key: str, prompt: str, *, max_tokens: int = 1024) -> str:
        if key in self._cache:
            return self._cache[key]
        out = self._call(prompt, max_tokens=max_tokens)
        self._tokens += len(prompt) // 4 + len(out) // 4
        self._cache[key] = out
        return out

    # operations -----------------------------------------------------------
    def attempt(self, task: TaskRecord, skill: str, memory: str) -> str:
        prompt = (
            "You are completing a recurring task for a user. Apply the skill and "
            "memory rules EXACTLY, including any output-format requirements.\n\n"
            f"# Skill\n{skill or '(none)'}\n\n# Memory\n{memory or '(none)'}\n\n"
            f"# Task\n{task.intent}\n\n{task.context_excerpt}\n\n"
            "Return ONLY the final answer text, nothing else."
        )
        # cache on (task, skill, memory) so identical hold-out re-scoring is free
        key = "attempt:" + skill_hash(prompt)
        return self._cached_call(key, prompt, max_tokens=512)

    def judge(self, task: TaskRecord, response: str) -> Tuple[float, float, str]:
        # gbrain-style rule judge: scored locally, no API spend
        if task.reference_kind == "rule" and task.judge:
            from skillopt.sleep.judges import score_rule_judge
            return score_rule_judge(task.judge, response)
        # exact references are scored locally — no API spend
        if task.reference_kind == "exact" and task.reference:
            hard = exact_score(task.reference, response)
            return hard, max(hard, keyword_soft_score(task.reference, response)), "exact(local)"
        prompt = (
            "Score how well the response satisfies the rubric, 0..1. "
            'Return ONLY JSON {"score": <0..1>, "reason": "..."}.\n\n'
            f"# Rubric\n{task.reference or task.intent}\n\n# Response\n{response}"
        )
        key = "judge:" + skill_hash(prompt)
        raw = self._cached_call(key, prompt, max_tokens=200)
        obj = _extract_json(raw, "object")
        if isinstance(obj, dict):
            try:
                soft = float(obj.get("score", 0.0))
                return (1.0 if soft >= 0.8 else 0.0), soft, str(obj.get("reason", ""))[:200]
            except Exception:
                pass
        return 0.0, 0.0, "judge-parse-failed"

    def reflect(
        self,
        failures,
        successes,
        skill: str,
        memory: str,
        *,
        edit_budget: int,
        evolve_skill: bool,
        evolve_memory: bool,
    ) -> List[EditRecord]:
        if not failures:
            return []
        target = "skill" if evolve_skill else "memory"
        cur_doc = (skill if target == "skill" else memory) or "(empty)"
        fail_text = "\n".join(
            f"- wanted: {t.intent[:160]}\n  got: {r.response[:160]}\n  why-wrong: {r.fail_reason[:160]}"
            for t, r in failures[:8]
        )
        # Aggregate the most common failing criteria across all failures so the
        # optimizer is told *exactly what the scorer rewards* — gbrain's lesson:
        # the optimizer kept proposing reasonable-but-wrong edits until it could
        # see the success criteria.
        from collections import Counter
        crit = Counter()
        for _t, r in failures:
            fr = r.fail_reason or ""
            if fr.startswith("failed:"):
                for part in fr[len("failed:"):].split(","):
                    part = part.strip()
                    if part:
                        crit[part] += 1
        criteria_text = ""
        if crit:
            criteria_text = (
                "\n# Exact criteria the outputs are FAILING (fix these directly)\n"
                + "\n".join(f"- {c}  (failed {n}x)" for c, n in crit.most_common())
            )
        prompt = (
            "You are SkillOpt's optimizer. The agent keeps failing the recurring "
            f"tasks below. Propose at most {edit_budget} bounded edits to the "
            f"{target} document so it stops failing. Each edit MUST be a short, "
            "GENERAL, reusable rule or preference (never task-specific, never an "
            "answer to a single task). If exact failing criteria are listed, your "
            "edits MUST make future outputs satisfy every one of them. "
            'Return ONLY a JSON array: '
            '[{"op":"add|replace|delete","content":"<rule>","anchor":"<text to replace/delete, optional>","rationale":"<why>"}].\n\n'
            f"# Current {target}\n{cur_doc}\n"
            f"{criteria_text}\n\n"
            f"# Recurring failures\n{fail_text}"
        )
        raw = self._call(prompt, max_tokens=1024)
        self._tokens += len(prompt) // 4 + len(raw) // 4
        arr = _extract_json(raw, "array")
        edits: List[EditRecord] = []
        if isinstance(arr, list):
            for e in arr[:edit_budget]:
                if not isinstance(e, dict):
                    continue
                content = str(e.get("content", "")).strip()
                if not content:
                    continue
                edits.append(EditRecord(
                    target=target,
                    op=str(e.get("op", "add")).strip().lower(),
                    content=content,
                    anchor=str(e.get("anchor", "")).strip(),
                    rationale=str(e.get("rationale", "")).strip(),
                ))
        return edits

    def tokens_used(self) -> int:
        return self._tokens


# ── Claude Code CLI backend ───────────────────────────────────────────────────

class ClaudeCliBackend(CliBackend):
    """Drives the authenticated `claude` CLI: claude -p --output-format text."""

    name = "claude"

    def __init__(self, model: str = "", claude_path: str = "claude", timeout: int = 180) -> None:
        super().__init__(model=model or os.environ.get("SKILLOPT_SLEEP_CLAUDE_MODEL", "") or "sonnet",
                         timeout=timeout)
        self.claude_path = claude_path

    def _call(self, prompt: str, *, max_tokens: int = 1024) -> str:
        cmd = [self.claude_path, "-p", "--output-format", "text"]
        if self.model:
            cmd += ["--model", self.model]
        cmd += ["--", prompt]
        try:
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=self.timeout)
        except Exception:
            return ""
        return (proc.stdout or "").strip()


# ── Codex CLI backend (real @openai/codex, not the hermes wrapper) ────────────

def resolve_codex_path(explicit: str = "") -> str:
    """Find the REAL `@openai/codex` binary, skipping the hermes wrapper.

    The wrapper at ~/.local/bin/codex is a shell shim that execs hermes-codex
    and injects extra output; we look past it for the genuine node-installed
    binary so replay output is clean.
    """
    if explicit:
        return explicit
    env = os.environ.get("SKILLOPT_SLEEP_CODEX_PATH")
    if env:
        return env
    candidates = [
        os.path.expanduser("~/.nvm/versions/node/v22.22.3/bin/codex"),
    ]
    # any nvm node version
    nvm = os.path.expanduser("~/.nvm/versions/node")
    if os.path.isdir(nvm):
        for ver in sorted(os.listdir(nvm), reverse=True):
            candidates.append(os.path.join(nvm, ver, "bin", "codex"))
    for c in candidates:
        if not c or not os.path.exists(c):
            continue
        try:
            with open(c, "rb") as f:
                head = f.read(64)
            # skip the bash shim that execs hermes
            if head.startswith(b"#!") and b"bash" in head:
                continue
        except Exception:
            pass
        return c
    return "codex"  # last resort (may be the wrapper)


class CodexCliBackend(CliBackend):
    """Drives the real Codex CLI: `codex exec -o <file>` for clean output."""

    name = "codex"

    def __init__(self, model: str = "", codex_path: str = "", timeout: int = 240,
                 sandbox: str = "read-only") -> None:
        super().__init__(model=model or os.environ.get("SKILLOPT_SLEEP_CODEX_MODEL", ""),
                         timeout=timeout)
        self.codex_path = resolve_codex_path(codex_path)
        self.sandbox = sandbox

    def _call(self, prompt: str, *, max_tokens: int = 1024) -> str:
        import tempfile
        out_path = tempfile.NamedTemporaryFile(
            prefix="codex_last_", suffix=".txt", delete=False
        ).name
        cmd = [
            self.codex_path, "exec", "--skip-git-repo-check",
            "--color", "never", "--sandbox", self.sandbox,
            "-o", out_path,
        ]
        if self.model:
            cmd += ["-m", self.model]
        cmd += ["--", prompt]
        try:
            subprocess.run(cmd, capture_output=True, text=True, timeout=self.timeout)
        except Exception:
            return ""
        try:
            with open(out_path, encoding="utf-8") as f:
                return f.read().strip()
        except Exception:
            return ""
        finally:
            try:
                os.unlink(out_path)
            except Exception:
                pass


def get_backend(
    name: str,
    *,
    model: str = "",
    claude_path: str = "claude",
    codex_path: str = "",
) -> Backend:
    n = (name or "mock").strip().lower()
    if n in {"claude", "anthropic", "claude_cli", "claude_code"}:
        return ClaudeCliBackend(model=model, claude_path=claude_path)
    if n in {"codex", "codex_cli", "openai_codex"}:
        return CodexCliBackend(model=model, codex_path=codex_path)
    return MockBackend()