Part E: Programmatic prompting

"""Part E: prompts as versioned code, automatic prompt search against the dev set, and meta-prompting.

The "model" that reads each triage prompt is a deterministic stand-in
(keyword_reader, through ScriptedLLM): it follows the prompt's category
definitions literally. That makes the search loop run for real here, with real
dev and test scores for THIS reader. With a key, pass --live and the same loop
scores candidates with supportdesk.llm.chat instead (48 calls per candidate).
"""
from __future__ import annotations

import hashlib
import json
import sys
import time
from dataclasses import dataclass, field, replace
from pathlib import Path

from examples.m05_helpers import TRIAGE_KEYWORDS, keyword_responder, parse_category, triage_messages, triage_prompt, wilson
from supportdesk.data import CATEGORIES, Ticket, load_tickets
from supportdesk.kb_search import tokenize
from supportdesk.llm import Usage
from supportdesk.pricing import cost_usd
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_messages

RUN_LOG = Path("runs/m05_prompt_search.jsonl")


@dataclass(frozen=True)
class PromptVersion:
    """A prompt as data: its parts are fields, the text is rendered, and the fingerprint names it."""
    name: str
    keywords: tuple[tuple[str, tuple[str, ...]], ...]
    fallback: str = "how_to"
    parent: str | None = None
    note: str = ""

    @classmethod
    def from_dict(cls, name: str, keywords: dict[str, list[str]], **kw) -> "PromptVersion":
        return cls(name, tuple((c, tuple(ws)) for c, ws in keywords.items()), **kw)

    def as_dict(self) -> dict[str, list[str]]:
        return {c: list(ws) for c, ws in self.keywords}

    @property
    def text(self) -> str:
        return triage_prompt(self.as_dict(), self.fallback)

    @property
    def fingerprint(self) -> str:
        return hashlib.sha256(self.text.encode()).hexdigest()[:12]


def validate(version: PromptVersion) -> list[str]:
    """Problems that make a candidate unusable, whoever wrote it (a person, the search, or a model)."""
    problems = []
    names = [c for c, _ in version.keywords]
    missing = [c for c in CATEGORIES if c not in names]
    if missing:
        problems.append(f"missing categories {missing}")
    if extra := [c for c in names if c not in CATEGORIES]:
        problems.append(f"unknown categories {extra}")
    if version.fallback not in CATEGORIES:
        problems.append(f"fallback {version.fallback!r} is not a category")
    if len(version.text) > 2000:
        problems.append("prompt longer than 2,000 characters")
    return problems


@dataclass
class Score:
    correct: int
    n: int
    calls: int
    failures: list[tuple[Ticket, str | None]] = field(default_factory=list)

    @property
    def acc(self) -> float:
        return self.correct / self.n


def score(chat_fn, version: PromptVersion, tickets: list[Ticket]) -> Score:
    """Run the prompt over labelled tickets and count exact category matches."""
    s = Score(0, len(tickets), 0)
    for t in tickets:
        predicted = parse_category(chat_fn(triage_messages(version.text, t.text), temperature=0.0).text)
        s.calls += 1
        if predicted == t.gold["category"]:
            s.correct += 1
        else:
            s.failures.append((t, predicted))
    return s


def log_run(version: PromptVersion, split: str, s: Score) -> None:
    RUN_LOG.parent.mkdir(exist_ok=True)
    with RUN_LOG.open("a") as f:
        f.write(json.dumps({"prompt": version.name, "fingerprint": version.fingerprint, "parent": version.parent,
                            "split": split, "correct": s.correct, "n": s.n, "note": version.note}) + "\n")


# Proposers: turn a version and its failures into candidate versions -----------------------

