scripts/generate_precomputed_examples.py

"""Generate bundled precomputed assets for example bills."""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, UTC
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config import (
    DEFAULT_CHUNK_OVERLAP,
    DEFAULT_CHUNK_SIZE,
    DEFAULT_EMBEDDING_MODEL,
    DEFAULT_PROVIDER,
    DEFAULT_QWEN_MODEL,
)
from services.example_bills import EXAMPLE_BILL_MANIFEST_PATH, ExampleBill, load_example_bills
from services.ingest import fetch_url_content
from services.precomputed_assets import write_precomputed_asset
from services.providers import get_default_api_key, instantiate_client
from services.rag_pipeline import document_hash_for_text, generate_analysis_once, split_into_chunks


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--bill-id",
        action="append",
        dest="bill_ids",
        help="Limit generation to one or more example bill ids.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Regenerate even when a bill already has bundled assets.",
    )
    parser.add_argument(
        "--fetch-timeout",
        type=int,
        default=120,
        help="HTTP timeout in seconds when downloading source bills.",
    )
    return parser.parse_args()


def _load_manifest() -> list[dict[str, object]]:
    return json.loads(EXAMPLE_BILL_MANIFEST_PATH.read_text(encoding="utf-8"))


def _write_manifest(items: list[dict[str, object]]) -> None:
    EXAMPLE_BILL_MANIFEST_PATH.write_text(json.dumps(items, indent=2) + "\n", encoding="utf-8")


def _has_complete_assets(bill: ExampleBill) -> bool:
    return (
        bill.analysis_path.exists()
        and bill.document_path.exists()
        and bill.chunks_path.exists()
        and bill.metadata_path.exists()
    )


def _metadata(document_hash: str, source_url: str) -> dict[str, object]:
    return {
        "generated_at": datetime.now(UTC).isoformat(),
        "provider": DEFAULT_PROVIDER,
        "model": DEFAULT_QWEN_MODEL,
        "source_url": source_url,
        "document_hash": document_hash,
        "chunk_size": DEFAULT_CHUNK_SIZE,
        "chunk_overlap": DEFAULT_CHUNK_OVERLAP,
        "embedding_model": DEFAULT_EMBEDDING_MODEL,
    }


def main() -> int:
    args = _parse_args()
    manifest_items = _load_manifest()
    manifest_by_id = {str(item["id"]): item for item in manifest_items}
    all_bills = load_example_bills()
    selected = [bill for bill in all_bills if not args.bill_ids or bill.id in set(args.bill_ids)]

    if args.bill_ids:
        missing = sorted(set(args.bill_ids) - {bill.id for bill in selected})
        if missing:
            raise SystemExit(f"Unknown bill ids: {', '.join(missing)}")

    api_key = get_default_api_key(DEFAULT_PROVIDER)
    if not api_key:
        raise SystemExit("HF_TOKEN is required to generate precomputed example analysis.")
    provider_client = instantiate_client(DEFAULT_PROVIDER, api_key)

    failures: list[tuple[str, str]] = []
    for bill in selected:
        if _has_complete_assets(bill) and not args.force:
            print(f"Skipping {bill.id}: bundled assets already exist.")
            continue

        try:
            print(f"Generating assets for {bill.id}...")
            document_text = fetch_url_content(bill.source_url, timeout_seconds=args.fetch_timeout)
            chunks = split_into_chunks(document_text)
            analysis = generate_analysis_once(provider_client, document_text)

            document_hash = document_hash_for_text(document_text)
            write_precomputed_asset(
                bill,
                analysis_payload=analysis.model_dump(mode="json"),
                document_text=document_text,
                chunks=chunks,
                metadata=_metadata(document_hash, bill.source_url),
            )

            manifest_entry = manifest_by_id[bill.id]
            manifest_entry["document_hash"] = document_hash
            print(f"Finished {bill.id}.")
        except Exception as exc:  # noqa: BLE001
            failures.append((bill.id, str(exc)))
            print(f"Failed {bill.id}: {exc}")

    _write_manifest(manifest_items)
    if failures:
        print("\nFailures:")
        for bill_id, message in failures:
            print(f"- {bill_id}: {message}")
        return 1
    return 0


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