Part E: Failure modes

python examples/m07_rot.py

question                                     wrong turns  P/tok right  P/tok wrong
my account is locked after too many attempts           0        0.999        0.000
my account is locked after too many attempts           1        0.998        0.000
can I get a refund on my annual Team plan?             0        0.999        0.001
can I get a refund on my annual Team plan?             1        0.999        0.001
how much does the Team plan cost?                      0        0.998        0.062
how much does the Team plan cost?                      1        0.997        0.049
does the Team plan include SSO?                        0        0.998        0.001
does the Team plan include SSO?                        1        0.994        0.001

"""Module 7: find which part of the context caused a bad answer, by ablation.

Given a context split into sections (each a list of items), a way to call the
model, and a check that says whether an answer is acceptable, the diagnostic
removes one section at a time, then one item at a time, reruns, and reports
which removals turn the bad answer into a good one. A small numeric-claim
check flags clashes between sections.

The responder used here is a ScriptedLLM with rules that react to specific
context contents (a planted wrong turn, a stale price, too many articles).
It is NOT a model: it exists so the diagnostic's plumbing can be run and
tested offline. Run the same harness with supportdesk.llm.chat for real results.
Run:  PYTHONPATH=. python examples/m07_diagnose.py
"""
from __future__ import annotations

import re
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from supportdesk.data import get_article
from supportdesk.kb_search import tokenize
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_tokens

ORDER = ("system", "profile", "history", "knowledge", "ticket")
SYSTEM = "You are the Brightlane support assistant. Answer only from the knowledge and profile. Keep it short."


def render(sections: dict[str, list[str]]) -> list[dict[str, str]]:
    """Sections to chat messages. Each section becomes a tagged block; empty sections vanish."""
    blocks = [f"<{name}>\n" + "\n".join(sections[name]) + f"\n</{name}>"
              for name in ORDER[1:] if sections.get(name)]
    return [{"role": "system", "content": "\n".join(sections.get("system", [SYSTEM]))},
            {"role": "user", "content": "\n\n".join(blocks)}]


# --- the scripted responder (rules, not intelligence) ---------------------------------------------
def scripted_responder(messages: list[dict[str, Any]], kwargs: dict[str, Any]) -> str:
    user = messages[-1]["content"]
    ticket = re.search(r"<ticket>\n(.*?)\n</ticket>", user, re.S)
    question = ticket.group(1) if ticket else ""
    articles = re.findall(r"<article[^>]*>\n(.*?)\n</article>", user, re.S)
    # Rule 1 (poisoning): an earlier claim of pro rata refunds is repeated as fact.
    if re.search(r"refunded pro rata for the remaining", user):
        return "Good news: as we said earlier, annual plans are refunded pro rata for the remaining months."
    # Rule 2 (clash): the first Team price that appears anywhere wins.
    if "Team" in question and ("save" in question or "cost" in question):
        price = re.search(r"Team (?:costs |at )?(\d+) USD per user per month", user)
        if price:
            return f"Team is {price.group(1)} USD per user per month."
    # Rule 3 (distraction): with more than 4 articles, only the first article is read.
    if len(articles) > 4:
        articles = articles[:1]
    # Otherwise: return the knowledge sentence that shares the most words with the question.
    q = set(tokenize(question))
    sentences = [s for a in articles for s in re.split(r"(?<=[.])\s", a)]
    best = max(sentences, key=lambda s: len(q & set(tokenize(s))), default="")
    return best if best and q & set(tokenize(best)) else "I could not find this in the help center."


# --- the diagnostic ---------------------------------------------------------------------------------
@dataclass
class Diagnosis:
    baseline: str
    baseline_ok: bool
    section_fixes: list[str] = field(default_factory=list)
    item_fixes: list[tuple[str, int, str]] = field(default_factory=list)
    clashes: list[str] = field(default_factory=list)
    calls: int = 0
    verdict: str = ""


PRICE = re.compile(r"(Team|Business)\D{0,20}?(\d+) USD per user per month")


def numeric_clashes(sections: dict[str, list[str]]) -> list[str]:
    """Flag the same plan quoted at different monthly prices in different sections."""
    seen: dict[str, set[tuple[str, str]]] = {}
    for name, items in sections.items():
        for item in items:
            for plan, amount in PRICE.findall(item):
                seen.setdefault(plan, set()).add((amount, name))
    return [f"{plan}: " + ", ".join(f"{a} USD in {s}" for a, s in sorted(v))
            for plan, v in seen.items() if len({a for a, _ in v}) > 1]


