Source code for src.orchestrator

"""Unified runner for GDELT and configured HTML/Scooper scrapers.

The orchestrator owns cross-pipeline CLI concerns so operators can run small
smoke tests or larger backfills from one command while each pipeline keeps its
source-specific defaults and implementation details.
"""

from __future__ import annotations

import argparse
import datetime
import json
import math
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor
import pandas as pd
from pathlib import Path

from .cli_reporter import CliReporter, PipelineStats
from .logging_utils import get_file_logger
from .GDELT.gdelt_seeds import backfill_seeds, fetch_gkg_links
from .GDELT.runner import load_seen, write_output_records
from .shared_utils import (
    DEBUG_DIR,
    NoiseCollector,
    clear_directory,
    collect_as_completed,
    ensure_model_available,
    exit_if_shutdown,
    get_config_bool,
    get_config_int,
    get_config_value,
    install_handler,
    model_unavailable_error,
    run_clean,
    shutdown_executor,
)
# from scripts.clean_gdelt import run_clean

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
GDELT_CACHE_DIR = _PROJECT_ROOT / "data" / "gdelt_cache"
_NORMALIZE_SCRIPT = (
    _PROJECT_ROOT / "tests" / "json_normilization" / "json_validation.py"
)
_DEDUP_SCRIPT = _PROJECT_ROOT / "scripts" / "dedup_file.py"
_DEDUP_OUTPUT_DIR = _PROJECT_ROOT / "data" / "output"
_SCOOPER_JSON = _PROJECT_ROOT / "data" / "processed" / "scooper.json"

LOG_FILE = _PROJECT_ROOT / "data" / "logs" / "orchestrator.log"
LOGGER = get_file_logger(__name__, LOG_FILE)
AI_MODEL = get_config_value("AI_MODEL", "llama3.2:latest")

try:
    from .supabase_function import has_supabase_creds, push_vulnerabilities

    SUPABASE_AVAILABLE = has_supabase_creds()
    if not SUPABASE_AVAILABLE:
        LOGGER.info("SUPABASE_URL or SUPABASE_KEY missing; DB writes disabled")
except Exception as e:
    LOGGER.info("Supabase unavailable, DB writes disabled: %s", e)
    SUPABASE_AVAILABLE = False
    push_vulnerabilities = None  # type: ignore[assignment]


def _split_date(
    start: datetime.date, end: datetime.date, k: int
) -> list[tuple[datetime.date, datetime.date]]:
    """
    Splits the date into K parts, where K is the number of threads
    NOTE: only used when both dates are passed into scooper

    Args:
        start:
        end:
        k: Number of threads

    Returns:
        Start-End window / Number of threads
    """
    if end > start:
        LOGGER.debug("date range given oldest-first; normalizing %s..%s", start, end)
        start, end = end, start

    num_days = (start - end).days + 1  # inclusive day count

    k = max(1, min(k, num_days))
    base, extra = divmod(num_days, k)

    windows: list[tuple[datetime.date, datetime.date]] = []
    cursor = end
    for i in range(k):
        size = base + (1 if i < extra else 0)
        window_end = cursor
        window_start = cursor + datetime.timedelta(days=size - 1)
        windows.append((window_start, window_end))
        cursor = window_start + datetime.timedelta(days=1)

    windows.reverse()
    return windows


def _parse_date(s: str | None) -> datetime.date | None:
    """Parse a YYYY-MM-DD or YYYYMMDD string into a date, or return None."""
    if s is None:
        return None
    for fmt in ("%Y-%m-%d", "%Y%m%d"):
        try:
            return datetime.datetime.strptime(s, fmt).date()
        except ValueError:
            continue
    LOGGER.warning("Could not parse date %r; ignoring for HTML engine", s)
    return None


