Spaces:

VAILL
/

legislation-tracker

Running on CPU Upgrade

legislation-tracker / backend /update_data.py

Deploy: unbuffer pipeline logs

4aa353b 10 days ago

11.3 kB

"""
update_data.py
--------------
CLI pipeline orchestrator for updating all legislation tracker data.
Usage:
    # Fully non-interactive (for scripts / CI):
    python update_data.py --pull --overwrite-pdf --skip-upload
    # Pull new data, skip HF upload, auto-continue on errors:
    python update_data.py --pull --overwrite-pdf --skip-upload --continue-on-error
    # Use existing local data only (no LegiScan pull):
    python update_data.py --no-pull --skip-upload
    # Resume from a specific step (1-indexed):
    python update_data.py --start-from 9 --skip-upload --continue-on-error
    # Interactive (prompts like before):
    python update_data.py
"""

import argparse
import subprocess
import sys
import os

import dotenv
dotenv.load_dotenv()

from data_updating_scripts.pipeline_status import PipelineStatus

ALL_SCRIPTS = [
    "data_updating_scripts/get_data.py",
    "data_updating_scripts/fix_pdf_bills.py",
    "data_updating_scripts/known_bills_status.py",
    "data_updating_scripts/migrate_iapp_categories.py",
    "data_updating_scripts/classify_bills.py",
    "data_updating_scripts/mark_no_text_bills.py",
    "data_updating_scripts/generate_summaries.py",
    "data_updating_scripts/generate_suggested_questions.py",
    "data_updating_scripts/generate_reports.py",
    "data_updating_scripts/build_bills_vectorstore.py",
    "data_updating_scripts/eu_vectorstore.py",
    "data_updating_scripts/detect_changes.py",
    "data_updating_scripts/generate_news_enrichment.py",
    "data_updating_scripts/build_calendar.py",
    "data_updating_scripts/generate_newsletter.py",
]


def _ask(prompt: str, default: bool = False) -> bool:
    """Ask a yes/no question interactively."""
    suffix = " [Y/n]: " if default else " [y/N]: "
    try:
        answer = input(prompt + suffix).strip().lower()
    except EOFError:
        return default
    if not answer:
        return default
    return answer in ("y", "yes")


LOG_FILE = "data_updating_scripts/logs/fetch_ai_bills.log"


