scripts/validate_gguf.py

"""
Validate the GGUF model's SQL quality against the held-out test set BEFORE deploying.

For each example in data/test_dataset.jsonl:
  1. Generate SQL from the GGUF model (greedy / temperature 0).
  2. Execute both predicted and expected SQL on data/company_sales.db.
  3. Count as a match if results are EXECUTION-EQUIVALENT (same rows, order-independent).

Acceptance bar: >= 95% match. If below, re-quantize at Q5_K_M and re-run.

Default backend: llama-server (recommended on Windows/local).
  The harness launches llama-server.exe (from the llama.cpp release you already
  downloaded), waits for it to load the model ONCE, then hits its OpenAI-compatible
  /v1/chat/completions endpoint over HTTP -- the same call shape production uses.
  This avoids:
    - the prebuilt llama-cpp-python wheel (illegal-instruction 0xc000001d on CPUs
      without the AVX-512 it was built for), and
    - llama-cli interactive/conversation-mode hangs.

Usage:
    python scripts/validate_gguf.py --gguf models/phi3-text-to-sql-Q4_K_M.gguf
    python scripts/validate_gguf.py --gguf models/...gguf --server .\\llama-b9637-bin-win-cpu-x64\\llama-server.exe
    python scripts/validate_gguf.py --gguf models/...gguf --use-python   # only if the wheel runs on your CPU

Requires: data/company_sales.db (run `python src/database.py` first if missing).
"""
import argparse
import glob
import json
import os
import re
import sqlite3
import subprocess
import sys
import time
import urllib.error
import urllib.request

DEFAULT_GGUF = "models/phi3-text-to-sql-Q4_K_M.gguf"
TEST_PATH = "data/test_dataset.jsonl"
DB_PATH = "data/company_sales.db"


def clean_sql(raw_output):
    """Same extraction logic as src/inference.py so validation matches serving."""
    if "<|assistant|>" in raw_output:
        raw_output = raw_output.split("<|assistant|>")[-1]
    cleaned = re.sub(r"```sql\s*", "", raw_output, flags=re.IGNORECASE)
    cleaned = re.sub(r"```", "", cleaned)
    cleaned = re.sub(r"<\|.*?\|>", "", cleaned).strip()
    m = re.search(r"(SELECT\s+.*?;)", cleaned, re.DOTALL | re.IGNORECASE)
    if m:
        return m.group(1).strip()
    if "select" in cleaned.lower():
        return cleaned.strip()
    return cleaned.strip()


def run_sql(conn, sql):
    """Returns (rows_as_sorted_list, error_or_None)."""
    if not sql or "select" not in sql.lower():
        return None, "invalid/blank SQL"
    try:
        cur = conn.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
        return sorted([tuple(r) for r in rows], key=lambda r: str(r)), None
    except Exception as e:
        return None, str(e)


def load_test_set(path):
    items = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            msgs = obj["messages"]
            user = next(m["content"] for m in msgs if m["role"] == "user")
            expected = next(m["content"] for m in msgs if m["role"] == "assistant")
            question = user.split("Question:", 1)[-1].strip() if "Question:" in user else user
            items.append({"user_prompt": user, "question": question, "expected_sql": expected})
    return items


# ---------- llama-server backend ----------

def find_llama_server():
    candidates = []
    candidates += glob.glob(os.path.join("llama-b*", "llama-server.exe"))
    candidates += glob.glob(os.path.join("llama-b*-bin-win-*", "llama-server.exe"))
    candidates += glob.glob(os.path.join("llama-bin", "llama-server.exe"))
    for name in ("llama-server.exe", "llama-server"):
        if os.path.exists(name):
            candidates.append(name)
    return candidates[0] if candidates else None


def start_server(server_path, gguf, n_ctx, port, log_path="llama-server.log"):
    cmd = [server_path, "-m", gguf, "-c", str(n_ctx),
           "--host", "127.0.0.1", "--port", str(port)]
    print(f"Starting llama-server (logs -> {log_path}) ...")
    log = open(log_path, "w", encoding="utf-8")
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)

    health = f"http://127.0.0.1:{port}/health"
    for _ in range(180):  # up to ~3 min for first model load
        if proc.poll() is not None:
            raise RuntimeError(f"llama-server exited early; see {log_path}")
        try:
            with urllib.request.urlopen(health, timeout=2) as r:
                if r.status == 200 and json.loads(r.read()).get("status") == "ok":
                    print("Server ready.")
                    return proc
        except Exception:
            pass
        time.sleep(1)
    proc.terminate()
    raise RuntimeError(f"llama-server did not become ready in time; see {log_path}")


