Part B: The context window

# examples/m02_budget.py
"""Module 2: a context budget for one Brightlane assistant request, checked against several windows."""
import json
from dataclasses import dataclass, field

from m02_conversation import SYSTEM, conversation

from supportdesk.data import get_article
from supportdesk.kb_search import KBSearch
from supportdesk.tokens import count_messages, count_tokens

TOOLS = [
    {"type": "function", "function": {
        "name": "lookup_invoice", "description": "Fetch one invoice by id, with amount, date, and payment status.",
        "parameters": {"type": "object", "properties": {"invoice_id": {"type": "string", "pattern": "^INV-\\d{4}-\\d{6}$"}},
                       "required": ["invoice_id"]}}},
    {"type": "function", "function": {
        "name": "create_refund_request", "description": "Open a refund request for a human billing agent to approve.",
        "parameters": {"type": "object", "properties": {"invoice_id": {"type": "string"}, "reason": {"type": "string"},
                                                        "amount_usd": {"type": "number"}},
                       "required": ["invoice_id", "reason"]}}},
]


@dataclass
class ContextBudget:
    """window >= system + tools + history + retrieved + user + output_reserve + safety margin."""
    window: int
    output_reserve: int
    safety: float = 0.05  # our tokenizer is not the model's tokenizer: keep 5% slack
    parts: dict[str, int] = field(default_factory=dict)

    def add(self, name: str, tokens: int) -> None:
        self.parts[name] = self.parts.get(name, 0) + tokens

    @property
    def usable(self) -> int:
        return int(self.window * (1 - self.safety)) - self.output_reserve

    @property
    def used(self) -> int:
        return sum(self.parts.values())

    def report(self, label: str) -> str:
        status = "fits" if self.used <= self.usable else f"OVER by {self.used - self.usable}"
        return f"{label:28} window {self.window:>9,}  usable input {self.usable:>9,}  used {self.used:>6,}  {status}"


messages = conversation()
history, question = messages[1:-1], messages[-1]["content"]
hits = KBSearch().search("duplicate charge refund invoice", k=3)
retrieved = "\n\n".join(f"[{h.article_id}]\n{get_article(h.article_id).body}" for h in hits)

parts = {
    "system": count_tokens(SYSTEM) + 4,
    "tools": count_tokens(json.dumps(TOOLS)),
    "history": count_messages(history) - 3,
    "retrieved": count_tokens(retrieved) + 4,
    "user": count_tokens(question) + 4 + 3,
}
print("Request parts (o200k_base estimates):")
for name, n in parts.items():
    print(f"  {name:10} {n:5}")
print(f"  {'total':10} {sum(parts.values()):5}   retrieved articles: {[h.article_id for h in hits]}\n")

for label, window, reserve in [
    ("gpt-oss-120b on Groq", 131_072, 4_096),       # reasoning model: reserve room to think
    ("gemini-3.5-flash", 1_048_576, 4_096),
    ("qwen3:8b on Ollama, 4k", 4_096, 1_024),         # Ollama default below 24 GiB VRAM
    ("a 2k-token small model", 2_048, 1_024),
    ("TinyLM", 128, 40),
]:
    budget = ContextBudget(window=window, output_reserve=reserve)
    for name, n in parts.items():
        budget.add(name, n)
    print(budget.report(label))

per_message = parts["history"] / len(history)
fixed = sum(parts.values()) - parts["history"]
for label, usable in [("qwen3:8b on Ollama, 4k", ContextBudget(4_096, 1_024).usable),
                      ("gpt-oss-120b on Groq", ContextBudget(131_072, 4_096).usable)]:
    print(f"{label}: room for about {int((usable - fixed) / per_message):,} history messages "
          f"at {per_message:.0f} tokens each")

  Request parts (o200k_base estimates):
    system        32
    tools        174
    history      505
    retrieved    218
    user          21
    total        950   retrieved articles: ['billing-refunds', 'billing-invoices']

  gpt-oss-120b on Groq         window   131,072  usable input   120,422  used    950  fits
  gemini-3.5-flash             window 1,048,576  usable input   992,051  used    950  fits
  qwen3:8b on Ollama, 4k       window     4,096  usable input     2,867  used    950  fits
  a 2k-token small model       window     2,048  usable input       921  used    950  OVER by 29
  TinyLM                       window       128  usable input        81  used    950  OVER by 869
  qwen3:8b on Ollama, 4k: room for about 105 history messages at 23 tokens each
  gpt-oss-120b on Groq: room for about 5,226 history messages at 23 tokens each

