data/prepare_triage_data.py

"""Generate & curate the triage classification dataset.

Why this script exists in its current form
------------------------------------------
The original ``data/raw/triage_refund_raw.jsonl`` was a duplicate of the
billing file — it contained **zero** refund examples. Training a 3-class
classifier on 2 classes silently caps accuracy and corrupts the confusion
matrix. This script regenerates clean, balanced, de-duplicated data for all
three categories and writes a stratified train/test split.

Pipeline
--------
1. For each category, call Claude (``claude-sonnet-4-6``) in batches of 20 to
   synthesise realistic NovaPay support messages with urgency + sentiment.
2. Validate every record against the allowed label vocabulary.
3. De-duplicate on the customer message.
4. Persist per-category raw files, then a stratified 80/20 split to
   ``data/processed/{triage_train,triage_test}.jsonl``.

Run:  ``python data/prepare_triage_data.py --per-category 400``
"""
from __future__ import annotations

import argparse
import json
import random
from collections import Counter
from pathlib import Path

from common.config import RAW_DIR, PROCESSED_DIR, settings
from common.llm import llm, LLMError
from common.logging_utils import get_logger
from common.parsing import extract_json, coerce_label

logger = get_logger("prepare_triage")

INSTRUCTION = "Classify the following customer query and extract urgency and sentiment"
BATCH_SIZE = 20

CATEGORY_HINTS = {
    "refund": (
        "disputed transactions, failed UPI transfers where money was debited, "
        "cashback not received, accidental double payments, refund delays"
    ),
    "technical": (
        "app login failures, biometric/fingerprint errors, OTP not received, "
        "app crashing, statement download failures, supported OS versions"
    ),
    "billing": (
        "EMI deduction errors, premium plan charges (Plus 199, Pro 499), "
        "interest calculation disputes, loan repayment schedule queries, invoices"
    ),
}

SYSTEM_PROMPT = (
    "You are generating training data for a fintech customer support triage "
    "model. Generate realistic customer queries for a digital banking app "
    "called NovaPay. Vary the phrasing, urgency, and tone. Output only a JSON "
    "array with no markdown formatting."
)


def _user_prompt(category: str) -> str:
    return (
        f"Generate {BATCH_SIZE} unique customer support messages for NovaPay in "
        f"the category '{category}' (themes: {CATEGORY_HINTS[category]}). For each, "
        "return a JSON object with fields: instruction (always set to "
        f"'{INSTRUCTION}'), input (the customer message), output (a JSON STRING "
        "with keys category, urgency, sentiment). category must be "
        f"'{category}'. urgency is one of low/medium/high. sentiment is one of "
        "neutral/frustrated/angry/calm. Mix urgency levels. Make the messages "
        "sound like real users texting support."
    )


def _validate(record: dict, category: str) -> dict | None:
    """Coerce a raw record into the canonical schema, or drop it."""
    msg = (record.get("input") or "").strip()
    if not msg:
        return None
    out = record.get("output")
    out = extract_json(out) if isinstance(out, str) else out
    if not isinstance(out, dict):
        return None

    canonical_out = {
        "category": category,  # trust the requested category, not the model
        "urgency": coerce_label(out.get("urgency"), settings.urgencies, "medium"),
        "sentiment": coerce_label(out.get("sentiment"), settings.sentiments, "neutral"),
    }
    return {
        "instruction": INSTRUCTION,
        "input": msg,
        "output": json.dumps(canonical_out, ensure_ascii=False),
    }


def generate_category(category: str, target: int) -> list[dict]:
    """Generate ``target`` validated, de-duplicated examples for one category."""
    seen: set[str] = set()
    records: list[dict] = []
    n_batches = (target + BATCH_SIZE - 1) // BATCH_SIZE

    for i in range(n_batches):
        try:
            result = llm.complete(
                system=SYSTEM_PROMPT,
                user=_user_prompt(category),
                model=settings.primary_model,
                max_tokens=2048,
                temperature=1.0,  # high diversity for synthetic data
            )
        except LLMError as exc:
            logger.error("[%s] batch %d failed: %s", category, i + 1, exc)
            continue

        parsed = extract_json(result.text)
        if not isinstance(parsed, list):
            logger.warning("[%s] batch %d returned non-list output, skipping", category, i + 1)
            continue

        for raw in parsed:
            rec = _validate(raw, category) if isinstance(raw, dict) else None
            if rec and rec["input"].lower() not in seen:
                seen.add(rec["input"].lower())
                records.append(rec)

        logger.info("[%s] batch %d/%d -> %d unique so far", category, i + 1, n_batches, len(records))
        if len(records) >= target:
            break

    return records[:target]


def stratified_split(by_category: dict[str, list[dict]], test_frac: float = 0.2):
    rng = random.Random(settings.random_seed)
    train, test = [], []
    for records in by_category.values():
        records = records[:]
        rng.shuffle(records)
        n_test = max(1, int(len(records) * test_frac))
        test.extend(records[:n_test])
        train.extend(records[n_test:])
    rng.shuffle(train)
    rng.shuffle(test)
    return train, test


def _write_jsonl(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")


def main(per_category: int) -> None:
    by_category: dict[str, list[dict]] = {}
    for category in settings.categories:
        logger.info("Generating %d examples for '%s'…", per_category, category)
        recs = generate_category(category, per_category)
        by_category[category] = recs
        _write_jsonl(RAW_DIR / f"triage_{category}_raw.jsonl", recs)
        logger.info("[%s] wrote %d examples", category, len(recs))

    train, test = stratified_split(by_category)
    _write_jsonl(PROCESSED_DIR / "triage_train.jsonl", train)
    _write_jsonl(PROCESSED_DIR / "triage_test.jsonl", test)

    logger.info("Train: %d | Test: %d", len(train), len(test))
    logger.info("Train class balance: %s", dict(Counter(json.loads(r["output"])["category"] for r in train)))
    logger.info("Test  class balance: %s", dict(Counter(json.loads(r["output"])["category"] for r in test)))


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Generate NovaPay triage training data.")
    ap.add_argument("--per-category", type=int, default=400, help="examples per category")
    args = ap.parse_args()
    main(args.per_category)