def _upload_pipeline_log() -> None:
    """Upload the pipeline log to HF Dataset for remote debugging."""
    import os as _os
    from pathlib import Path

    log = Path(LOG_FILE)
    if not log.exists():
        return
    try:
        from huggingface_hub import HfApi
        token = _os.getenv("HUGGINGFACE_HUB_TOKEN")
        repo = _os.getenv("HF_REPO_ID")
        if not token or not repo:
            return
        api = HfApi(token=token)
        # Upload last 500KB of log to keep size manageable
        content = log.read_bytes()
        if len(content) > 500_000:
            content = content[-500_000:]
        tmp = Path("/tmp/pipeline_log_tail.log")
        tmp.write_bytes(content)
        api.upload_file(
            path_or_fileobj=str(tmp),
            path_in_repo="logs/fetch_ai_bills.log",
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message="Update pipeline log",
        )
        print("✓ Pipeline log uploaded to HF")
    except Exception as e:
        print(f"⚠ Failed to upload pipeline log: {e}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run the full AI legislation data-update pipeline."
    )
    pull_group = parser.add_mutually_exclusive_group()
    pull_group.add_argument(
        "--pull", action="store_true", default=None,
        help="Pull fresh data from LegiScan (costs API quota).",
    )
    pull_group.add_argument(
        "--no-pull", action="store_true",
        help="Skip LegiScan pull; use existing local data.",
    )
    parser.add_argument(
        "--overwrite-pdf", action="store_true",
        help="Overwrite known_bills.json after fixing PDF bills.",
    )
    parser.add_argument(
        "--skip-upload", action="store_true",
        help="Skip HuggingFace upload at the end.",
    )
    parser.add_argument(
        "--continue-on-error", action="store_true",
        help="Continue running remaining scripts if one fails.",
    )
    parser.add_argument(
        "--test", action="store_true",
        help="Test mode: only pull a few bills from CA to verify the pipeline works.",
    )
    parser.add_argument(
        "--test-count", type=int, default=3,
        help="Number of bills to pull in test mode (default: 3).",
    )
    parser.add_argument(
        "--timeout", type=int, default=7200,
        help="Per-script timeout in seconds (default: 7200 = 2 hours).",
    )
    parser.add_argument(
        "--start-from", type=int, default=None,
        help="Start from step N (1-indexed). Skips earlier steps.",
    )
    parser.add_argument(
        "--stop-after", type=int, default=None,
        help="Stop after step N (1-indexed). Skips later steps.",
    )
    parser.add_argument(
        "--reset-vectorstore", action="store_true",
        help="Delete all Pinecone vectors and re-embed from scratch.",
    )
    args = parser.parse_args()

    # --start-from / --stop-after overrides pull logic
    start_idx = 0
    if args.start_from is not None:
        start_idx = args.start_from - 1
        if start_idx < 0 or start_idx >= len(ALL_SCRIPTS):
            print(f"Invalid --start-from: must be 1-{len(ALL_SCRIPTS)}")
            sys.exit(1)

    stop_idx = len(ALL_SCRIPTS)
    if args.stop_after is not None:
        stop_idx = args.stop_after
        if stop_idx < 1 or stop_idx > len(ALL_SCRIPTS):
            print(f"Invalid --stop-after: must be 1-{len(ALL_SCRIPTS)}")
            sys.exit(1)

    if args.start_from is not None:
        scripts_to_run = ALL_SCRIPTS[start_idx:stop_idx]
        print(f"\n> Running steps {start_idx + 1}-{stop_idx}: {', '.join(scripts_to_run)}")
    elif args.pull or args.test:
        print("\n> Will pull new data from LegiScan")
        scripts_to_run = ALL_SCRIPTS
    elif args.no_pull:
        print("\n> Skipping data pull, using existing local data")
        scripts_to_run = ALL_SCRIPTS[2:]  # skip get_data + fix_pdf_bills
    else:
        # Interactive fallback
        do_pull = _ask("Pull new data from LegiScan?")
        if do_pull:
            print("\n> Will pull new data from LegiScan")
            scripts_to_run = ALL_SCRIPTS
        else:
            print("\n> Skipping data pull, using existing local data")
            scripts_to_run = ALL_SCRIPTS[2:]  # skip get_data + fix_pdf_bills

    print(f"\nWill run {len(scripts_to_run)} scripts:")
    for script in scripts_to_run:
        print(f"  - {script}")

    # Download existing data from HuggingFace so generate scripts can skip
    # already-processed bills (avoids re-generating reports/summaries after rebuild)
    print("\nDownloading existing data from HuggingFace...")
    try:
        from huggingface_upload import download_from_huggingface
        downloaded = download_from_huggingface()
        print(f"✓ Downloaded {len(downloaded)} files from HuggingFace")
    except Exception as e:
        print(f"⚠ HuggingFace download skipped: {e}")

    # Initialize pipeline status tracking
    status = PipelineStatus()
    triggered_by = os.environ.get("PIPELINE_TRIGGERED_BY", "cli")
    status.start_pipeline(scripts_to_run, triggered_by=triggered_by)

    # Build environment for child processes
    env = os.environ.copy()
    # Child stdout is redirected to a log file rather than a terminal, so
    # Python block-buffers it and output only lands when the buffer fills or
    # the process exits. That makes a long-running step look hung and hides
    # the error when one dies early. Unbuffer so the log tails live.
    env["PYTHONUNBUFFERED"] = "1"
    if args.overwrite_pdf:
        env["FIX_PDF_OVERWRITE"] = "yes"
    if args.test:
        env["PIPELINE_TEST_MODE"] = "true"
        env["PIPELINE_TEST_MAX_BILLS"] = str(args.test_count)
        print(f"\n>>> TEST MODE: only pulling {args.test_count} bill(s) from CA to verify pipeline <<<\n")

    # Ensure log directory exists for capturing subprocess output
    os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

    print("\n" + "=" * 50)

    for script in scripts_to_run:
        print(f"\n--- Running {script} ---")
        print("=" * 50)

        # Mark step as running and tell child script its identity
        status.start_step(script)
        env["PIPELINE_CURRENT_SCRIPT"] = script

        # Give generate_reports.py a much longer timeout since each report
        # takes 1-3 min and there can be hundreds of new bills to process.
        script_timeout = 86400 if "generate_reports" in script else args.timeout

        try:
            # Build command with any script-specific flags
            cmd = [sys.executable, script]
            if args.reset_vectorstore and "build_bills_vectorstore" in script:
                cmd.append("--reset")

            # Capture all subprocess output to the pipeline log file
            # so errors are visible when the log is uploaded to HF Dataset.
            with open(LOG_FILE, "a") as log_fh:
                log_fh.write(f"\n{'='*50}\n--- {script} ---\n{'='*50}\n")
                log_fh.flush()
                result = subprocess.run(
                    cmd, env=env,
                    timeout=script_timeout,
                    stdout=log_fh, stderr=subprocess.STDOUT,
                )
        except subprocess.TimeoutExpired:
            status.fail_step(script, message=f"Timed out after {args.timeout}s")
            print(f"\n✗ Script {script} timed out after {args.timeout}s")
            if args.continue_on_error:
                print("  (--continue-on-error set, moving on)")
                continue
            if _ask("Continue with the remaining scripts?"):
                continue
            print("Stopping pipeline execution.")
            status.fail_pipeline("Stopped by user after script timeout")
            sys.exit(1)

        if result.returncode != 0:
            status.fail_step(script, message=f"Exit code {result.returncode}")
            print(f"\n✗ Script {script} failed with return code {result.returncode}")
            if args.continue_on_error:
                print("  (--continue-on-error set, moving on)")
                continue
            # Interactive fallback
            if _ask("Continue with the remaining scripts?"):
                continue
            print("Stopping pipeline execution.")
            status.fail_pipeline("Stopped by user after script failure")
            sys.exit(1)
        else:
            status.complete_step(script)
            print(f"✓ {script} completed successfully")

    # Mark pipeline as complete
    status.complete_pipeline()

    print("\n" + "=" * 50)
    print("✓ Pipeline execution completed!")

    # HuggingFace upload
    if args.skip_upload:
        print("\n> Skipping HuggingFace upload (--skip-upload)")
    else:
        print("\nUploading all JSON datasets to HuggingFace...")
        try:
            from huggingface_upload import upload_all_to_huggingface
            upload_all_to_huggingface()
            print("✓ HuggingFace upload complete!")
        except KeyError as e:
            print(f"✗ HuggingFace config error: {e}")
        except Exception as e:
            print(f"✗ Upload failed: {e}")

        # Upload pipeline log so it's accessible remotely
        _upload_pipeline_log()


