scripts/build_kcc_testcases.py

"""Convert a curated KCC Animal Husbandry slice into CeRAI DataPoints.json.

KCC is field data: questions and answers are often in Romanized Hindi /
mixed-script — *we keep those rows*, that's actually closer to how a farmer
would phrase a query to Amul AI. We do strip out:
  - rows whose answer is essentially a phone number (very common — "call
    1800-180-5141 for cattle queries")
  - empty / single-word queries
  - non-dairy/animal-husbandry rows (broad keyword filter, English + romanized
    Hindi terms)

Reads:  data/processed/kcc_animal_husbandry.parquet
Writes: data/processed/kcc_dairy_datapoints.json
        data/processed/kcc_dairy_plan.json

Usage:
    python scripts/build_kcc_testcases.py --max-rows 20
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "data" / "processed" / "kcc_animal_husbandry.parquet"
OUT_DATAPOINTS = ROOT / "data" / "processed" / "kcc_dairy_datapoints.json"
OUT_PLAN = ROOT / "data" / "processed" / "kcc_dairy_plan.json"

# Broad dairy / animal-husbandry vocabulary spanning English + romanized Hindi.
# `\b` word boundaries so we don't false-match substrings.
KEYWORDS = re.compile(
    r"\b("
    # English
    r"cow|cows|bull|bulls|buffalo|buffaloes|calf|calves|cattle|bovine|"
    r"dairy|milk|milking|milker|udder|teat|mastitis|colostrum|"
    r"lactation|heifer|holstein|jersey|gir|sahiwal|murrah|"
    r"ghee|curd|butterfat|fodder|silage|chara|deworming|insemination|ai|"
    r"fmd|brucellosis|theileria|hoof|"
    # Romanized Hindi / common KCC terms
    r"pasu|pashu|gay|gaay|bhains|bhainsh|doodh|dudh|dudhha|dugdh|"
    r"bachhda|bachhra|bachra|bachhdi|"
    r"thaan|thanela|paseena|thanaila|"
    r"ghas|charae|chare|charaa|"
    r"hara|sukha|"
    r"khurpaka|munhpaka|"
    r"sankarn|sangraman|sankran|"
    r"khali|khaali"
    r")\b",
    re.IGNORECASE,
)

# Match if the answer is mostly digits/punctuation/whitespace — i.e. a bare
# phone number, an extension code, etc.
PHONE_ONLY = re.compile(r"^[\s\d\-+()/.,]*$")
# Match an isolated phone number anywhere — KCC has lots of "call XXXXXXX"
# answers; we drop those too.
HAS_LONG_NUMBER = re.compile(r"\b\d{6,}\b")

METRICS_AND_STRATEGY: dict[str, tuple[str, list[str]]] = {
    "14": ("Lexical_Diversity",   ["18"]),
    "38": ("BLEU",                ["42"]),
    "39": ("ROUGE",               ["42"]),
    "40": ("METEOR",              ["42"]),
    "41": ("Turn_Around_Time",    ["25"]),
    "47": ("Error_Rate",          ["13"]),
}
METRICS = {met_id: meta[0] for met_id, meta in METRICS_AND_STRATEGY.items()}

SYSTEM_PROMPT = (
    "You are an expert assistant for Indian dairy and animal-husbandry farmers. "
    "The farmer may write in English, Hindi, or romanized Hindi (Hinglish). "
    "Answer the question accurately and concisely."
)


def _looks_like_real_answer(ans: str) -> bool:
    a = ans.strip()
    if len(a) < 30:
        return False
    if PHONE_ONLY.match(a):
        return False
    if HAS_LONG_NUMBER.search(a):
        return False
    if " " not in a:
        return False
    return True


def _looks_like_real_query(q: str) -> bool:
    return len(q.strip()) >= 8


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-rows", type=int, default=20)
    parser.add_argument("--seed", type=int, default=11)
    args = parser.parse_args()

    df = pd.read_parquet(SRC)
    n0 = len(df)
    df = df.dropna(subset=["QueryText", "KccAns"])
    df = df[df["QueryText"].apply(_looks_like_real_query)]
    df = df[df["KccAns"].apply(_looks_like_real_answer)]
    df = df[df["QueryText"].astype(str).str.contains(KEYWORDS, regex=True)
            | df["KccAns"].astype(str).str.contains(KEYWORDS, regex=True)]
    df = df.drop_duplicates(subset=["QueryText"]).reset_index(drop=True)

    print(f"From {n0:,} rows ➜ {len(df):,} after filters (dairy-relevant, "
          f"non-phone answer, dedup)")
    if len(df) == 0:
        raise SystemExit("filter is too strict; relax KEYWORDS or thresholds.")

    df = df.sample(min(args.max_rows, len(df)), random_state=args.seed).reset_index(drop=True)

    base_cases = []
    for i, row in df.iterrows():
        q = str(row["QueryText"]).strip()
        a = str(row["KccAns"]).strip()
        base_cases.append({
            "PROMPT_ID": f"KCC-AH-{i:04d}",
            "LLM_AS_JUDGE": "No",
            "SYSTEM_PROMPT": SYSTEM_PROMPT,
            "PROMPT": q,
            "EXPECTED_OUTPUT": a,
            "DOMAIN": "animal_husbandry",
        })

    # One testcase row per (prompt, metric) — see the long comment in
    # build_cerai_testcases.py. Same prompt across metric copies, distinct
    # strategy per copy. Executor's cache dedupes API calls by user_prompt.
    datapoints: dict[str, dict] = {}
    for met_id, (_, strategy_ids) in METRICS_AND_STRATEGY.items():
        suffixed = []
        for case in base_cases:
            c = dict(case)
            c["PROMPT_ID"] = f"{case['PROMPT_ID']}-M{met_id}"
            c["STRATEGY"] = strategy_ids
            suffixed.append(c)
        datapoints[met_id] = {"cases": suffixed}

    plan = {
        "T_AmulKCC": {
            "TestPlan_name": "Amul_KCC_AnimalHusbandry",
            "metrics": METRICS,
        }
    }

    OUT_DATAPOINTS.write_text(json.dumps(datapoints, indent=2, ensure_ascii=False))
    OUT_PLAN.write_text(json.dumps(plan, indent=2, ensure_ascii=False))

    print(f"Wrote {OUT_DATAPOINTS} ({len(base_cases)} base cases x {len(METRICS)} metrics)")
    print(f"Wrote {OUT_PLAN}")
    print()
    print("--- 4 sample cases ---")
    for c in base_cases[:4]:
        print(f"Q: {c['PROMPT']}")
        print(f"A: {c['EXPECTED_OUTPUT']}")
        print()


if __name__ == "__main__":
    main()