scripts/_dump_bfs_traces.py

"""One-shot helper for issue #8: dump BFS trace per test event.

Runs ``CascadePredictor.predict_stream`` for every test event and writes
``data/evaluation/v02_bfs_traces/{event_id}.json`` containing the
``trace`` list (per-layer records with ``stop_reason``, frontier ids,
evidence ids, produced ids). Used to populate §4 of
``technical_report/v0.2/evaluation/v02_alignment.md``; not part of the
production pipeline.
"""

from __future__ import annotations

import json
import logging
from datetime import date
from pathlib import Path

from src.data.cascade_extractor import _infer_severity
from src.eval.evaluator import _build_description
from src.llm import create_llm_client
from src.llm.client import load_config
from src.models.schemas import FloodEvent
from src.rag.predictor import CascadePredictor

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


def main() -> None:
    config = load_config()
    out_dir = Path("data/evaluation/v02_bfs_traces")
    out_dir.mkdir(parents=True, exist_ok=True)

    test_events = [
        FloodEvent(**e)
        for e in json.loads(Path(config["paths"]["test_events"]).read_text())
    ]

    llm = create_llm_client(config)
    predictor = CascadePredictor(llm, config)

    for i, event in enumerate(test_events, start=1):
        logger.info("[%d/%d] %s (%s)", i, len(test_events), event.event_id, event.country)
        chunks = list(
            predictor.predict_stream(
                country=event.country,
                iso=event.iso,
                location=event.location or event.country,
                event_date=str(event.start_date),
                severity=_infer_severity(event),
                description=_build_description(event),
            )
        )
        final = chunks[-1]
        result = final["result"]

        trace = result.trace
        # Layer 0 always issues an LLM call; layer ≥1 calls LLM unless a
        # pre-LLM termination fires (similarity_below_threshold or
        # safety:max_total_nodes / safety:max_layers).
        llm_calls = 0
        for rec in trace:
            sr = rec.get("stop_reason")
            if rec["layer"] == 0:
                llm_calls += 1
            elif sr in (
                "similarity_below_threshold",
                "safety:max_total_nodes",
                "safety:max_layers",
            ):
                continue
            else:
                llm_calls += 1

        terminal_stop = trace[-1].get("stop_reason") if trace else None
        layer_count = max((r["layer"] for r in trace), default=-1) + 1
        node_count = len(result.predicted_chain.cascade_events)

        out_path = out_dir / f"{event.event_id}.json"
        out_path.write_text(
            json.dumps(
                {
                    "event_id": event.event_id,
                    "country": event.country,
                    "trace": trace,
                    "summary": {
                        "layer_count": layer_count,
                        "node_count": node_count,
                        "llm_calls": llm_calls,
                        "terminal_stop_reason": terminal_stop,
                    },
                    "dumped_at": str(date.today()),
                },
                indent=2,
                ensure_ascii=False,
            )
        )
        logger.info(
            "  → %s  layers=%d  nodes=%d  llm_calls=%d  stop=%s",
            out_path.name,
            layer_count,
            node_count,
            llm_calls,
            terminal_stop,
        )


if __name__ == "__main__":
    main()