legal-ai/scripts/backfill_lesson_synthesis.py

#!/usr/bin/env python3
"""Backfill — LLM synthesis of decision_lessons (#158 / INV-LRN8).

WHAT THIS DOES
--------------
Walks (practice_area, category) shards that hold ≥2 live APPROVED style lessons,
clusters near-duplicates (cosine), and asks a local ``claude_session`` model (Opus
by default) to merge each cluster into ONE richer, generalised "super-lesson" —
grounded in the source lessons (INV-AH) with a drift guard. Accepted merges are
written as ``source='synthesis'`` rows (review_status='approved', graduated gate)
and their sources flip to ``superseded`` (provenance, no longer writer-fed). The
writer-fed set shrinks, so its limit=15 stops truncating (#157).

All logic lives in services/lesson_synthesis.py (G2) — this is the batch driver:
shard ordering, throttling, dry-run reporting and a CSV audit trail.

IDEMPOTENCY / RESUME
--------------------
Re-running is safe: superseded sources are excluded from candidates, and an accepted
merge that matches an existing synthesis (cosine) is skipped (duplicate_skipped).

USAGE
-----
cd ~/legal-ai/mcp-server
.venv/bin/python ../scripts/backfill_lesson_synthesis.py --dry-run                       # all shards, no writes
.venv/bin/python ../scripts/backfill_lesson_synthesis.py --dry-run --practice-area rishuy_uvniya --category style
.venv/bin/python ../scripts/backfill_lesson_synthesis.py --apply                         # full throttled run
"""
from __future__ import annotations

import argparse
import asyncio
import csv
import os
import sys
from collections import Counter
from datetime import datetime, timezone

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "mcp-server", "src"))

from legal_mcp.services import db, lesson_synthesis  # noqa: E402

try:  # stdlib-only module, importable from system python too
    from legal_mcp.services import usage_limits
except Exception:  # pragma: no cover
    usage_limits = None

AUDIT_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "audit")


def _throttled() -> tuple[bool, str]:
    if usage_limits is None:
        return False, "usage_limits unavailable"
    usage = usage_limits.subscription_usage()
    if usage is None:
        return False, "usage read failed (proceeding)"
    over, _reset, detail = usage_limits.ceiling_status(usage)
    return over, detail


def _short(s: str, n: int = 100) -> str:
    s = (s or "").replace("\n", " ")
    return s if len(s) <= n else s[: n - 1] + "…"


async def _run(apply: bool, practice_area: str, category: str,
               throttle: bool, verbose: bool) -> int:
    if practice_area and category:
        shards = [{"practice_area": practice_area, "category": category}]
    else:
        shards = await db.synthesis_shards(min_size=2)
        if practice_area:
            shards = [s for s in shards if s["practice_area"] == practice_area]
        if category:
            shards = [s for s in shards if s["category"] == category]

    mode = "APPLY" if apply else "DRY-RUN"
    print(f"[{mode}] {len(shards)} shards with ≥2 approved lessons "
          f"(throttle={'on' if throttle else 'off'})\n")
    if not shards:
        print("nothing to do.")
        return 0

    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    os.makedirs(AUDIT_DIR, exist_ok=True)
    audit_path = os.path.join(
        AUDIT_DIR, f"lesson-synthesis-{'apply' if apply else 'dryrun'}-{stamp}.csv")
    counts: Counter[str] = Counter()
    stopped = False

    with open(audit_path, "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["practice_area", "category", "cluster_size", "status",
                    "drift_cosine", "applied", "reason", "after", "member_ids"])
        for n, s in enumerate(shards, 1):
            if throttle:
                over, detail = _throttled()
                if over:
                    print(f"\n⏸  usage ceiling reached ({detail}) — stopping at "
                          f"shard {n - 1}/{len(shards)}. Re-run to resume.")
                    stopped = True
                    break
            pa, cat = s["practice_area"], s["category"]
            res = await lesson_synthesis.run_shard(pa, cat, apply=apply)
            print(f"[{n}/{len(shards)}] {pa}/{cat}: {res['candidates']} candidates → "
                  f"{len(res['clusters'])} clusters")
            for c in res["clusters"]:
                counts[c["status"]] += 1
                w.writerow([pa, cat, len(c["members"]), c["status"],
                            c.get("drift_cosine"), c.get("applied"),
                            c.get("reason", ""), c.get("proposed", ""),
                            "|".join(c["members"])])
                mark = {"accepted": "✓", "duplicate_skipped": "=", "abstained": "·",
                        "drift_rejected": "✗", "llm_error": "!", "too_small": "·"}.get(c["status"], "?")
                print(f"      {mark} {c['status']:<18} size={len(c['members'])} "
                      f"drift={c.get('drift_cosine')}{' [written]' if c.get('applied') else ''}")
                if verbose and c.get("proposed"):
                    print(f"        → {_short(c['proposed'])}")

    processed = sum(counts.values())
    print(f"\n── summary ({mode}) — {processed} clusters"
          f"{' (stopped early)' if stopped else ''} ──")
    for status, c in counts.most_common():
        print(f"   {status:<18} {c}")
    print(f"\naudit CSV: {audit_path}")
    if not apply:
        print("dry-run — nothing written. Re-run with --apply to commit.")
    return 0


def main() -> int:
    p = argparse.ArgumentParser(description="LLM synthesis of decision_lessons (#158 / INV-LRN8)")
    p.add_argument("--apply", action="store_true", help="commit to the DB (default: dry-run)")
    p.add_argument("--dry-run", action="store_true", help="explicit dry-run (default)")
    p.add_argument("--practice-area", default="", help="limit to one practice_area")
    p.add_argument("--category", default="", help="limit to one category (style/structure/lexicon/tabular)")
    p.add_argument("--no-throttle", action="store_true", help="skip usage-ceiling checks")
    p.add_argument("--verbose", action="store_true", help="print merged text per cluster")
    args = p.parse_args()
    return asyncio.run(_run(
        apply=args.apply, practice_area=args.practice_area, category=args.category,
        throttle=not args.no_throttle, verbose=args.verbose,
    ))


if __name__ == "__main__":
    raise SystemExit(main())