# examples/m02_tinylm_limit.py
"""Module 2: what TinyLM does with a prompt longer than its 128-token context window."""
import torch

from supportdesk.tinylm import SamplingParams, generate, load

torch.manual_seed(0)
torch.set_num_threads(2)
model, tok = load()
window = model.cfg.context

prompt = (
    "Customer (Ana): Hi, I was charged twice for invoice INV-2026-004512 on the Business plan. Please refund it.\n"
    "Agent (Omar): Sorry about that. A billing agent will review the duplicate charge.\n"
    "Customer (Ana): Also, Slack notifications stopped arriving in our channel.\n"
    "Agent (Omar): Choose the channel again in the Slack integration settings.\n"
    "Customer (Ana): Thanks, that worked. How do I export a board to CSV?\n"
    "Agent (Omar): Open the board and choose More, then Export, then CSV.\n"
    "Customer (Ana): And where are we on my refund for the duplicate charge?\n"
    "Agent (Omar):"
)
ids = tok.encode(prompt).ids
params = SamplingParams(max_new_tokens=40, temperature=0)
kept = ids[-(window - params.max_new_tokens):]  # the same slice generate() applies
print(f"prompt tokens: {len(ids)}   window: {window}   reserved for output: {params.max_new_tokens}")
print(f"kept: last {len(kept)} tokens, dropped: first {len(ids) - len(kept)} tokens")
print(f"dropped text: {tok.decode(ids[:len(ids) - len(kept)])!r}")
print(f"kept text starts: {tok.decode(kept)[:60]!r}")
print(f"invoice id still in view: {'INV-2026-004512' in tok.decode(kept)}")

out = generate(model, tok, prompt, params)
print(f"\nTinyLM continues ({out.stop_reason}, {len(out.token_ids)} tokens): {out.text.strip()[:160]!r}")

print("\nAsking for the whole window as output:")
try:
    generate(model, tok, prompt, SamplingParams(max_new_tokens=window, temperature=0))
except IndexError as err:
    print(f"  IndexError: {err}")

  prompt tokens: 154   window: 128   reserved for output: 40
  kept: last 88 tokens, dropped: first 66 tokens
  dropped text: 'Customer (Ana): Hi, I was charged twice for invoice INV-2026-004512 on the Business plan. Please refund it.\nAgent (Omar): Sorry about that. A billing agent will review the duplicate charge.\nCustomer (Ana): Also, Slack notifications stopped arriving in our channel.\n'
  kept text starts: 'Agent (Omar): Choose the channel again in the Slack integrat'
  invoice id still in view: False

  TinyLM continues (length, 40 tokens): 'Sorry about the sign-in page to 10 business days to the original payment method.\nTwo-factor authentication (2FA) recovery: Team costs 12 USD per user per user p'

  Asking for the whole window as output:
    IndexError: index out of range in self

# examples/m02_overflow.py
"""Module 2: stay inside the context window before sending, and recover if the provider still says no."""
import httpx
import openai
from m02_context import drop_oldest
from m02_conversation import conversation

from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_messages

OVERFLOW_HINTS = ("context_length_exceeded", "context length", "maximum context", "too many tokens",
                  "exceeds the maximum", "reduce the length")


def is_context_overflow(err: Exception) -> bool:
    """Providers word this differently; check the code first, then the message."""
    if not isinstance(err, openai.BadRequestError):
        return False
    text = f"{getattr(err, 'code', '')} {err.message} {err.body}".lower()
    return any(hint in text for hint in OVERFLOW_HINTS)