def propose_edits(version: PromptVersion, failures: list[tuple[Ticket, str | None]], words_per_ticket: int = 3) -> list[PromptVersion]:
    """Deterministic proposer: for each failure, add a word from the ticket to the gold category,
    or remove the keyword that pulled it into the wrong category."""
    kw = version.as_dict()
    used = {w for ws in kw.values() for w in ws}
    seen, out = set(), []
    for ticket, predicted in failures:
        gold = ticket.gold["category"]
        words = [w for w in tokenize(ticket.text) if len(w) >= 4 and w.isascii() and w not in used]
        for w in list(dict.fromkeys(words))[:words_per_ticket]:
            if (gold, "+", w) not in seen:
                seen.add((gold, "+", w))
                new = {c: ws + [w] if c == gold else ws for c, ws in kw.items()}
                out.append(PromptVersion.from_dict("", new, fallback=version.fallback, note=f"add '{w}' to {gold}"))
        if predicted in kw:
            for w in kw[predicted]:
                if w in ticket.text.lower() and (predicted, "-", w) not in seen:
                    seen.add((predicted, "-", w))
                    new = {c: [x for x in ws if x != w] if c == predicted else ws for c, ws in kw.items()}
                    out.append(PromptVersion.from_dict("", new, fallback=version.fallback, note=f"remove '{w}' from {predicted}"))
    for c in CATEGORIES:
        if c != version.fallback:
            out.append(replace(version, fallback=c, note=f"fallback {c}"))
    return out


def search(chat_fn, start: PromptVersion, dev: list[Ticket], rounds: int = 8, min_gain: int = 1,
           holdout: list[Ticket] | None = None) -> tuple[PromptVersion, list[str], int]:
    """Greedy hill climbing: score every candidate on dev, keep the best if it gains at least min_gain.

    With a holdout set, each kept version is also scored there; the search stops if the holdout
    score drops, and returns the EARLIEST version with the best holdout score (the simplest prompt
    that the unseen tickets support), not the last one.
    """
    best, best_score = start, score(chat_fn, start, dev)
    calls, history = best_score.calls, [f"round 0  {best_score.correct}/{len(dev)}  {start.name} (start)"]
    held = score(chat_fn, start, holdout).correct if holdout else None
    kept = [(held, 0, start)]
    if holdout:
        calls += len(holdout)
        history[0] += f"  holdout {held}/{len(holdout)}"
    for r in range(1, rounds + 1):
        candidates = [c for c in propose_edits(best, best_score.failures) if not validate(c)]
        scored = []
        for c in candidates:
            s = score(chat_fn, c, dev)
            calls += s.calls
            scored.append((s.correct, -len(c.text), c, s))    # ties go to the shorter prompt
        top_correct, _, top, top_score = max(scored, key=lambda x: x[:2])
        if top_correct - best_score.correct < min_gain:
            history.append(f"round {r}  no candidate gains {min_gain}+ of {len(candidates)} tried; stop")
            break
        line = f"round {r}  {top_correct}/{len(dev)}  {len(candidates):>3} candidates  kept: {top.note}"
        if holdout:
            h = score(chat_fn, top, holdout).correct
            calls += len(holdout)
            line += f"  holdout {h}/{len(holdout)}"
            if h < held:
                history.append(line + "  holdout dropped; stop and keep the previous version")
                break
            held = h
        best = replace(top, name=f"{start.name}.s{r}", parent=best.name)
        best_score = top_score
        kept.append((held, r, best))
        history.append(line)
    if holdout:
        top_held = max(h for h, _, _ in kept)
        choice = min((r, v) for h, r, v in kept if h == top_held)
        history.append(f"returning round {choice[0]} ({choice[1].name}): earliest version with the best holdout score {top_held}/{len(holdout)}")
        return choice[1], history, calls
    return best, history, calls


# Meta-prompting: a model proposes the next version ----------------------------------------

META = ("You improve a ticket-classification prompt. Below is the current prompt and tickets it got wrong, "
        "with the correct category. Rewrite ONLY the category lines, one per category, in the same format "
        "'- category: word, word, ...'. Keep all six categories. Return the lines and nothing else.\n\n"
        "<current_prompt>\n{prompt}\n</current_prompt>\n<failures>\n{failures}\n</failures>")


