Part F: Optimization

"""Module 7: context optimization, measured.

  1. cache-aware ordering: shared-prefix tokens across consecutive tickets
  2. just-in-time loading vs preloading: tokens and number of calls
  3. compaction triggers over a long conversation
  4. an A/B harness: did the added context improve the output?
Parts 2 and 4 drive a ScriptedLLM (rules, not a model) to exercise the plumbing;
pass --real to run part 4 with supportdesk.llm.chat and a real provider.
Run:  PYTHONPATH=. python examples/m07_optimize.py [--real]
"""
from __future__ import annotations

import json
import math
import random
import re
import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))  # this module's own example files

from m07_context import (FEW_SHOT, SYSTEM, TOOLS, ContextBuilder, extractive_summary,  # noqa: E402
                         tool_tokens)
from supportdesk.data import Ticket, get_article, load_tickets  # noqa: E402
from supportdesk.kb_search import KBSearch  # noqa: E402
from supportdesk.llm import ChatResult, ToolCall, Usage  # noqa: E402
from supportdesk.pricing import PRICES, cost_usd  # noqa: E402
from supportdesk.stand_in import ScriptedLLM  # noqa: E402
from supportdesk.tokens import count_messages, count_tokens, encode  # noqa: E402

KB = KBSearch()


# --- 1. cache-aware ordering ---------------------------------------------------------------------
def serialize(messages: list[dict[str, Any]], tools: list[dict[str, Any]], tools_first: bool) -> str:
    """Roughly how a provider flattens a request into one prompt. Tool placement differs by provider."""
    body = "".join(f"<|{m['role']}|>{m['content']}<|end|>" for m in messages)
    tool_text = f"<|tools|>{json.dumps(tools, separators=(',', ':'))}<|end|>" if tools else ""
    return tool_text + body if tools_first else body.replace("<|end|>", "<|end|>" + tool_text, 1)


def common_prefix(a: list[int], b: list[int]) -> int:
    n = 0
    for x, y in zip(a, b):
        if x != y:
            break
        n += 1
    return n


def layout(ticket: Ticket, order: str, per_category_tools: bool) -> tuple[list[dict], list[dict]]:
    """Build the same content in two orders: 'stable_first' (cache-aware) or 'volatile_first'."""
    category = ticket.gold["category"]
    b = (ContextBuilder().system().examples()
         .with_tools(category=category if per_category_tools else None)
         .profile({"customer_tier": ticket.customer_tier, "language": ticket.language})
         .knowledge(KB.search(ticket.text, k=3))
         .state({"ticket": ticket.id, "step": "draft_reply"}).ticket(ticket))
    ctx = b.build()
    if order == "stable_first":
        return ctx.messages, ctx.tools
    volatile = ctx.messages[-1]["content"]
    return [{"role": "system", "content": volatile}, {"role": "user", "content": ctx.stable_prefix}], ctx.tools


def cache_report(tickets: list[Ticket]) -> None:
    print("1. Shared prefix between each ticket and the one before it (dev split, arrival order)")
    print(f"   {'layout':52} {'prompt':>6} {'shared':>6} {'share':>6} {'billed':>6}")
    for order, per_cat, tools_first in [("volatile_first", False, False), ("stable_first", False, False),
                                        ("stable_first", True, False), ("stable_first", True, True),
                                        ("stable_first", False, True)]:
        ids = [encode(serialize(*layout(t, order, per_cat), tools_first)) for t in tickets]
        shared = [common_prefix(a, b) for a, b in zip(ids, ids[1:])]
        mean_len = sum(map(len, ids[1:])) / len(shared)
        label = f"{order}, {'category' if per_cat else 'all'} tools, {'tools first' if tools_first else 'tools after system'}"
        mean_shared = sum(shared) / len(shared)
        billed = mean_len - 0.5 * mean_shared  # input tokens billed at full price, 50% cache discount
        print(f"   {label:52} {mean_len:>6.0f} {mean_shared:>6.0f} {mean_shared / mean_len:>6.0%} {billed:>6.0f}")


# --- 2. just-in-time loading vs preloading -------------------------------------------------------
SEARCH_TOOL = [t for t in TOOLS if t["function"]["name"] == "search_kb"]


