configs/variants/base.py

"""TeacherVariant — the interface every variant folder implements.

A variant is a self-contained folder under agents/variants/:
    <variant>/utils.py      a TeacherVariant subclass (+ VARIANT = TheSubclass)
    <variant>/prompts.yaml  its maker strategy / system (+ optional teacher wrapper)
    <variant>/README.md     what it is, its bet, its prediction

The SHARED machinery — HTTP transport (ChatClient.score_topk / complete), the MC
loop, the divergence metrics, and the batched training loop — is NOT here and is
never copied. The probe/trainer call the methods below, so adding a variant never
adds an `if variant == ...` branch to core code.

Each variant overrides ONLY what differs; the base is the current live behavior
(reference in the user turn, analyze inside a closed <think>, then derive).
"""

import os
import re
import difflib

import yaml

from agents.utils import (_fill, teacher_user_msg, ANALYSIS_LEADIN,
                          POST_THINK_TRANSITION)
from agents.maker import build_prefix
from agents.verify import is_acceptable_prefix, answer_matches

# Shared maker user-template — variants supply only strategy + system, and inherit
# this fixed scaffold (it is plumbing, not variant logic, so it is not duplicated).
_MAKER_TEMPLATE = yaml.safe_load(
    open(os.path.join(os.path.dirname(__file__), "..", "prompts", "maker_prompts.yaml"))
)["maker_user_template"]


class TeacherVariant:
    """Base = the live 'thinking_on' behavior. Subclasses override selectively."""

    name = "base"
    description = "reference in the user turn; analyze in a closed <think>, then derive"
    thinking = True          # does the teacher analyze inside a <think> block?
    uses_maker = True        # does resolve_prefix call the maker endpoint? (P4: False)

    def __init__(self, variant_dir=None):
        if variant_dir is None:
            import sys
            mod = sys.modules.get(self.__class__.__module__)
            variant_dir = os.path.dirname(getattr(mod, "__file__", "") or ".")
        self.dir = variant_dir
        self.prompts_path = os.path.join(self.dir, "prompts.yaml")
        with open(self.prompts_path) as f:
            self.prompts = yaml.safe_load(f) or {}

    # ------------------------------------------------------------------ #
    # 1. prefix resolution — what fills the reference slot of the context
    # ------------------------------------------------------------------ #
    def _maker_prompts(self):
        return {"strategy": self.prompts["strategy"],
                "maker_system_prompt": self.prompts["maker_system_prompt"],
                "maker_user_template": _MAKER_TEMPLATE}

    def resolve_prefix(self, problem, solution, rollout, maker_client):
        """Maker rewrite of the reference in the student's voice; falls back to the
        reference on any failure. Uses THIS variant's strategy + accept policy."""
        if maker_client is None or not (solution or "").strip():
            return solution
        return build_prefix(problem, solution, rollout, maker_client,
                            prompts=self._maker_prompts(), accept=self.accepts)

    # ------------------------------------------------------------------ #
    # 2. accept gate — is a candidate prefix good enough to use?
    # ------------------------------------------------------------------ #
    def accepts(self, prefix, reference):
        return is_acceptable_prefix(prefix, reference)

    # ------------------------------------------------------------------ #
    # 3. teacher context (probe transport: needs the tokenizer + HTTP client)
    # ------------------------------------------------------------------ #
    def build_ctx(self, tok, client, problem, prefix, cfg):
        wrapper = self.prompts.get("teacher_user")
        user_msg = _fill(wrapper, problem, prefix) if wrapper else teacher_user_msg(problem, prefix)
        base = tok.apply_chat_template(
            [{"role": "user", "content": user_msg}], tokenize=False,
            add_generation_prompt=True, enable_thinking=self.thinking)
        if not self.thinking:
            return base
        lead = self.prompts.get("analysis_leadin", ANALYSIS_LEADIN)
        gen_prompt = base + "<think>\n" + lead + "\n\n"
        out = client.complete(gen_prompt, n=1, temperature=cfg["analysis_temperature"],
                              max_tokens=cfg["analysis_max_tokens"], stop=["</think>"])
        analysis = (out[0] if out else "").split("</think>")[0].rstrip()
        trans = self.prompts.get("post_think_transition", POST_THINK_TRANSITION)
        return gen_prompt + analysis + "\n</think>\n\n" + trans

    # ------------------------------------------------------------------ #
    # 4. error localization — t* (token index into the rollout)
    # ------------------------------------------------------------------ #
    def t_star(self, prefix, rollout, tok):
        return self._difflib_t_star(prefix, rollout, tok)

    # -------- shared t* helpers (a variant picks one) -------- #
    @staticmethod
    def _char_to_tok(text, char_idx, tok):
        if char_idx is None or char_idx < 0:
            return None
        return len(tok(text[:char_idx], add_special_tokens=False)["input_ids"])

    @classmethod
    def _difflib_t_star(cls, prefix, rollout, tok):
        """First block where the rollout stops matching the prefix (approx)."""
        sm = difflib.SequenceMatcher(a=prefix or "", b=rollout or "", autojunk=False)
        for tag, i1, i2, j1, j2 in sm.get_opcodes():
            if tag != "equal":
                return cls._char_to_tok(rollout, j1, tok)
        return None

    @classmethod
    def _lcp_t_star(cls, prefix, rollout, tok):
        """Longest common prefix boundary — the splice point for verbatim-keep r."""
        a, b = prefix or "", rollout or ""
        i = 0
        while i < min(len(a), len(b)) and a[i] == b[i]:
            i += 1
        return cls._char_to_tok(rollout, i, tok)

    @classmethod
    def _quote_t_star(cls, quote, rollout, tok):
        """Locate a quoted first-wrong-sentence inside the rollout (whitespace-tolerant)."""
        if not quote:
            return None
        pat = re.compile(re.escape(quote.strip()).replace(r"\ ", r"\s+"))
        m = pat.search(rollout or "")
        return cls._char_to_tok(rollout, m.start(), tok) if m else None