def meta_messages(version: PromptVersion, failures: list[tuple[Ticket, str | None]], k: int = 8) -> list[dict]:
    lines = "\n".join(f"- ticket: {t.subject!r} | predicted: {p} | correct: {t.gold['category']}" for t, p in failures[:k])
    return [{"role": "user", "content": META.format(prompt=version.text, failures=lines)}]


def parse_meta(text: str, base: PromptVersion, name: str) -> PromptVersion:
    """Read '- category: words' lines from a model reply into a new version (then validate it)."""
    kw: dict[str, list[str]] = {}
    for line in text.splitlines():
        if line.startswith("- ") and ":" in line:
            cat, words = line[2:].split(":", 1)
            kw[cat.strip()] = [w.strip().lower() for w in words.split(",") if w.strip()]
    return PromptVersion.from_dict(name, kw, fallback=base.fallback, parent=base.name, note="meta-prompt proposal")


def scripted_meta(messages, kwargs):
    """ScriptedLLM stand-in for the meta-prompt call (plumbing only): it adds each failure's first
    subject word to the correct category, and, to show why validation exists, drops feature_request."""
    content = messages[-1]["content"]
    current = dict(line[2:].split(": ", 1) for line in content.splitlines() if line.startswith("- ") and ": " in line
                   and not line.startswith("- ticket"))
    for line in content.split("<failures>\n", 1)[1].splitlines():
        if line.startswith("- ticket"):
            subject = line.split("'")[1].lower()
            correct = line.rsplit("correct: ", 1)[1]
            current[correct] += ", " + subject.split()[0]
    current.pop("feature_request")
    return "\n".join(f"- {c}: {w}" for c, w in current.items())


if __name__ == "__main__":
    live = "--live" in sys.argv
    if live:
        from supportdesk.llm import chat as reader
    else:
        reader = ScriptedLLM(responder=keyword_responder)
    dev, test = load_tickets("dev"), load_tickets("test")
    v1 = PromptVersion.from_dict("triage-kw-v1", TRIAGE_KEYWORDS, note="hand-written keyword definitions")

    print(f"{v1.name} fingerprint {v1.fingerprint}, {len(v1.text)} characters, problems: {validate(v1) or 'none'}")
    started = time.perf_counter()
    best, history, calls = search(reader, v1, dev)
    print("\n".join(history))
    elapsed = time.perf_counter() - started
    tokens = sum(count_messages(c["messages"]) for c in getattr(reader, "calls", []))  # stand-in keeps its calls
    usd = cost_usd(Usage(input_tokens=tokens), "openai/gpt-oss-120b")
    print(f"search: {calls} reader calls, {tokens:,} input tokens (o200k estimate), "
          f"{usd:.2f} USD of input at gpt-oss-120b prices if they were real calls; {elapsed:.1f} s\n")

    # Same search, but it only sees 32 dev tickets and early-stops on the other 16.
    early, history2, calls2 = search(reader, replace(v1, name="triage-kw-v1e"), dev[:32], holdout=dev[32:])
    print("\n".join(history2))
    print(f"search with holdout: {calls2} reader calls\n")

    print(f"{'version':<16} {'fingerprint':<13} {'dev (n=48)':<22} {'test (n=24)':<22}")
    for v in dict.fromkeys((v1, best, early)):
        d, t = score(reader, v, dev), score(reader, v, test)
        log_run(v, "dev", d)
        log_run(v, "test", t)
        cells = [f"{s.correct}/{s.n} ({wilson(s.correct, s.n)[0]:.2f}-{wilson(s.correct, s.n)[1]:.2f})" for s in (d, t)]
        print(f"{v.name:<16} {v.fingerprint:<13} {cells[0]:<22} {cells[1]:<22}")
    added = {c: [w for w in ws if w not in dict(v1.as_dict()).get(c, [])] for c, ws in best.as_dict().items()}
    print("words the search added:", {c: ws for c, ws in added.items() if ws})

    # Meta-prompting: ask a model to write the next version; validate before scoring.
    meta_llm = reader if live else ScriptedLLM(responder=scripted_meta)
    failures = score(reader, v1, dev).failures
    proposal = parse_meta(meta_llm(meta_messages(v1, failures), temperature=0.7).text, v1, "triage-kw-meta1")
    problems = validate(proposal)
    print(f"\nmeta-prompt proposal {proposal.fingerprint}: problems {problems or 'none'}")
    if problems and not live:
        repaired = PromptVersion.from_dict("triage-kw-meta1r", {**proposal.as_dict(), "feature_request": list(dict(v1.keywords)["feature_request"])},
                                           parent=v1.name, note="meta proposal with the dropped category restored")
        proposal = repaired
        print(f"restored the dropped category from the parent: problems {validate(proposal) or 'none'}")
    if not validate(proposal):
        d, t = score(reader, proposal, dev), score(reader, proposal, test)
        print(f"{proposal.name:<16} dev {d.correct}/{d.n}  test {t.correct}/{t.n}")