def jit_responder(messages: list[dict[str, Any]], kwargs: dict[str, Any]) -> ChatResult:
    """Scripted: first turn asks search_kb (unless the ticket needs no lookup), second turn answers."""
    if messages[-1]["role"] == "tool":
        return ChatResult(text='{"reply": "Draft based on the article.", "cited_articles": [], "confidence": "medium"}')
    ticket = messages[-1]["content"]
    if "feature=True" in ticket:  # stands in for "the model decided no lookup is needed"
        return ChatResult(text='{"reply": "Thanks for the idea.", "cited_articles": [], "confidence": "low"}')
    subject = re.search(r"Subject: (.*)", ticket).group(1)
    return ChatResult(text="", tool_calls=[ToolCall("call_1", "search_kb", {"query": subject, "k": 1})])


def run_jit(ticket: Ticket, llm: ScriptedLLM) -> tuple[int, int]:
    """Returns (calls, input tokens including tool definitions) for the just-in-time loop."""
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": f"<ticket>{ticket.text}</ticket>\nfeature={ticket.gold['category'] == 'feature_request'}"}]
    calls = tokens = 0
    for _ in range(3):
        result = llm(messages, tools=SEARCH_TOOL)
        calls += 1
        tokens += result.usage.input_tokens + tool_tokens(SEARCH_TOOL)
        if not result.tool_calls:
            break
        messages.append(result.as_message())
        for call in result.tool_calls:
            hits = KB.search(call.arguments["query"], k=call.arguments.get("k", 1))
            content = "\n".join(get_article(h.article_id).body for h in hits)
            messages.append({"role": "tool", "tool_call_id": call.id, "content": content})
    return calls, tokens


def jit_report(tickets: list[Ticket]) -> None:
    print("\n2. Preload top-3 articles vs just-in-time search_kb (dev split, input tokens per ticket)")
    llm = ScriptedLLM(responder=jit_responder)  # plumbing only: it decides by rule when to search
    rows = {"preload k=3": [0, 0], "preload k=1": [0, 0], "just-in-time": [0, 0]}
    for t in tickets:
        for k in (3, 1):
            knowledge = "\n".join(get_article(h.article_id).body for h in KB.search(t.text, k=k))
            msgs = [{"role": "system", "content": SYSTEM},
                    {"role": "user", "content": f"<knowledge>{knowledge}</knowledge>\n<ticket>{t.text}</ticket>"}]
            rows[f"preload k={k}"][0] += 1
            rows[f"preload k={k}"][1] += count_messages(msgs)
        calls, tokens = run_jit(t, llm)
        rows["just-in-time"][0] += calls
        rows["just-in-time"][1] += tokens
    print(f"   {'strategy':14} {'calls':>6} {'input tokens':>13} {'per ticket':>11}")
    for name, (calls, tokens) in rows.items():
        print(f"   {name:14} {calls:>6} {tokens:>13} {tokens / len(tickets):>11.0f}")


# --- 3. compaction triggers ----------------------------------------------------------------------
def long_conversation(turns: int, seed: int = 0) -> list[dict[str, str]]:
    """A synthetic multi-issue chat built from real tickets and article sentences."""
    rng = random.Random(seed)
    pool = load_tickets("dev")
    convo = []
    for i in range(turns):
        t = rng.choice(pool)
        convo.append({"role": "user", "content": t.body})
        art = t.gold["kb_article"] or "feature-requests"
        convo.append({"role": "assistant", "content": get_article(art).body.split("\n")[0]})
    return convo


def compaction_report() -> None:
    print("\n3. Compaction over a 30-exchange conversation (history tokens sent per turn)")
    convo = long_conversation(30)
    budget, keep = 1200, 6
    for name, trigger in [("never", lambda tok, n: False),
                          ("every 10 exchanges", lambda tok, n: n % 20 == 0),
                          ("at 70% of 1200", lambda tok, n: tok > 0.7 * budget),
                          ("at 95% of 1200", lambda tok, n: tok > 0.95 * budget)]:
        history: list[dict[str, str]] = []
        sent, peak, compactions, rewrites = 0, 0, 0, 0
        for n, message in enumerate(convo):
            history.append(message)
            tokens = sum(count_tokens(m["content"]) for m in history)
            if trigger(tokens, n + 1) and len(history) > keep:
                summary = extractive_summary(history[:-keep])
                history = [{"role": "user", "content": summary}] + history[-keep:]
                compactions += 1
                rewrites += 1  # the prefix changed, so the next call misses the prompt cache
                tokens = sum(count_tokens(m["content"]) for m in history)
            if message["role"] == "user":  # a model call happens after each customer message
                sent += tokens
                peak = max(peak, tokens)
        print(f"   {name:18} peak {peak:>5}  total sent {sent:>6}  compactions {compactions:>2}  cache resets {rewrites:>2}")