def send_within_window(messages, llm, window: int, max_tokens: int, safety: float = 0.05):
    """Trim to fit before sending; on an overflow error, trim 20% harder and retry once."""
    budget = int(window * (1 - safety)) - max_tokens
    for attempt in (1, 2):
        trimmed = drop_oldest(messages, budget)
        print(f"  attempt {attempt}: sending {count_messages(trimmed)} estimated tokens "
              f"({len(messages) - len(trimmed)} messages dropped)")
        try:
            result = llm(trimmed, max_tokens=max_tokens)
        except openai.BadRequestError as err:
            if not is_context_overflow(err) or attempt == 2:
                raise
            print(f"  provider refused: {err.body['message']}")
            budget = int(budget * 0.8)
            continue
        if result.finish_reason == "length":
            print("  warning: the reply hit max_tokens and is cut off")
        return result


def fake_provider(limit: int, tokenizer_ratio: float):
    """A ScriptedLLM that behaves like a provider whose tokenizer counts more tokens than ours (plumbing only)."""
    def respond(messages, kwargs):
        real = int(count_messages(messages) * tokenizer_ratio)
        if real + kwargs["max_tokens"] > limit:
            request = httpx.Request("POST", "https://example.invalid/v1/chat/completions")
            raise openai.BadRequestError(
                "Error code: 400", response=httpx.Response(400, request=request),
                body={"message": f"This model's maximum context length is {limit} tokens. "
                                 f"Your request used {real + kwargs['max_tokens']} tokens.",
                      "type": "invalid_request_error", "code": "context_length_exceeded"})
        return "(scripted reply: request accepted)"
    return ScriptedLLM(responder=respond)


chat_log = conversation()
print(f"Conversation: {count_messages(chat_log)} estimated tokens")
print("Provider A (tokenizer counts like ours):")
reply = send_within_window(chat_log, fake_provider(limit=512, tokenizer_ratio=1.0), window=512, max_tokens=128)
print(f"  reply: {reply.text!r}")
print("Provider B (tokenizer counts 30% more than ours):")
reply = send_within_window(chat_log, fake_provider(limit=512, tokenizer_ratio=1.3), window=512, max_tokens=128)
print(f"  reply: {reply.text!r}")

# With a real provider the call is the same; only the window changes, for example:
#   from supportdesk.llm import chat
#   send_within_window(chat_log, chat, window=131_072, max_tokens=1_024)   # gpt-oss-120b on Groq

  Conversation: 558 estimated tokens
  Provider A (tokenizer counts like ours):
    attempt 1: sending 336 estimated tokens (9 messages dropped)
    reply: '(scripted reply: request accepted)'
  Provider B (tokenizer counts 30% more than ours):
    attempt 1: sending 336 estimated tokens (9 messages dropped)
    provider refused: This model's maximum context length is 512 tokens. Your request used 564 tokens.
    attempt 2: sending 272 estimated tokens (12 messages dropped)
    reply: '(scripted reply: request accepted)'

# examples/m02_needle.py
"""Module 2: a needle-in-a-haystack harness for position effects.

Run with a stand-in (default) to test the plumbing, or with a real model:
  PYTHONPATH=. python examples/m02_needle.py --real
"""
import argparse
import random
import re

from supportdesk.data import load_articles
from supportdesk.llm import Usage, chat, resolve
from supportdesk.pricing import PRICES, cost_usd
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_tokens

DEPTHS = (0.0, 0.25, 0.5, 0.75, 1.0)
QUESTION = "What is the refund approval code for workspace {ws}? Reply with the code only."
NEEDLE = "Internal billing note: the refund approval code for workspace {ws} is {code}."


def build_haystack(target_tokens: int, rng: random.Random) -> list[str]:
    """Help-center paragraphs, shuffled and repeated until the text reaches about target_tokens."""
    paragraphs = [p for a in load_articles() for p in a.body.split("\n") if p.strip()]
    out, total = [], 0
    while total < target_tokens:
        p = rng.choice(paragraphs)
        out.append(p)
        total += count_tokens(p) + 1
    return out