PYTHONPATH=. python examples/m05_prompt_search.py

"""Module 5 lab: one pipeline over the week's tickets and the 20 refund scenarios.

1. Route all 72 tickets (router prompt, conditional paths).
2. Refund decisions: sample up to N answers with the scaffold prompt, accept the first one the
   billing-record verifier confirms, escalate if none (best-of-n with early stopping).
3. Map-reduce the week into report counts (map: one triage call per ticket; reduce: code counts).
4. Print calls, tokens, dollars, and the prompt fingerprints that produced the run.

Offline, every model call goes to a ScriptedLLM stand-in (plumbing only): the refund sampler follows
the simulated error model from m05_best_of_n.py, so its accuracy is the simulation's, not a model's.
With a key:  LLM_PROVIDER=groq GROQ_API_KEY=... PYTHONPATH=. python examples/m05_lab.py --live
"""
import hashlib
import json
import random
import sys
from collections import Counter

from examples.m05_helpers import (CASES, LABELS, SCAFFOLD, TRIAGE_KEYWORDS, case_from_messages, decide,
                                  keyword_responder, parse_category, parse_decision, prompt_scaffold,
                                  shortcut_baseline, triage_messages, triage_prompt, wilson)
from examples.m05_router import ROUTER, gold_route, route, rule_router
from supportdesk.data import load_tickets
from supportdesk.llm import Usage
from supportdesk.pricing import cost_usd
from supportdesk.stand_in import ScriptedLLM

MODEL = "openai/gpt-oss-120b"
MAX_SAMPLES = 5
rng = random.Random(7)


def simulated_sampler(messages, kwargs):
    """Stand-in: right with p=0.9 on easy cases, 0.45 on trap cases; traps usually give the shortcut answer."""
    case = case_from_messages(messages)
    trap = shortcut_baseline(case.text) != case.gold
    if rng.random() < (0.45 if trap else 0.9):
        label = case.gold
    elif trap and rng.random() < 0.8:
        label = shortcut_baseline(case.text)
    else:
        label = rng.choice([x for x in LABELS if x != case.gold])
    return f"(scripted stand-in reasoning)\nDECISION: {label}"


def verified_decision(chat_fn, case, usage: Usage) -> tuple[str, int]:
    """Sample until the billing-record check accepts an answer; escalate after MAX_SAMPLES."""
    for i in range(1, MAX_SAMPLES + 1):
        reply = chat_fn(prompt_scaffold(case), temperature=0.8, max_tokens=2000)
        usage.input_tokens += reply.usage.input_tokens
        usage.output_tokens += reply.usage.output_tokens
        label = parse_decision(reply.text)
        if label is not None and label == decide(case.facts):   # the external check
            return label, i
    return "escalate", MAX_SAMPLES