# --- 4. A/B harness --------------------------------------------------------------------------------
def extractive_answerer(messages: list[dict[str, Any]], kwargs: dict[str, Any]) -> str:
    """Scripted, not a model: cites the article whose text shares the most words with the ticket."""
    user = messages[-1]["content"]
    ticket = set(re.findall(r"[a-z]+", user.split("<ticket")[-1].lower()))
    best, best_overlap = None, 0
    for aid, body in re.findall(r'<article id="([^"]+)"[^>]*>\n(.*?)\n</article>', user, re.S):
        overlap = len(ticket & set(re.findall(r"[a-z]+", body.lower())))
        if overlap > best_overlap:
            best, best_overlap = aid, overlap
    return json.dumps({"reply": "...", "cited_articles": [best] if best else [], "confidence": "medium"})


def mcnemar_exact(b: int, c: int) -> float:
    """Two-sided exact p-value for b wins vs c losses among discordant pairs."""
    n = b + c
    if n == 0:
        return 1.0
    tail = sum(math.comb(n, i) for i in range(0, min(b, c) + 1)) / 2 ** n
    return min(1.0, 2 * tail)


def ab_test(variant_a: Callable[[Ticket], Any], variant_b: Callable[[Ticket], Any], tickets: list[Ticket],
            call: Callable[[list[dict]], str], judge: Callable[[Ticket, str], bool], model: str) -> None:
    """Run both context variants on every ticket; report paired wins, losses, p-value, tokens, cost."""
    passes = {"A": [], "B": []}
    tokens = {"A": 0, "B": 0}
    for t in tickets:
        for name, variant in (("A", variant_a), ("B", variant_b)):
            ctx = variant(t)
            passes[name].append(judge(t, call(ctx.messages)))
            tokens[name] += ctx.total_tokens
    wins = sum(b and not a for a, b in zip(passes["A"], passes["B"]))
    losses = sum(a and not b for a, b in zip(passes["A"], passes["B"]))
    n = len(tickets)
    for name in ("A", "B"):
        usage = Usage(input_tokens=tokens[name] // n)
        print(f"   variant {name}: pass {sum(passes[name])}/{n}  mean prompt tokens {tokens[name] / n:.0f}"
              f"  input cost per 1,000 tickets on {model}: {cost_usd(usage, model) * 1000:.3f} USD")
    print(f"   B vs A: {wins} wins, {losses} losses, {n - wins - losses} ties;"
          f" exact McNemar p = {mcnemar_exact(wins, losses):.3f}")


def draft_judge(ticket: Ticket, text: str) -> bool:
    """Pass if the draft parses and cites the gold article (unanswerable: cites nothing or says low)."""
    try:
        draft = json.loads(re.search(r"\{.*\}", text, re.S).group(0))
    except (AttributeError, json.JSONDecodeError):
        return False
    gold = ticket.gold["kb_article"]
    return gold in draft.get("cited_articles", []) if gold else draft.get("confidence") == "low"


def variant(k: int) -> Callable[[Ticket], Any]:
    def build(t: Ticket):
        return (ContextBuilder().system().examples().with_tools(category=t.gold["category"])
                .profile({"customer_tier": t.customer_tier, "language": t.language})
                .knowledge(KB.search(t.text, k=k)).ticket(t).build())
    return build


if __name__ == "__main__":
    dev = load_tickets("dev")
    cache_report(dev)
    jit_report(dev)
    compaction_report()
    answerable = [t for t in dev if t.gold["answerable"]]
    model = "gemini-3.5-flash"
    if "--real" in sys.argv:
        from supportdesk.llm import chat, resolve
        model = resolve()[1]
        print(f"\n4. A/B on {len(answerable)} answerable dev tickets with a real model ({model})")
        ab_test(variant(1), variant(3), answerable, lambda m: chat(m, temperature=0.0).text, draft_judge,
                model if model in PRICES else "openai/gpt-oss-120b")
    else:
        print(f"\n4. A/B harness, knowledge k=1 (A) vs k=3 (B), {len(answerable)} answerable dev tickets,"
              " ScriptedLLM extractive answerer (plumbing, not model quality)")
        llm = ScriptedLLM(responder=extractive_answerer)
        ab_test(variant(1), variant(3), answerable, lambda m: llm(m).text, draft_judge, model)

1. Shared prefix between each ticket and the one before it (dev split, arrival order)
   layout                                               prompt shared  share billed
   volatile_first, all tools, tools after system          1379     18     1%   1370
   stable_first, all tools, tools after system            1379    966    70%    896
   stable_first, category tools, tools after system       1165    664    57%    833
   stable_first, category tools, tools first              1165    291    25%   1020
   stable_first, all tools, tools first                   1379    966    70%    896

2. Preload top-3 articles vs just-in-time search_kb (dev split, input tokens per ticket)
   strategy        calls  input tokens  per ticket
   preload k=3        48         24032         501
   preload k=1        48         15066         314
   just-in-time       92         30587         637

3. Compaction over a 30-exchange conversation (history tokens sent per turn)
   never              peak  1480  total sent  22146  compactions  0  cache resets  0
   every 10 exchanges peak   994  total sent  16346  compactions  3  cache resets  3
   at 70% of 1200     peak   829  total sent  14265  compactions  2  cache resets  2
   at 95% of 1200     peak  1108  total sent  19170  compactions  2  cache resets  2

4. A/B harness, knowledge k=1 (A) vs k=3 (B), 42 answerable dev tickets, ScriptedLLM extractive answerer (plumbing, not model quality)
   variant A: pass 35/42  mean prompt tokens 916  input cost per 1,000 tickets on gemini-3.5-flash: 1.373 USD
   variant B: pass 30/42  mean prompt tokens 1137  input cost per 1,000 tickets on gemini-3.5-flash: 1.706 USD
   B vs A: 0 wins, 5 losses, 37 ties; exact McNemar p = 0.062

export LLM_PROVIDER=groq GROQ_API_KEY=...   # or gemini / ollama, as in Module 1
python examples/m07_optimize.py --real

"""Module 7 Lab: the full context pipeline over the dev split.

For each ticket: expand the query and retrieve articles, load remembered facts,
assemble every section with the ContextBuilder, call the model, validate the
draft, and record tokens per section. Then report context recall, pass rate,
tokens, shared prefix, estimated cost, and diagnose the first failure by ablation.

Default model: a ScriptedLLM extractive answerer (plumbing only, not a model).
With --real it calls supportdesk.llm.chat on your configured provider.
Run:  PYTHONPATH=. python examples/m07_lab.py [--real]
"""
from __future__ import annotations

import json
import re
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from m07_context import SYSTEM, ContextBuilder  # noqa: E402
from m07_diagnose import article, diagnose  # noqa: E402
from m07_memory import CONVERSATIONS, SCRIPTED_EXTRACTIONS, MemoryStore  # noqa: E402
from m07_optimize import common_prefix, extractive_answerer, serialize  # noqa: E402
from m07_retrieval import expand_query  # noqa: E402
from pydantic import ValidationError  # noqa: E402
from supportdesk.data import Ticket, load_tickets  # noqa: E402
from supportdesk.kb_search import KBSearch  # noqa: E402
from supportdesk.llm import Usage  # noqa: E402
from supportdesk.pricing import PRICES, cost_usd  # noqa: E402
from supportdesk.schemas import DraftReply  # noqa: E402
from supportdesk.stand_in import ScriptedLLM  # noqa: E402
from supportdesk.tokens import encode  # noqa: E402

NOW = datetime(2026, 9, 21, 9, 0)
CUSTOMER_OF = {"T-1004": "cust-acme", "T-1041": "cust-acme", "T-1067": "cust-acme"}  # demo mapping


def parse_draft(text: str) -> DraftReply | None:
    match = re.search(r"\{.*\}", text, re.S)
    if not match:
        return None
    try:
        return DraftReply.model_validate(json.loads(match.group(0)))
    except (json.JSONDecodeError, ValidationError):
        return None


def main(real: bool) -> None:
    kb = KBSearch()
    store = MemoryStore()
    seed = ScriptedLLM(replies=list(SCRIPTED_EXTRACTIONS[:3]))  # scripted extractor, see m07_memory.py
    for ticket_id, when, convo in CONVERSATIONS[:3]:
        store.extract(convo, "cust-acme", ticket_id, when, seed)

    if real:
        from supportdesk.llm import chat, resolve
        model = resolve()[1]
        call = lambda msgs, tools: chat(msgs, tools=tools or None, temperature=0.0, max_tokens=600)  # noqa: E731
    else:
        model = "gemini-3.5-flash"
        stand_in = ScriptedLLM(responder=extractive_answerer)
        call = lambda msgs, tools: stand_in(msgs, tools=tools)  # noqa: E731

    tickets = load_tickets("dev")
    section_tokens: dict[str, list[int]] = defaultdict(list)
    recall = passed = valid = 0
    answerable_ok = {True: 0, False: 0}
    prompts, failures, total = [], [], Usage()
    for t in tickets:
        hits = kb.search(expand_query(t), k=3)
        customer = CUSTOMER_OF.get(t.id, f"cust-{t.id}")
        memories = [m.line() for m in store.retrieve(customer, t.text, NOW)]
        ctx = (ContextBuilder().system().examples().with_tools(category=t.gold["category"])
               .profile({"customer_tier": t.customer_tier, "language": t.language}, memories)
               .knowledge(hits)
               .state({"ticket": t.id, "step": "draft_reply", "plan": ["read", "check policy", "draft", "cite"]})
               .ticket(t).build())
        for r in ctx.report:
            section_tokens[r.name].append(r.tokens)
        prompts.append(encode(serialize(ctx.messages, ctx.tools, tools_first=False)))
        result = call(ctx.messages, ctx.tools)
        total.input_tokens += result.usage.input_tokens if real else ctx.total_tokens  # stand-in: our estimate
        total.output_tokens += result.usage.output_tokens
        total.cached_tokens += result.usage.cached_tokens
        draft = parse_draft(result.text)
        gold = t.gold["kb_article"]
        recall += gold is None or gold in [h.article_id for h in hits]
        valid += draft is not None
        ok = draft is not None and (gold in draft.cited_articles if gold else draft.confidence == "low")
        passed += ok
        answerable_ok[gold is not None] += ok
        if not ok:
            failures.append((t, hits))

    n = len(tickets)
    print(f"Lab over {n} dev tickets ({'real model ' + model if real else 'ScriptedLLM extractive answerer: plumbing only'})")
    print(f"\n{'section':10} {'mean':>6} {'max':>5}")
    for name, values in section_tokens.items():
        print(f"{name:10} {sum(values) / n:>6.0f} {max(values):>5}")
    shared = [common_prefix(a, b) for a, b in zip(prompts, prompts[1:])]
    print(f"\ncontext recall (gold article in context, unanswerable counted as ok): {recall}/{n}")
    n_answerable = sum(t.gold["answerable"] for t in tickets)
    print(f"valid DraftReply JSON: {valid}/{n}; answerable drafts citing the gold article: "
          f"{answerable_ok[True]}/{n_answerable}; unanswerable drafts marked low confidence: "
          f"{answerable_ok[False]}/{n - n_answerable}")
    print(f"mean prompt {sum(map(len, prompts)) / n:.0f} tokens, mean shared prefix with previous ticket "
          f"{sum(shared) / len(shared):.0f} tokens")
    price_model = model if model in PRICES else "openai/gpt-oss-120b"
    print(f"input+output tokens {total.input_tokens}+{total.output_tokens}; estimated cost on {price_model}: "
          f"{cost_usd(total, price_model):.4f} USD for {n} tickets")

    if failures:
        t, hits = failures[0]
        sections = {"system": [SYSTEM], "profile": [f"customer_tier: {t.customer_tier}"],
                    "knowledge": [article(h.article_id) for h in hits], "ticket": [t.text]}

        def check(answer: str) -> bool:
            draft = parse_draft(answer)
            return draft is not None and t.gold["kb_article"] in draft.cited_articles

        def rerun(messages: list[dict]) -> str:
            return call(messages, None).text

        d = diagnose(sections, rerun, check)
        print(f"\nfirst failure {t.id} ({t.language}) gold={t.gold['kb_article']} retrieved={[h.article_id for h in hits]}")
        print(f"  verdict: {d.verdict}  ({d.calls} calls)")


if __name__ == "__main__":
    main(real="--real" in sys.argv)

Lab over 48 dev tickets (ScriptedLLM extractive answerer: plumbing only)

section      mean   max
system        171   171
examples      369   369
tools         180   292
profile        22   118
knowledge     349   396
state          30    30
ticket         41    57

context recall (gold article in context, unanswerable counted as ok): 46/48
valid DraftReply JSON: 48/48; answerable drafts citing the gold article: 30/42; unanswerable drafts marked low confidence: 0/6
mean prompt 1191 tokens, mean shared prefix with previous ticket 664 tokens
input+output tokens 56546+1006; estimated cost on gemini-3.5-flash: 0.0939 USD for 48 tickets

first failure T-1001 (en) gold=billing-refunds retrieved=['billing-invoices', 'billing-plans', 'boards-automations']
  verdict: no single removal fixes it: check retrieval (missing knowledge) or the model itself  (7 calls)

python -m pytest -q tests/test_m07_context.py

..............                                                           [100%]
14 passed in 4.40s