if __name__ == "__main__":
    main()

	"""
	update_data.py
	--------------
	CLI pipeline orchestrator for updating all legislation tracker data.

	Usage:
	# Fully non-interactive (for scripts / CI):
	python update_data.py --pull --overwrite-pdf --skip-upload

	# Pull new data, skip HF upload, auto-continue on errors:
	python update_data.py --pull --overwrite-pdf --skip-upload --continue-on-error

	# Use existing local data only (no LegiScan pull):
	python update_data.py --no-pull --skip-upload

	# Resume from a specific step (1-indexed):
	python update_data.py --start-from 9 --skip-upload --continue-on-error

	# Interactive (prompts like before):
	python update_data.py
	"""

	import argparse
	import subprocess
	import sys
	import os

	import dotenv
	dotenv.load_dotenv()

	from data_updating_scripts.pipeline_status import PipelineStatus

	ALL_SCRIPTS = [
	"data_updating_scripts/get_data.py",
	"data_updating_scripts/fix_pdf_bills.py",
	"data_updating_scripts/known_bills_status.py",
	"data_updating_scripts/migrate_iapp_categories.py",
	"data_updating_scripts/classify_bills.py",
	"data_updating_scripts/mark_no_text_bills.py",
	"data_updating_scripts/generate_summaries.py",
	"data_updating_scripts/generate_suggested_questions.py",
	"data_updating_scripts/generate_reports.py",
	"data_updating_scripts/build_bills_vectorstore.py",
	"data_updating_scripts/eu_vectorstore.py",
	"data_updating_scripts/detect_changes.py",
	"data_updating_scripts/generate_news_enrichment.py",
	"data_updating_scripts/build_calendar.py",
	"data_updating_scripts/generate_newsletter.py",
	]


	def _ask(prompt: str, default: bool = False) -> bool:
	"""Ask a yes/no question interactively."""
	suffix = " [Y/n]: " if default else " [y/N]: "
	try:
	answer = input(prompt + suffix).strip().lower()
	except EOFError:
	return default
	if not answer:
	return default
	return answer in ("y", "yes")


	LOG_FILE = "data_updating_scripts/logs/fetch_ai_bills.log"


	def _upload_pipeline_log() -> None:
	"""Upload the pipeline log to HF Dataset for remote debugging."""
	import os as _os
	from pathlib import Path

	log = Path(LOG_FILE)
	if not log.exists():
	return
	try:
	from huggingface_hub import HfApi
	token = _os.getenv("HUGGINGFACE_HUB_TOKEN")
	repo = _os.getenv("HF_REPO_ID")
	if not token or not repo:
	return
	api = HfApi(token=token)
	# Upload last 500KB of log to keep size manageable
	content = log.read_bytes()
	if len(content) > 500_000:
	content = content[-500_000:]
	tmp = Path("/tmp/pipeline_log_tail.log")
	tmp.write_bytes(content)
	api.upload_file(
	path_or_fileobj=str(tmp),
	path_in_repo="logs/fetch_ai_bills.log",
	repo_id=repo,
	repo_type="dataset",
	token=token,
	commit_message="Update pipeline log",
	)
	print("✓ Pipeline log uploaded to HF")
	except Exception as e:
	print(f"⚠ Failed to upload pipeline log: {e}")


	def main() -> None:
	parser = argparse.ArgumentParser(
	description="Run the full AI legislation data-update pipeline."
	)
	pull_group = parser.add_mutually_exclusive_group()
	pull_group.add_argument(
	"--pull", action="store_true", default=None,
	help="Pull fresh data from LegiScan (costs API quota).",
	)
	pull_group.add_argument(
	"--no-pull", action="store_true",
	help="Skip LegiScan pull; use existing local data.",
	)
	parser.add_argument(
	"--overwrite-pdf", action="store_true",
	help="Overwrite known_bills.json after fixing PDF bills.",
	)
	parser.add_argument(
	"--skip-upload", action="store_true",
	help="Skip HuggingFace upload at the end.",
	)
	parser.add_argument(
	"--continue-on-error", action="store_true",
	help="Continue running remaining scripts if one fails.",
	)
	parser.add_argument(
	"--test", action="store_true",
	help="Test mode: only pull a few bills from CA to verify the pipeline works.",
	)
	parser.add_argument(
	"--test-count", type=int, default=3,
	help="Number of bills to pull in test mode (default: 3).",
	)
	parser.add_argument(
	"--timeout", type=int, default=7200,
	help="Per-script timeout in seconds (default: 7200 = 2 hours).",
	)
	parser.add_argument(
	"--start-from", type=int, default=None,
	help="Start from step N (1-indexed). Skips earlier steps.",
	)
	parser.add_argument(
	"--stop-after", type=int, default=None,
	help="Stop after step N (1-indexed). Skips later steps.",
	)
	parser.add_argument(
	"--reset-vectorstore", action="store_true",
	help="Delete all Pinecone vectors and re-embed from scratch.",
	)
	args = parser.parse_args()

	# --start-from / --stop-after overrides pull logic
	start_idx = 0
	if args.start_from is not None:
	start_idx = args.start_from - 1
	if start_idx < 0 or start_idx >= len(ALL_SCRIPTS):
	print(f"Invalid --start-from: must be 1-{len(ALL_SCRIPTS)}")
	sys.exit(1)

	stop_idx = len(ALL_SCRIPTS)
	if args.stop_after is not None:
	stop_idx = args.stop_after
	if stop_idx < 1 or stop_idx > len(ALL_SCRIPTS):
	print(f"Invalid --stop-after: must be 1-{len(ALL_SCRIPTS)}")
	sys.exit(1)

	if args.start_from is not None:
	scripts_to_run = ALL_SCRIPTS[start_idx:stop_idx]
	print(f"\n> Running steps {start_idx + 1}-{stop_idx}: {', '.join(scripts_to_run)}")
	elif args.pull or args.test:
	print("\n> Will pull new data from LegiScan")
	scripts_to_run = ALL_SCRIPTS
	elif args.no_pull:
	print("\n> Skipping data pull, using existing local data")
	scripts_to_run = ALL_SCRIPTS[2:] # skip get_data + fix_pdf_bills
	else:
	# Interactive fallback
	do_pull = _ask("Pull new data from LegiScan?")
	if do_pull:
	print("\n> Will pull new data from LegiScan")
	scripts_to_run = ALL_SCRIPTS
	else:
	print("\n> Skipping data pull, using existing local data")
	scripts_to_run = ALL_SCRIPTS[2:] # skip get_data + fix_pdf_bills

	print(f"\nWill run {len(scripts_to_run)} scripts:")
	for script in scripts_to_run:
	print(f" - {script}")

	# Download existing data from HuggingFace so generate scripts can skip
	# already-processed bills (avoids re-generating reports/summaries after rebuild)
	print("\nDownloading existing data from HuggingFace...")
	try:
	from huggingface_upload import download_from_huggingface
	downloaded = download_from_huggingface()
	print(f"✓ Downloaded {len(downloaded)} files from HuggingFace")
	except Exception as e:
	print(f"⚠ HuggingFace download skipped: {e}")

	# Initialize pipeline status tracking
	status = PipelineStatus()
	triggered_by = os.environ.get("PIPELINE_TRIGGERED_BY", "cli")
	status.start_pipeline(scripts_to_run, triggered_by=triggered_by)

	# Build environment for child processes
	env = os.environ.copy()
	# Child stdout is redirected to a log file rather than a terminal, so
	# Python block-buffers it and output only lands when the buffer fills or
	# the process exits. That makes a long-running step look hung and hides
	# the error when one dies early. Unbuffer so the log tails live.
	env["PYTHONUNBUFFERED"] = "1"
	if args.overwrite_pdf:
	env["FIX_PDF_OVERWRITE"] = "yes"
	if args.test:
	env["PIPELINE_TEST_MODE"] = "true"
	env["PIPELINE_TEST_MAX_BILLS"] = str(args.test_count)
	print(f"\n>>> TEST MODE: only pulling {args.test_count} bill(s) from CA to verify pipeline <<<\n")

	# Ensure log directory exists for capturing subprocess output
	os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

	print("\n" + "=" * 50)

	for script in scripts_to_run:
	print(f"\n--- Running {script} ---")
	print("=" * 50)

	# Mark step as running and tell child script its identity
	status.start_step(script)
	env["PIPELINE_CURRENT_SCRIPT"] = script

	# Give generate_reports.py a much longer timeout since each report
	# takes 1-3 min and there can be hundreds of new bills to process.
	script_timeout = 86400 if "generate_reports" in script else args.timeout

	try:
	# Build command with any script-specific flags
	cmd = [sys.executable, script]
	if args.reset_vectorstore and "build_bills_vectorstore" in script:
	cmd.append("--reset")

	# Capture all subprocess output to the pipeline log file
	# so errors are visible when the log is uploaded to HF Dataset.
	with open(LOG_FILE, "a") as log_fh:
	log_fh.write(f"\n{'='50}\n--- {script} ---\n{'='50}\n")
	log_fh.flush()
	result = subprocess.run(
	cmd, env=env,
	timeout=script_timeout,
	stdout=log_fh, stderr=subprocess.STDOUT,
	)
	except subprocess.TimeoutExpired:
	status.fail_step(script, message=f"Timed out after {args.timeout}s")
	print(f"\n✗ Script {script} timed out after {args.timeout}s")
	if args.continue_on_error:
	print(" (--continue-on-error set, moving on)")
	continue
	if _ask("Continue with the remaining scripts?"):
	continue
	print("Stopping pipeline execution.")
	status.fail_pipeline("Stopped by user after script timeout")
	sys.exit(1)

	if result.returncode != 0:
	status.fail_step(script, message=f"Exit code {result.returncode}")
	print(f"\n✗ Script {script} failed with return code {result.returncode}")
	if args.continue_on_error:
	print(" (--continue-on-error set, moving on)")
	continue
	# Interactive fallback
	if _ask("Continue with the remaining scripts?"):
	continue
	print("Stopping pipeline execution.")
	status.fail_pipeline("Stopped by user after script failure")
	sys.exit(1)
	else:
	status.complete_step(script)
	print(f"✓ {script} completed successfully")

	# Mark pipeline as complete
	status.complete_pipeline()

	print("\n" + "=" * 50)
	print("✓ Pipeline execution completed!")

	# HuggingFace upload
	if args.skip_upload:
	print("\n> Skipping HuggingFace upload (--skip-upload)")
	else:
	print("\nUploading all JSON datasets to HuggingFace...")
	try:
	from huggingface_upload import upload_all_to_huggingface
	upload_all_to_huggingface()
	print("✓ HuggingFace upload complete!")
	except KeyError as e:
	print(f"✗ HuggingFace config error: {e}")
	except Exception as e:
	print(f"✗ Upload failed: {e}")

	# Upload pipeline log so it's accessible remotely
	_upload_pipeline_log()


	if __name__ == "__main__":
	main()