def make_case(target_tokens: int, depth: float, rng: random.Random) -> tuple[list[dict], str]:
    ws, code = f"ACME-{rng.randint(10, 99)}", f"PLUM-{rng.randint(1000, 9999)}"
    paragraphs = build_haystack(target_tokens, rng)
    paragraphs.insert(round(depth * len(paragraphs)), NEEDLE.format(ws=ws, code=code))
    context = "\n".join(paragraphs)
    messages = [
        {"role": "system", "content": "Answer using only the reference text."},
        {"role": "user", "content": f"Reference text:\n{context}\n\n{QUESTION.format(ws=ws)}"},
    ]
    return messages, code


def run(llm, lengths, trials: int, seed: int = 0) -> dict[tuple[int, float], float]:
    rng = random.Random(seed)
    scores = {}
    for length in lengths:
        for depth in DEPTHS:
            hits = 0
            for _ in range(trials):
                messages, code = make_case(length, depth, rng)
                # Reasoning models spend max_tokens on thinking first, so leave room beyond the short answer.
                hits += code in llm(messages, max_tokens=512, temperature=0).text
            scores[(length, depth)] = hits / trials
    return scores


def perfect_reader(messages, kwargs):
    """Stand-in that finds the code with a regex: it can only prove the plumbing works."""
    found = re.search(r"code for workspace \S+ is (PLUM-\d+)", messages[-1]["content"])
    return found.group(1) if found else "not found"


def head_and_tail_reader(messages, kwargs, keep: float = 0.3):
    """Stand-in with a planted blind spot: it only reads the first and last 30% of the text."""
    text = messages[-1]["content"]
    cut = int(len(text) * keep)
    return perfect_reader([{"role": "user", "content": text[:cut] + text[-cut:]}], kwargs)


def show(title: str, scores: dict, trials: int) -> None:
    lengths = sorted({k[0] for k in scores})
    print(f"{title}\n  share answered correctly, n = {trials} per cell")
    print(f"  {'tokens':>7}" + "".join(f"{f'depth {d:.2f}':>12}" for d in DEPTHS))
    for length in lengths:
        print(f"  {length:>7}" + "".join(f"{scores[(length, d)]:>12.2f}" for d in DEPTHS))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--real", action="store_true", help="call supportdesk.llm.chat (needs a provider key)")
    parser.add_argument("--trials", type=int, default=5)
    args = parser.parse_args()
    lengths = (2_000, 8_000, 32_000)
    if args.real:
        calls = len(lengths) * len(DEPTHS) * args.trials
        input_tokens = sum(lengths) * len(DEPTHS) * args.trials
        _, model = resolve()
        estimate = cost_usd(Usage(input_tokens=input_tokens, output_tokens=calls * 512), model) if model in PRICES else float("nan")
        print(f"{calls} calls to {model}, about {input_tokens:,} input tokens, at most {estimate:.2f} USD")
        show("Real model (supportdesk.llm.chat):", run(chat, lengths, args.trials), args.trials)
    else:
        show("Stand-in: perfect regex reader (plumbing check, not a model)",
             run(ScriptedLLM(responder=perfect_reader), lengths, args.trials), args.trials)
        show("Stand-in: planted head-and-tail blind spot (does the harness detect it?)",
             run(ScriptedLLM(responder=head_and_tail_reader), lengths, args.trials), args.trials)

  Stand-in: perfect regex reader (plumbing check, not a model)
    share answered correctly, n = 5 per cell
     tokens  depth 0.00  depth 0.25  depth 0.50  depth 0.75  depth 1.00
       2000        1.00        1.00        1.00        1.00        1.00
       8000        1.00        1.00        1.00        1.00        1.00
      32000        1.00        1.00        1.00        1.00        1.00
  Stand-in: planted head-and-tail blind spot (does the harness detect it?)
    share answered correctly, n = 5 per cell
     tokens  depth 0.00  depth 0.25  depth 0.50  depth 0.75  depth 1.00
       2000        1.00        1.00        0.00        1.00        1.00
       8000        1.00        1.00        0.00        1.00        1.00
      32000        1.00        1.00        0.00        1.00        1.00

     tokens  depth 0.00  depth 0.25  depth 0.50  depth 0.75  depth 1.00
       2000        1.00        1.00        1.00        1.00        1.00
       8000        1.00        1.00        1.00        1.00        1.00
      32000        1.00        1.00        0.80        1.00        1.00