stack/dev/scripts/classify_comments.py

"""Classify skin-substitute-relevant rulemaking comments (#254).

Pipeline per docket:
  1. Pull every comment item for the docket from bib — full extracted
     text (the "Comment text" note, i.e. combined.md) when present,
     inline body (abstract) otherwise.
  2. Keep the skin-substitute-relevant subset (rex.comments.classify
     RELEVANCE_PAT) — OPPS rules draw thousands of comments on
     unrelated provisions.
  3. LLM-classify each relevant comment (position / themes /
     stakeholder / provisions / commenter+org) via prisma.llm.
     Resumable: results append to a JSONL cache keyed by comment key;
     re-runs only classify new comments.
  4. Rebuild the docket's slice of skin_subs.rulemaking_comments,
     merging the coordination cache (analyze_coordination.py) when
     present.

Usage:
    uv run python dev/scripts/classify_comments.py --docket CMS-2025-0304
    uv run python dev/scripts/classify_comments.py --limit 5 --dry-run
"""

from __future__ import annotations

import argparse
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
BIB_PATH = ROOT / "data" / "bib.sqlite"

DEFAULT_DOCKET = "CMS-2025-0304"
DEFAULT_MODEL = "claude-haiku-4-5"


def cache_paths(docket: str) -> tuple[Path, Path]:
    base = ROOT / "data" / "cms"
    return (
        base / f"comments-classified-{docket}.jsonl",
        base / f"comments-coordination-{docket}.jsonl",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--docket", default=DEFAULT_DOCKET)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--limit", type=int, default=0, help="cap LLM calls this run")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="relevance-filter and report only; no LLM calls, no DB write",
    )
    parser.add_argument(
        "--no-load", action="store_true", help="classify but skip the DuckDB load"
    )
    args = parser.parse_args()

    import json

    from rex.comments.classify import is_relevant
    from rex.comments.table import load_comments, read_cache

    comments = load_comments(args.docket, BIB_PATH)
    relevant = [c for c in comments if is_relevant(c["text"])]
    print(f"{args.docket}: {len(comments)} comments, {len(relevant)} skin-sub relevant")
    if args.dry_run:
        for c in relevant[:10]:
            print(f"  {c['comment_id'][:40]:40s} {len(c['text']):>7} chars")
        return 0

    classified_path, coordination_path = cache_paths(args.docket)
    classified_path.parent.mkdir(parents=True, exist_ok=True)
    cache = read_cache(classified_path)
    todo = [c for c in relevant if c["key"] not in cache]
    if args.limit:
        todo = todo[: args.limit]
    print(f"cached: {len(cache)}, to classify: {len(todo)}")

    if todo:
        import os

        os.environ["PRISMA_LLM_MODEL"] = args.model
        from prisma.llm import make_provider
        from rex.comments.classify import classify

        provider = make_provider()
        errors = 0
        with classified_path.open("a") as fh:
            for i, c in enumerate(todo, 1):
                try:
                    rec = classify(provider, c["text"], title=c["comment_id"])
                except Exception as e:  # noqa: BLE001 — skip and continue the batch
                    errors += 1
                    print(f"  [{i}/{len(todo)}] {c['comment_id'][:36]} ERROR: {e}")
                    continue
                rec.update(key=c["key"], docket_id=args.docket)
                cache[c["key"]] = rec
                fh.write(json.dumps(rec) + "\n")
                fh.flush()
                print(
                    f"  [{i}/{len(todo)}] {c['comment_id'][:36]:36s} "
                    f"{rec['position']:16s} {rec['stakeholder_type']}"
                )
        if errors:
            print(f"{errors} comments failed to classify (rerun to retry)")

    if args.no_load:
        return 0

    from conf.connect import duckdb_batch
    from rex.comments.table import load_table

    with duckdb_batch("aco") as con:
        n = load_table(
            con,
            args.docket,
            relevant,
            classified=cache,
            coordination=read_cache(coordination_path),
        )
        print(f"skin_subs.rulemaking_comments: {n} rows for {args.docket}")
    return 0


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