Part B: Retrieved knowledge

"""A small BM25 keyword search over the help-center articles (no external services)."""
from __future__ import annotations

import math
import re
from collections import Counter
from dataclasses import dataclass

from supportdesk.data import Article, load_articles

WORD = re.compile(r"[a-z0-9]+")
STOPWORDS = frozenset("a an and are as at be but by can do does for from has have how i if in is it its me my no not of on or our so that the their this to was we what when where which who why will with you your".split())


def tokenize(text: str) -> list[str]:
    return [w for w in WORD.findall(text.lower()) if w not in STOPWORDS]


@dataclass(frozen=True)
class Hit:
    article_id: str
    title: str
    score: float


class KBSearch:
    """BM25 ranking: rewards rare query words that appear often in a short article."""

    def __init__(self, articles: list[Article] | None = None, k1: float = 1.2, b: float = 0.75) -> None:
        self.articles = articles if articles is not None else load_articles()
        self.k1, self.b = k1, b
        self.docs = [tokenize(a.title + " " + a.title + " " + a.body) for a in self.articles]
        self.avg_len = sum(len(d) for d in self.docs) / len(self.docs)
        df = Counter(w for d in self.docs for w in set(d))
        n = len(self.docs)
        self.idf = {w: math.log(1 + (n - c + 0.5) / (c + 0.5)) for w, c in df.items()}
        self.tf = [Counter(d) for d in self.docs]

    def search(self, query: str, k: int = 3) -> list[Hit]:
        terms = tokenize(query)
        scored = []
        for article, tf, doc in zip(self.articles, self.tf, self.docs):
            score = 0.0
            for t in terms:
                if t in tf:
                    f = tf[t]
                    score += self.idf[t] * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * len(doc) / self.avg_len))
            if score > 0:
                scored.append(Hit(article.id, article.title, round(score, 3)))
        scored.sort(key=lambda h: (-h.score, h.article_id))
        return scored[:k]

"""Module 7: measure KB retrieval by language, then try one improvement.

Run from the repo root:  PYTHONPATH=. python examples/m07_retrieval.py
"""
from __future__ import annotations

import math
from collections import defaultdict

from supportdesk.data import Ticket, get_article, load_tickets
from supportdesk.kb_search import Hit, KBSearch
from supportdesk.tokens import count_tokens

# A small glossary built from the help-center vocabulary (not from the tickets):
# each English KB term with the words a Spanish, German, Japanese, or Hindi
# customer would use for it. Keys are matched as lowercase substrings, because
# the canonical tokenizer keeps only [a-z0-9] and silently drops Japanese and
# Hindi characters (and splits words at accented letters).
GLOSSARY: dict[str, tuple[str, ...]] = {
    "refund refunded": ("reembolso", "rueckerstattung", "rückerstattung", "erstatt", "返金", "रिफंड", "धनवापसी"),
    "duplicate charge charged twice": ("duplicado", "dos veces", "doppelt", "zweimal", "二重", "दो बार"),
    "cancel cancellation": ("cancelar", "kündig", "kuendig", "解約", "キャンセル", "रद्द"),
    "password reset link": ("contraseña", "contrasena", "passwort", "パスワード", "पासवर्ड"),
    "locked attempts": ("bloquead", "gesperrt", "ロック", "लॉक"),
    "reset": ("restablecer", "zuruecksetzen", "zurücksetzen", "リセット", "रीसेट"),
    "price cost plans pricing": ("precio", "cuesta", "preis", "kosten", "料金", "価格", "कीमत", "मूल्य"),
    "per user": ("por usuario", "pro benutzer", "ユーザーあたり", "प्रति उपयोगकर्ता"),
    "invoice invoices": ("factura", "rechnung", "請求書", "चालान", "इनवॉइस"),
    "annual renewal": ("anual", "jährlich", "jaehrlich", "年間", "वार्षिक"),
    "export csv": ("exportar", "exportier", "エクスポート", "निर्यात"),
    "notifications slack channel": ("notificaciones", "benachrichtigung", "通知", "सूचना"),
    "sign in login": ("iniciar sesión", "anmelden", "ログイン", "लॉगिन"),
    "subscription plan": ("suscripción", "abonnement", "サブスクリプション", "プラン", "सदस्यता", "प्लान"),
    "delete deletion data": ("eliminar", "borrar", "löschen", "loeschen", "削除", "हटा"),
}


def expand_query(ticket: Ticket, subject_weight: int = 1, glossary: bool = True) -> str:
    """Build the search query: the subject repeated, the body, and English glossary terms."""
    text = f"{ticket.subject} " * subject_weight + ticket.body
    extra = []
    if glossary:
        lowered = text.lower()
        extra = [english for english, foreign in GLOSSARY.items() if any(f in lowered for f in foreign)]
    return " ".join([text, *extra])


def answerable(split: str | None = None) -> list[Ticket]:
    return [t for t in load_tickets(split) if t.gold["answerable"]]


def evaluate(search, tickets: list[Ticket]) -> dict[str, dict[str, int]]:
    """Top-1 and top-3 hits per language. `search(ticket)` returns a ranked list of Hits."""
    table: dict[str, dict[str, int]] = defaultdict(lambda: {"n": 0, "top1": 0, "top3": 0})
    for t in tickets:
        ids = [h.article_id for h in search(t)]
        for key in (t.language, "all"):
            row = table[key]
            row["n"] += 1
            row["top1"] += ids[:1] == [t.gold["kb_article"]]
            row["top3"] += t.gold["kb_article"] in ids[:3]
    return dict(table)


