scripts/04b_extract_edges.py

"""Step 4b: Offline edge extraction from train cascade chains (v0.2 BFS RAG).

Reads ``cascade_chains/{event_id}.json`` for every event in the train split
and writes per-event edge JSON arrays under ``data/processed/cascade_edges/``,
plus a flat ``cascade_edges_index.json`` summary. **Train split only** — test
events are skipped to keep them out of any v0.2 retrieval index.

Then ingests the freshly written edges into the ``cascade_edges`` ChromaDB
collection at ``data/vectordb_v2/`` so the BFS predictor sees them. v0.3
issue #57 caught the prior split (extract-then-forget-to-ingest) — the two
steps now run as one to remove that footgun. Pass ``--no-ingest`` to keep
the legacy extract-only behaviour.
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

from src.llm.client import load_config
from src.models.schemas import CascadeChain
from src.rag.fragment_extractor import extract_edges
from src.rag.ingestion import build_edge_vectordb

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

EDGES_DIR = Path("data/processed/cascade_edges")
EDGES_INDEX = Path("data/processed/cascade_edges_index.json")


def _load_train_event_ids(train_events_path: Path) -> set[str]:
    with train_events_path.open() as f:
        events = json.load(f)
    return {ev["event_id"] for ev in events}


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--no-ingest",
        action="store_true",
        help="Skip the ChromaDB ingestion step; produce edge JSON files only.",
    )
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    config = load_config()
    chains_dir = Path(config["paths"]["cascade_chains_dir"])
    train_events_path = Path(config["paths"]["train_events"])
    train_ids = _load_train_event_ids(train_events_path)

    EDGES_DIR.mkdir(parents=True, exist_ok=True)

    index: list[dict] = []
    total_edges = 0

    for chain_file in sorted(chains_dir.glob("*.json")):
        event_id = chain_file.stem
        if event_id not in train_ids:
            logger.info("skip non-train event: %s", event_id)
            continue

        with chain_file.open() as f:
            chain = CascadeChain.model_validate(json.load(f))

        edges = extract_edges(chain)
        out_path = EDGES_DIR / f"{event_id}.json"
        with out_path.open("w") as f:
            json.dump([e.model_dump() for e in edges], f, indent=2, ensure_ascii=False)

        num_first = sum(1 for e in edges if e.is_first_level)
        num_node = len(edges) - num_first
        index.append(
            {
                "event_id": event_id,
                "num_edges": len(edges),
                "num_first_level": num_first,
                "num_node_to_node": num_node,
            }
        )
        total_edges += len(edges)
        logger.info(
            "%s: %d edges (%d first-level + %d node-to-node)",
            event_id,
            len(edges),
            num_first,
            num_node,
        )

    with EDGES_INDEX.open("w") as f:
        json.dump(index, f, indent=2, ensure_ascii=False)

    logger.info("=" * 60)
    logger.info("wrote %d train events → %s/", len(index), EDGES_DIR)
    logger.info("total edges: %d", total_edges)
    logger.info("index: %s", EDGES_INDEX)

    if args.no_ingest:
        logger.info("--no-ingest set; skipping ChromaDB ingestion.")
        logger.info(
            "Run `scripts/04_build_vectordb.py --build-edges` later to ingest."
        )
        return

    logger.info("=" * 60)
    logger.info("ingesting edges into ChromaDB at %s/...", config["rag"]["edge_vectordb_dir"])
    edge_count = build_edge_vectordb(config)
    logger.info("indexed %d cascade edges into the cascade_edges collection", edge_count)


if __name__ == "__main__":
    main()