def _post_process_and_push(
    json_files: list[Path],
    reporter: CliReporter,
    use_sb: bool = False,
) -> None:
    """Normalize, dedup, then optionally push processed JSON files to Supabase.

    Runs ``json_validation.py -n``, then ``dedup_file.py``. When ``use_sb`` is
    true and credentials are present, also runs ``push_vulnerabilities`` on each
    resulting ``*_d.json``.
    """
    if not use_sb:
        return
    files = [p.resolve() for p in json_files if p.exists()]
    if not files:
        return

    reporter.phase("Normalize, dedup, and push to Supabase")

    for path in files:
        reporter.info(f"Normalizing {path}")
        result = subprocess.run(
            [sys.executable, str(_NORMALIZE_SCRIPT), "-n", str(path)],
            cwd=str(_PROJECT_ROOT),
        )
        if result.returncode != 0:
            LOGGER.warning(
                "json_validation -n failed for %s (exit %s)", path, result.returncode
            )

    reporter.info(f"Deduplicating {len(files)} file(s)")
    result = subprocess.run(
        [sys.executable, str(_DEDUP_SCRIPT), *[str(p) for p in files]],
        cwd=str(_PROJECT_ROOT),
    )
    if result.returncode != 0:
        LOGGER.warning("dedup_file failed (exit %s)", result.returncode)
        reporter.info("Dedup failed; skipping Supabase push")
        return

    # TODO remove this or the if stmt on  line 128
    if not use_sb:
        reporter.info("Skipping Supabase push (--use-sb not set)")
        return

    if not SUPABASE_AVAILABLE or push_vulnerabilities is None:
        reporter.info("Supabase unavailable; skipping push")
        return

    for path in files:
        deduped = _DEDUP_OUTPUT_DIR / f"{path.stem}_d.json"
        if not deduped.exists():
            LOGGER.warning("Deduped file missing: %s", deduped)
            continue
        sources = json.loads(deduped.read_text(encoding="utf-8")).get("sources") or []
        if not sources:
            continue
        try:
            n = push_vulnerabilities(sources)
            reporter.info(f"Pushed {n} records from {deduped.name}")
        except Exception as e:
            LOGGER.warning("Supabase push failed for %s: %s", deduped, e)
            reporter.info(f"Supabase push failed for {deduped.name}: {e}")