def wilson(hits: int, n: int, z: float = 1.96) -> tuple[float, float]:
    """95 percent Wilson interval for a proportion: honest error bars for small n."""
    if n == 0:
        return (0.0, 0.0)
    p = hits / n
    centre = (p + z * z / (2 * n)) / (1 + z * z / n)
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return (centre - half, centre + half)


def print_table(title: str, table: dict[str, dict[str, int]]) -> None:
    print(title)
    print(f"  {'lang':5} {'n':>3} {'top1':>6} {'top3':>6}   top1 95% CI")
    for lang in ("en", "es", "de", "ja", "hi", "all"):
        if lang in table:
            r = table[lang]
            lo, hi = wilson(r["top1"], r["n"])
            print(f"  {lang:5} {r['n']:>3} {r['top1']:>3}/{r['n']:<2} {r['top3']:>3}/{r['n']:<2}   {lo:.2f} to {hi:.2f}")


def format_articles(hits: list[Hit]) -> str:
    """How retrieved articles are pasted into a prompt (same format the ContextBuilder uses)."""
    return "\n\n".join(f'<article id="{h.article_id}">\n{get_article(h.article_id).body}\n</article>' for h in hits)


if __name__ == "__main__":
    kb = KBSearch()
    tickets = answerable()
    base = evaluate(lambda t: kb.search(t.text, k=3), tickets)
    print_table("Baseline KBSearch on ticket.text (all 62 answerable tickets)", base)

    print("\nMisses at top-1 (baseline):")
    for t in tickets:
        ids = [h.article_id for h in kb.search(t.text, k=3)]
        if ids[:1] != [t.gold["kb_article"]]:
            print(f"  {t.id} {t.language} {t.split:4} gold={t.gold['kb_article']:18} got={ids}")

    print("\nWhat the tokenizer keeps from a Japanese ticket:")
    ja = next(t for t in tickets if t.language == "ja")
    from supportdesk.kb_search import tokenize
    print(f"  {ja.text!r}\n  -> {tokenize(ja.text)}")

    variants = {
        "baseline": lambda t: kb.search(t.text, k=3),
        "subject x2": lambda t: kb.search(expand_query(t, 2, glossary=False), k=3),
        "glossary": lambda t: kb.search(expand_query(t, 1, glossary=True), k=3),
        "subject x2 + glossary": lambda t: kb.search(expand_query(t, 2, glossary=True), k=3),
    }
    print("\nAblation, top-1 hits (dev n=42, test n=20, non-English n=8):")
    for name, fn in variants.items():
        dev = evaluate(fn, answerable("dev"))["all"]
        test = evaluate(fn, answerable("test"))["all"]
        non_en = evaluate(fn, [t for t in tickets if t.language != "en"])["all"]
        print(f"  {name:22} dev {dev['top1']:>2}/{dev['n']}  test {test['top1']:>2}/{test['n']}"
              f"  non-en {non_en['top1']}/{non_en['n']}  (top-3 all: "
              f"{evaluate(fn, tickets)['all']['top3']}/{len(tickets)})")

    best = variants["glossary"]
    print("\nPer-ticket flips at top-1, baseline vs glossary:")
    for t in tickets:
        before = [h.article_id for h in kb.search(t.text, k=1)] == [t.gold["kb_article"]]
        after = [h.article_id for h in best(t)][:1] == [t.gold["kb_article"]]
        if before != after:
            print(f"  {t.id} {t.language} {t.split:4} {'gained' if after else 'LOST'}")
    print_table("\nGlossary search (all 62 answerable tickets)", evaluate(best, tickets))

    print("\nRecall@k vs tokens added (glossary search, 62 tickets, o200k_base):")
    print(f"  {'k':>2} {'recall@k':>9} {'mean tokens':>12} {'extra tokens per extra hit':>27}")
    prev_hits, prev_tokens = 0, 0.0
    for k in range(1, 6):
        hits_found, tokens = 0, 0
        for t in tickets:
            hits = kb.search(expand_query(t), k=k)
            hits_found += t.gold["kb_article"] in [h.article_id for h in hits]
            tokens += count_tokens(format_articles(hits))
        mean_tokens = tokens / len(tickets)
        gained = hits_found - prev_hits
        per_hit = f"{(mean_tokens - prev_tokens) * len(tickets) / gained:.0f}" if gained and k > 1 else "-"
        print(f"  {k:>2} {hits_found:>3}/{len(tickets)} ({hits_found / len(tickets):.2f}) {mean_tokens:>9.0f}"
              f" {per_hit:>27}")
        prev_hits, prev_tokens = hits_found, mean_tokens

# Where does the best article go? PYTHONPATH=.:examples python examples/m07_order.py
from m07_context import ContextBuilder
from m07_retrieval import expand_query
from supportdesk.data import load_tickets
from supportdesk.kb_search import KBSearch

kb = KBSearch()
t = next(x for x in load_tickets() if x.id == "T-1062")
hits = kb.search(expand_query(t), k=3)
for best_last in (False, True):
    b = ContextBuilder().knowledge(hits, best_last=best_last)
    print(f"best_last={best_last}: order sent = {b.reports['knowledge'].note}")