def make_http_generator(port):
    url = f"http://127.0.0.1:{port}/v1/chat/completions"

    def gen(user_content):
        body = json.dumps({
            "messages": [{"role": "user", "content": user_content}],
            "temperature": 0.0,
            "max_tokens": 150,
        }).encode("utf-8")
        req = urllib.request.Request(url, data=body, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=180) as r:
            out = json.loads(r.read())
        return out["choices"][0]["message"]["content"]
    return gen


def make_python_generator(gguf, n_ctx):
    from llama_cpp import Llama
    llm = Llama(model_path=gguf, n_ctx=n_ctx, n_threads=os.cpu_count() or 2, verbose=False)

    def gen(user_content):
        out = llm.create_chat_completion(
            messages=[{"role": "user", "content": user_content}],
            max_tokens=150, temperature=0.0,
        )
        return out["choices"][0]["message"]["content"]
    return gen


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gguf", default=DEFAULT_GGUF, help="Path to the GGUF file")
    ap.add_argument("--threshold", type=float, default=0.95, help="Acceptance match rate")
    ap.add_argument("--n_ctx", type=int, default=2048)
    ap.add_argument("--server", default=None, help="Path to llama-server(.exe). Auto-detected if omitted.")
    ap.add_argument("--port", type=int, default=8080)
    ap.add_argument("--use-python", action="store_true", help="Use llama-cpp-python bindings instead of llama-server")
    args = ap.parse_args()

    if not os.path.exists(args.gguf):
        sys.exit(f"GGUF not found: {args.gguf}")
    if not os.path.exists(TEST_PATH):
        sys.exit(f"Test set not found: {TEST_PATH}")
    if not os.path.exists(DB_PATH):
        sys.exit(f"Database not found: {DB_PATH}. Run `python src/database.py` first.")

    server_proc = None
    try:
        if args.use_python:
            print("Backend: llama-cpp-python bindings")
            generate = make_python_generator(args.gguf, args.n_ctx)
        else:
            server = args.server or find_llama_server()
            if not server:
                sys.exit("Could not find llama-server(.exe). Pass --server <path>, "
                         "or use --use-python if the wheel runs on your CPU.")
            print(f"Backend: llama-server -> {server}")
            server_proc = start_server(server, args.gguf, args.n_ctx, args.port)
            generate = make_http_generator(args.port)

        tests = load_test_set(TEST_PATH)
        conn = sqlite3.connect(DB_PATH)

        passed = 0
        failures = []
        for i, t in enumerate(tests, 1):
            raw = generate(t["user_prompt"])
            pred_sql = clean_sql(raw)

            exp_rows, exp_err = run_sql(conn, t["expected_sql"])
            pred_rows, pred_err = run_sql(conn, pred_sql)

            if exp_err is not None:
                ok = pred_sql.strip().lower().rstrip(";") == t["expected_sql"].strip().lower().rstrip(";")
            else:
                ok = (pred_err is None) and (pred_rows == exp_rows)

            status = "PASS" if ok else "FAIL"
            if ok:
                passed += 1
            else:
                failures.append({
                    "n": i, "question": t["question"],
                    "expected": t["expected_sql"], "predicted": pred_sql, "pred_err": pred_err,
                })
            print(f"[{i:02d}/{len(tests)}] {status}  {t['question'][:70]}")

        conn.close()

        rate = passed / len(tests) if tests else 0.0
        print("\n" + "=" * 60)
        print(f"Execution-equivalent: {passed}/{len(tests)} = {rate:.1%}")
        print(f"Acceptance bar      : {args.threshold:.0%}")
        print("Result              :", "ACCEPTED" if rate >= args.threshold else "BELOW BAR (try Q5_K_M)")
        print("=" * 60)

        if failures:
            print("\nFailures:")
            for f in failures:
                print(f"\n  Q{f['n']}: {f['question']}")
                print(f"    expected : {f['expected']}")
                print(f"    predicted: {f['predicted']}")
                if f["pred_err"]:
                    print(f"    sql error: {f['pred_err']}")

        exit_code = 0 if rate >= args.threshold else 1
    finally:
        if server_proc is not None:
            server_proc.terminate()
            try:
                server_proc.wait(timeout=10)
            except Exception:
                server_proc.kill()

    sys.exit(exit_code)


if __name__ == "__main__":
    main()