[docs] def chunk_list(items, num_chunks): """ Split a list of items into a specified number of chunks, as evenly as possible. Args: items: The list of items to split. num_chunks: The number of chunks to create. Returns: A list of lists, where each sublist is a chunk of the original items. """ if num_chunks is None or num_chunks <= 0: num_chunks = 1 if not items: return [] chunk_size = math.ceil(len(items) / num_chunks) if chunk_size <= 0: return [] return [items[i : i + chunk_size] for i in range(0, len(items), chunk_size)]
def _collect_gdelt_seeds( *, num_files: int, start_date: str | None, end_date: str | None, cache_dir: Path, reporter: CliReporter, stats: PipelineStats, seed_threads: int, sector: str | None = None, ) -> list[dict]: """Collect GDELT seeds, optionally splitting work across threads. Parameters: num_files: Number of GDELT GKG files to scan. start_date: Ceiling date (YYYYMMDD or YYYY-MM-DD) for articles; newer articles are skipped. end_date: Floor date (YYYYMMDD or YYYY-MM-DD) for articles; crawling stops at older articles. cache_dir: Directory to cache downloaded GKG files. reporter: CliReporter for logging progress and warnings. stats: PipelineStats for recording statistics. seed_threads: Number of threads to use for seed collection. sector: Sector key into SECTOR_THEMES (default: health). Returns: A list of candidate seed records (dicts) collected from GDELT. """ if start_date is not None and end_date is not None: num_days = (_parse_date(start_date) - _parse_date(end_date)).days + 1 reporter.info( f"Collecting GDELT seeds for date range {_parse_date(start_date)} - {_parse_date(end_date)} ({num_days} days)" ) else: reporter.info( f"Collecting GDELT seeds for {num_files} files ({num_files / 4} hours)" ) seed_threads = max(1, seed_threads) if seed_threads > 1: reporter.info(f"Collecting GDELT seeds with {seed_threads} threads") if seed_threads == 1: LOGGER.debug("Collecting GDELT seeds in single-threaded mode") return list( backfill_seeds( num_files=num_files, start_date=start_date, end_date=end_date, cache_dir=cache_dir, reporter=reporter, stats=stats, sector=sector, ) ) quiet_reporter = CliReporter(verbose=False) raw_seeds: list[dict] = [] # Resolve the concrete GKG file list for either a date range or a recent # file count, then split that list across threads. Splitting by file (not # by calendar day) honors seed_threads even when the requested date range # spans fewer days than threads. links = fetch_gkg_links( num_files=num_files, start_date=start_date, end_date=end_date ) chunks = chunk_list(links, seed_threads) if not chunks: return [] executor = ThreadPoolExecutor(max_workers=len(chunks)) try: futures = [] # Submit tasks for each chunk of links for link_chunk in chunks: futures.append( executor.submit( backfill_seeds, links=link_chunk, cache_dir=cache_dir, reporter=quiet_reporter, sector=sector, ) ) collect_as_completed(futures, raw_seeds.extend) finally: shutdown_executor(executor) LOGGER.debug( "Collected %d GDELT seeds across %d link chunks with %d threads", len(raw_seeds), len(chunks), seed_threads, ) return raw_seeds
[docs] def main(argv: list[str] | None = None) -> int: """Parse CLI options, run selected pipeline stages, and report summaries. Args: argv: Optional argument list for tests and programmatic callers. When omitted, argparse reads from the process command line. Returns: Process exit code. A successful orchestrated run returns ``0``. """ parser = argparse.ArgumentParser( description="Unified runner for GDELT and HTML scrapers" ) # Shared parser.add_argument( "--use-bert", "-b", action="store_true", default=get_config_bool("USE_BERT", False), help="Run BERT pre-filter before LLM field extraction in both pipelines", ) parser.add_argument( "--skip-gdelt", action="store_true", default=get_config_bool("SKIP_GDELT", False), help="Skip the GDELT pipeline", ) parser.add_argument( "--skip-html", action="store_true", default=get_config_bool("SKIP_HTML", False), help="Skip the HTML scraper pipeline", ) parser.add_argument( "--verbose", "-v", action="store_true", default=get_config_bool("VERBOSE", False), help="Show detailed per-article pipeline output", ) parser.add_argument( "--debug", "-d", action="store_true", default=get_config_bool("DEBUG", False), help="Log all rejected/skipped articles (noise) to JSON files in data/noise/", ) parser.add_argument( "--start-date", default=get_config_value("START_DATE", None), help=( "Ceiling date (YYYYMMDD or YYYY-MM-DD): articles newer than this are " "skipped. Applied to both GDELT files and HTML article dates." ), ) parser.add_argument( "--end-date", default=get_config_value("END_DATE", None), help=( "Floor date (YYYYMMDD or YYYY-MM-DD): crawling stops at articles older " "than this. Applied to both GDELT files and HTML article dates." ), ) parser.add_argument( "--use-sb", action="store_true", default=get_config_bool("USE_SB", False), help="Use supabase along with local JSON", ) # GDELT-specific parser.add_argument( "--num-files", "-n", type=int, default=get_config_int("GDELT_NUM_FILES", 2), help="GDELT GKG files to scan (default: 2)", ) parser.add_argument( "--limit", "-l", type=int, default=get_config_int("GDELT_LIMIT", None), help="Cap on seeds to process; defaults to 3 unless --num-files is provided", ) parser.add_argument( "--output-path", "-o", default=get_config_value("OUTPUT_PATH", "data/output/results.json"), help="Output JSON file or directory for GDELT results", ) parser.add_argument( "--seen-urls-file", default=get_config_value("SEEN_URLS_FILE", None), help="Path to store/load seen URLs JSON file", ) parser.add_argument( "--clean", action="store_true", default=get_config_bool("CLEAN", False), help="Clear all modified directories and files before running", ) parser.add_argument( "--gdelt-seed-threads", type=int, default=get_config_int("GDELT_SEED_THREADS", 1), help="Number of threads to use for GDELT seed collection (default: 1)", ) parser.add_argument( "--models", type=int, default=get_config_int("MODELS", 1), help=("Number of model instances to run concurrently."), ) parser.add_argument( "--threads-per-model", type=int, default=get_config_int("THREADS_PER_MODEL", 1), help=("Number of threads to use per model instance."), ) parser.add_argument( "--starting-port", type=int, default=get_config_int("STARTING_PORT", 11434), help=( "Starting port number for LLM instances. Each instance is expected to " "run on a consecutive port (e.g. 11434, 11435, etc.)" ), ) parser.add_argument( "--seeds_only", action="store_true", default=get_config_bool("SEEDS_ONLY", False), help="Process only seed articles, skipping full scraping and processing. Also skips the scooper pipeline.", ) parser.add_argument( "--sector", default=get_config_value("SECTOR", "health"), help="The sector to process (default: health)", ) start = time.time() args = parser.parse_args(argv) install_handler() reporter = CliReporter(verbose=args.verbose) summaries: list[PipelineStats] = [] # Overall progress bar: one unit of work per pipeline phase that will run. # When only one phase runs, the overall bar would just mirror the phase bar, # so we suppress it entirely. phases = int(not args.skip_gdelt) + int(not args.skip_html) show_overall = phases > 1 if show_overall: reporter.set_overall_total(phases) reporter.set_overall_step("Initializing") threads = max(1, args.models) * max(1, args.threads_per_model) json_files: list[Path] = [] if not args.skip_gdelt: import src.GDELT.runner as runner gdelt_start = time.time() n_provided = ( any(opt in sys.argv[1:] for opt in ("-n", "--num-files")) or args.num_files is not None ) l_provided = args.limit is not None effective_limit = args.limit if not l_provided: effective_limit = None if n_provided else 3 gdelt_noise = ( NoiseCollector(DEBUG_DIR / "debug_noise_gdelt.json") if args.debug else None ) gdelt_stats = PipelineStats("GDELT") reporter.phase(f"Running GDELT pipeline for sector: {args.sector}") if show_overall: reporter.set_overall_step("GDELT") LOGGER.info("Running GDELT pipeline with args: %s", args) if args.clean: run_clean() raw_seeds = _collect_gdelt_seeds( num_files=args.num_files, start_date=args.start_date, end_date=args.end_date, cache_dir=GDELT_CACHE_DIR, reporter=reporter, stats=gdelt_stats, seed_threads=args.gdelt_seed_threads, sector=args.sector, ) LOGGER.info( f"Seed collection complete in {(time.time() - gdelt_start) / 60:.2f} minutes" ) gdelt_processing_start = time.time() if args.seeds_only: LOGGER.info("Seeds-only mode enabled; skipping full GDELT processing") exit(0) seen = load_seen(args.seen_urls_file) chunks = chunk_list(raw_seeds, threads) port = args.starting_port if not chunks: chunks = [[]] try: ensure_model_available() reporter.info(f"LLM model: {AI_MODEL}") except model_unavailable_error as exc: LOGGER.error("Model availability check failed: %s", exc) print(exc, file=sys.stderr) return 1 # Size the phase bar to the seeds that will actually process (each worker # applies effective_limit to its own chunk). Per-seed advances from the # workers then drive both the phase bar and, smoothly, the overall bar. gdelt_units = sum( min(len(c), effective_limit) if effective_limit else len(c) for c in chunks ) reporter.start_phase("GDELT", total=gdelt_units) reporter.info(f"\nProcessing {effective_limit} seeds") if threads > 1: reporter.info(f"Models: {args.models}") reporter.info(f"Threads per model: {args.threads_per_model}") reporter.info(f"Total threads: {threads}") all_records = [] executor = ThreadPoolExecutor(max_workers=threads) try: futures = [] n = 0 for chunk in chunks: futures.append( executor.submit( runner.run, num_files=args.num_files, limit=effective_limit, output_path=None, start_date=args.start_date, end_date=args.end_date, seen=seen, use_bert=args.use_bert, verbose=args.verbose, reporter=reporter, stats=PipelineStats("GDELT"), raw_seeds=chunk, debug_noise=gdelt_noise, port=port, clear_seeds=False, ) ) n += 1 if n == args.threads_per_model: port += 1 n = 0 def _merge_gdelt(result): worker_stats, records = result gdelt_stats.merge(worker_stats) all_records.extend(records) collect_as_completed(futures, _merge_gdelt) finally: shutdown_executor(executor) if not gdelt_stats.paused: clear_directory(runner.SEEDS_DIR) reporter.detail(f"Cleared seed staging directory: {runner.SEEDS_DIR}") LOGGER.debug("Cleared seeds directory: %s", runner.SEEDS_DIR) else: reporter.detail(f"Preserved seed staging directory: {runner.SEEDS_DIR}") out_file = write_output_records( all_records, args.output_path, reporter, gdelt_stats ) json_files.append(out_file) if gdelt_noise: out = gdelt_noise.flush() if out: reporter.info(f"Debug noise (GDELT): {out}") LOGGER.info( f"GDELT processing complete in {(time.time() - gdelt_processing_start) / 60:.2f} minutes" ) # Whole-phase wall clock for the run summary. gdelt_start is set # at the top of the GDELT branch; gdelt_processing_start above only covers # the LLM-processing sub-phase. gdelt_stats.elapsed_seconds = time.time() - gdelt_start summaries.append(gdelt_stats) if gdelt_stats.paused: reporter.info("GDELT pipeline paused; skipping remaining pipelines.") _post_process_and_push(json_files, reporter, use_sb=args.use_sb) reporter.summary(summaries) LOGGER.info("GDELT pipeline paused; skipping remaining pipelines") return exit_if_shutdown(0) if not args.skip_html: import src.HTML.scooper as scooper html_start = time.time() html_stats = PipelineStats("HTML") reporter.phase("Running HTML/Scooper pipeline") if show_overall: reporter.set_overall_step("HTML") LOGGER.info("Running HTML/Scooper pipeline with args %s", args) scooper.setup_scooper(sb_only=args.use_sb) vuln_dfs: list[pd.DataFrame] = [] noise_dfs: list[pd.DataFrame] = [] # K split — one scooper instance per date window, run in parallel. start_date = _parse_date(args.start_date) end_date = _parse_date(args.end_date) if start_date is not None and end_date is not None: dates: list[tuple[datetime.date, datetime.date]] = _split_date( end_date, start_date, threads ) # One phase unit per date window; advanced as each window completes. reporter.start_phase("HTML", total=len(dates)) # One scooper instance per date window port = args.starting_port results = [] try: ensure_model_available() reporter.info(f"LLM model: {AI_MODEL}") except model_unavailable_error as exc: LOGGER.error("Model availability check failed: %s", exc) print(exc, file=sys.stderr) return 1 executor = ThreadPoolExecutor(max_workers=len(dates)) try: futures = [] n = 0 for win_start, win_end in dates: futures.append( executor.submit( scooper.run_scooper, use_bert=args.use_bert, verbose=args.verbose, start_date=win_start, end_date=win_end, reporter=reporter, stats=PipelineStats("HTML"), # per-window; merged below sb_only=args.use_sb, port=port, ) ) n += 1 if n == args.threads_per_model: port += 1 n = 0 def _append_html(result): results.append(result) reporter.advance(1) collect_as_completed(futures, _append_html) finally: shutdown_executor(executor) vuln_lists: list = [] for window_stats, w_vuln_list, v_df, n_df in results: html_stats.merge(window_stats) vuln_dfs.append(v_df) noise_dfs.append(n_df) vuln_lists.extend(w_vuln_list) scooper.save_results(vuln_lists, vuln_dfs, noise_dfs, sb_only=args.use_sb) # Default: one thread per site. run_scooper fans out internally and # returns frames merged across sites (disjoint); we persist them here. # Per-site HTML progress lives inside scooper (deferred to #229); treat # the whole phase as one unit for now. else: reporter.start_phase("HTML", total=1) html_stats, vuln_list, v_df, n_df = scooper.run_scooper( use_bert=args.use_bert, verbose=args.verbose, start_date=_parse_date(args.end_date), # swapped here end_date=_parse_date(args.start_date), # swapped here reporter=reporter, stats=html_stats, sb_only=args.use_sb, site_split=True, ) scooper.save_results(vuln_list, [v_df], [n_df], sb_only=args.use_sb) reporter.advance(1) json_files.append(_SCOOPER_JSON) html_stats.elapsed_seconds = time.time() - html_start summaries.append(html_stats) if html_stats.paused: reporter.info("HTML scraper paused; skipping remaining pipelines.") _post_process_and_push(json_files, reporter, use_sb=args.use_sb) reporter.summary(summaries) LOGGER.info("HTML scraper paused; skipping remaining pipelines") return exit_if_shutdown(0) LOGGER.info( f"HTML/Scooper processing complete in {(time.time() - html_start) / 60:.2f} minutes" ) _post_process_and_push(json_files, reporter, use_sb=args.use_sb) if summaries: reporter.summary(summaries) LOGGER.info("Orchestrator run complete with summaries: %s", summaries) LOGGER.info(f"Total execution time: {(time.time() - start) / 60:.2f} minutes") return exit_if_shutdown(0)
if __name__ == "__main__": raise SystemExit(main())