File size: 4,576 Bytes
4076efa 6d78ee1 4076efa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """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())
|