def fingerprint(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()[:12]


if __name__ == "__main__":
    live = "--live" in sys.argv
    if live:
        from supportdesk.llm import chat
        router_llm = sampler = chat
    else:
        router_llm, sampler = ScriptedLLM(responder=rule_router), ScriptedLLM(responder=simulated_sampler)
    print(f"prompts: router {fingerprint(ROUTER)}  refund scaffold {fingerprint(SCAFFOLD)}  "
          f"({'live model' if live else 'scripted stand-ins'})\n")

    # 1. Route the week.
    tickets = load_tickets()
    routes = [route(router_llm, t) for t in tickets]
    agree = sum(r == gold_route(t) for r, t in zip(routes, tickets))
    print(f"1. routing: {dict(Counter(routes))}; agrees with hand labels on {agree}/{len(tickets)}")

    # 2. Refund decisions with a verifier.
    usage = Usage()
    outcomes = [verified_decision(sampler, case, usage) for case in CASES]
    decided = [(case, label) for case, (label, _) in zip(CASES, outcomes) if label != "escalate"]
    calls = sum(n for _, n in outcomes)
    wrong = sum(label != case.gold for case, label in decided)
    lo, hi = wilson(len(decided), len(CASES))
    print(f"2. refunds: decided {len(decided)}/{len(CASES)} (95% CI {lo:.2f}-{hi:.2f}), wrong among decided {wrong}, "
          f"escalated {len(CASES) - len(decided)}; {calls} calls ({calls / len(CASES):.2f} per case)")
    print(f"   tokens in/out {usage.input_tokens}/{usage.output_tokens}; "
          f"USD {cost_usd(usage, MODEL):.4f} at {MODEL} prices")

    # 3. Map step: classify every ticket (keyword stand-in offline); reduce step: count in code.
    mapper = chat if live else ScriptedLLM(responder=keyword_responder)
    triage_text = triage_prompt(TRIAGE_KEYWORDS)
    predicted = [parse_category(mapper(triage_messages(triage_text, t.text), temperature=0.0).text) for t in tickets]
    counts = Counter(predicted)
    gold_counts = Counter(t.gold["category"] for t in tickets)
    right = sum(p == t.gold["category"] for p, t in zip(predicted, tickets))
    print(f"3. week: {len(tickets)} map calls; predicted {dict(counts.most_common())}")
    print(f"   gold      {dict(gold_counts.most_common())}; per-ticket agreement {right}/{len(tickets)}")

    # 4. Run record.
    record = {"router": fingerprint(ROUTER), "scaffold": fingerprint(SCAFFOLD), "routing_agree": agree,
              "refund_decided": len(decided), "refund_wrong": wrong, "refund_calls": calls, "map_agree": right,
              "usd": round(cost_usd(usage, MODEL), 6), "live": live}
    print("4. run record:", json.dumps(record))

"""Tests for Module 5: the refund rule, parsers, voting math, chains, routing, self-correction, prompt search."""
import random
from collections import Counter
from datetime import date

import pytest
from pydantic import ValidationError

from examples.m05_chain import RefundFacts, run_chain
from examples.m05_helpers import (CASES, Facts, PROMPTS, TODAY, case_from_messages, decide, evaluate,
                                  parse_decision, wilson)
from examples.m05_least_to_most import least_to_most, settled, stand_in
from examples.m05_prompt_search import PromptVersion, score, search, validate
from examples.m05_router import route
from examples.m05_self_correct import critique_and_revise, make_stand_in
from examples.m05_voting_math import majority_exact
from examples.m05_helpers import TRIAGE_KEYWORDS, keyword_responder
from supportdesk.data import load_tickets
from supportdesk.stand_in import ScriptedLLM


def test_rule_order_and_14_day_boundary():
    assert decide(Facts("member", "annual", TODAY, True)) == "needs_owner"       # role first
    assert decide(Facts("owner", "annual", date(2026, 1, 1), True)) == "refund_duplicate"
    assert decide(Facts("owner", "annual", date(2026, 9, 7), False)) == "refund_full"    # 14 days: inside
    assert decide(Facts("owner", "annual", date(2026, 9, 6), False)) == "no_refund"      # 15 days: outside
    assert decide(Facts("billing_admin", "monthly", TODAY, False)) == "no_refund"


def test_gold_distribution_is_stable():
    assert len(CASES) == 20
    assert Counter(c.gold for c in CASES) == {"no_refund": 8, "refund_full": 6, "refund_duplicate": 3, "needs_owner": 3}


def test_parse_decision_takes_last_known_label():
    assert parse_decision("DECISION: no_refund\n...\nDECISION: refund_full") == "refund_full"
    assert parse_decision("**DECISION:** needs_owner") == "needs_owner"
    assert parse_decision("Decision - Refund Full") is None
    assert parse_decision("DECISION: maybe") is None


def test_wilson_contains_point_estimate():
    lo, hi = wilson(14, 20)
    assert lo < 0.7 < hi and 0 <= lo and hi <= 1


def test_majority_math():
    assert majority_exact(0.7, 5) == pytest.approx(0.83692, abs=1e-5)
    assert all(majority_exact(0.5, n) == pytest.approx(0.5) for n in (1, 2, 3, 8, 15))
    assert majority_exact(0.3, 9) < 0.3  # voting amplifies a minority-correct model's errors


def test_harness_with_oracle_scores_all():
    def oracle(messages, kwargs):
        return f"DECISION: {decide(case_from_messages(messages).facts)}"
    result = evaluate(ScriptedLLM(responder=oracle), PROMPTS["direct"])
    assert (result.correct, result.unparsed) == (20, 0)


def test_chain_escalates_on_missing_facts():
    null = ScriptedLLM(replies=['{"role": null, "billing": null, "charged_on": null, "duplicate": null}'])
    trace = run_chain(CASES[0], null, ScriptedLLM(replies=[]))
    assert trace["outcome"] == "escalate" and trace["stage"] == "extract" and trace["calls"] == 1
    with pytest.raises(ValidationError):
        RefundFacts.model_validate_json('{"role": "owner", "billing": "annual", "charged_on": "3 September", "duplicate": false}')


def test_router_sends_unexpected_output_to_a_human():
    ticket = load_tickets()[0]
    assert route(ScriptedLLM(replies=['{"route": "refund_everything"}']), ticket) == "human"
    assert route(ScriptedLLM(replies=["not json"]), ticket) == "human"


def test_least_to_most_stops_once_settled():
    llm = ScriptedLLM(responder=stand_in)
    solved, calls = least_to_most(llm, CASES[3].text, settled)   # R04: a member, so question 1 settles it
    assert calls == 2 and solved[0][1] == "no"


def test_verifier_never_lets_a_correct_answer_become_wrong():
    llm = make_stand_in(random.Random(0))
    for _ in range(20):
        for case in CASES:
            history, _ = critique_and_revise(llm, case, 4, use_verifier=True)
            first_right = next((i for i, h in enumerate(history) if h == case.gold), None)
            if first_right is not None:
                assert all(h == case.gold for h in history[first_right:])


def test_without_verifier_correct_answers_can_flip():
    llm = make_stand_in(random.Random(0))
    flips = 0
    for case in CASES * 10:
        history, _ = critique_and_revise(llm, case, 1, use_verifier=False)
        flips += history[0] == case.gold and history[1] != case.gold
    assert flips > 0


def test_prompt_versions_and_search():
    v1 = PromptVersion.from_dict("v1", TRIAGE_KEYWORDS)
    assert v1.fingerprint == PromptVersion.from_dict("renamed", TRIAGE_KEYWORDS).fingerprint  # text, not name
    broken = PromptVersion.from_dict("bad", {k: v for k, v in TRIAGE_KEYWORDS.items() if k != "bug"})
    assert validate(broken) and not validate(v1)
    reader = ScriptedLLM(responder=keyword_responder)
    dev = load_tickets("dev")
    best, _, _ = search(reader, v1, dev, rounds=2)
    assert score(reader, best, dev).correct >= score(reader, v1, dev).correct + 2