def diagnose(sections: dict[str, list[str]], call: Callable[[list[dict]], str],
             ok: Callable[[str], bool], protected: tuple[str, ...] = ("system", "ticket")) -> Diagnosis:
    """Ablate each section, then each item of every section, and record which removals fix the answer."""
    d = Diagnosis(baseline=call(render(sections)), baseline_ok=False)
    d.baseline_ok, d.calls = ok(d.baseline), 1
    if d.baseline_ok:
        d.verdict = "answer passes; nothing to diagnose"
        return d
    for name in sections:
        if name in protected:
            continue
        d.calls += 1
        if ok(call(render({k: v for k, v in sections.items() if k != name}))):
            d.section_fixes.append(name)
        for i, item in enumerate(sections[name]):
            trimmed = {**sections, name: sections[name][:i] + sections[name][i + 1:]}
            d.calls += 1
            if ok(call(render(trimmed))):
                d.item_fixes.append((name, i, item[:70]))
    d.clashes = numeric_clashes(sections)
    by_section: dict[str, int] = {}
    for name, _, _ in d.item_fixes:
        by_section[name] = by_section.get(name, 0) + 1
    if any(n >= 2 for n in by_section.values()):
        name = max(by_section, key=by_section.get)
        d.verdict = (f"distraction: removing any of {by_section[name]} items in '{name}' fixes it, "
                     "so volume, not one bad item, is the problem")
    elif len(d.item_fixes) == 1 and d.clashes:
        d.verdict = f"clash: one item in '{d.item_fixes[0][0]}' contradicts another section ({d.clashes[0]})"
    elif len(d.item_fixes) == 1:
        d.verdict = f"poisoning: one item in '{d.item_fixes[0][0]}' carries a wrong claim the answer repeats"
    elif d.section_fixes:
        d.verdict = f"section-level: removing {d.section_fixes} fixes it but no single item does"
    else:
        d.verdict = "no single removal fixes it: check retrieval (missing knowledge) or the model itself"
    return d


def article(article_id: str) -> str:
    return f'<article id="{article_id}">\n{get_article(article_id).body}\n</article>'


CASES = {
    "T-1005 refund after 2 months": (
        {"system": [SYSTEM],
         "profile": ["customer_tier: team", "remembered: pays annually (customer_said, 2026-07-02)"],
         "history": ["user: We paid annually in July. Can we stop now?",
                     "assistant: Annual plans can be refunded pro rata for the remaining months, I will check.",
                     "user: OK, how does that work?"],
         "knowledge": [article("billing-refunds"), article("billing-invoices")],
         "ticket": ["We paid annually in July and want to stop now. Can I get a refund for the remaining months?"]},
        lambda a: "pro rata" not in a and ("14 days" in a or "not refunded" in a)),
    "T-1067 annual discount on Team": (
        {"system": [SYSTEM],
         "profile": ["customer_tier: team", "remembered: was quoted Team at 15 USD per user per month (agent_note, 2025-02-01)"],
         "history": [],
         "knowledge": [article("billing-plans")],
         "ticket": ["How much do we save per user per month going annual on Team? Team cost today?"]},
        lambda a: "12" in a),
    "T-1016 webhook retries": (
        {"system": [SYSTEM],
         "profile": ["customer_tier: business"],
         "history": [],
         "knowledge": [article(a) for a in ("integrations-slack", "billing-plans", "boards-automations",
                                              "status-incidents", "mobile-app")],
         "ticket": ["Our webhook endpoint was down for 2 minutes. Will the automation retry the calls?"]},
        lambda a: "retr" in a.lower()),
}


if __name__ == "__main__":
    for name, (sections, ok) in CASES.items():
        llm = ScriptedLLM(responder=scripted_responder)  # rules that react to planted content, not a model
        d = diagnose(sections, lambda msgs: llm(msgs).text, ok)
        tokens = count_tokens(render(sections)[1]["content"])
        print(f"== {name} (user message {tokens} tokens)")
        print(f"  bad answer: {d.baseline}")
        print(f"  sections whose removal fixes it: {d.section_fixes or 'none'}")
        for section, i, text in d.item_fixes:
            print(f"  item fix: {section}[{i}] {text!r}")
        if d.clashes:
            print(f"  numeric clash: {d.clashes}")
        print(f"  verdict: {d.verdict}")
        print(f"  model calls spent on the diagnosis: {d.calls}\n")