tests/test_iterative_predictor.py

"""Tests for the BFS CascadePredictor (v0.2 issue #6).

These tests use:

* a deterministic ``FakeEmbedder`` (8-dim one-hot keyed by ``tag:N``
  substring), reused from the EdgeRetriever test fixture style — both
  ingestion and retrieval are patched onto the same class so cosine
  similarity is exact;
* a scripted ``FakeLLMClient`` that returns a queued JSON response per
  ``call`` and which detects which of the two prompts is being asked by
  inspecting whether ``"BFS anchor rules"`` (only present in the
  iterative prompt) is in the template;
* a tmp ChromaDB seeded with a tiny edge corpus per test, plus on-disk
  ``cascade_edges/{event_id}.json`` so ``EdgeRetriever._load_edge`` can
  rehydrate full ``CascadeEdge`` objects.

No real LLM is invoked.
"""

from __future__ import annotations

import json
import logging
import re
from pathlib import Path

import pytest

from src.llm.client import LLMClient
from src.models.schemas import CascadeEdge, PredictionResult
from src.rag.ingestion import (
    add_cascade_edge,
    get_edge_chroma_client,
    get_or_create_edge_collection,
)


# ----------------------------------------------------------------------
# Test doubles
# ----------------------------------------------------------------------
_TAG_RE = re.compile(r"tag:(\d+)")


class FakeEmbedder:
    """Deterministic 8-dim one-hot embedder keyed by ``tag:N`` substrings."""

    DIM = 8

    def __init__(self, *_args, **_kwargs):
        self.calls: list[str] = []

    def embed_text(self, text: str) -> list[float]:
        self.calls.append(text)
        m = _TAG_RE.search(text)
        idx = int(m.group(1)) % self.DIM if m else 0
        v = [0.0] * self.DIM
        v[idx] = 1.0
        return v

    def embed_texts(self, texts: list[str]) -> list[list[float]]:
        return [self.embed_text(t) for t in texts]


class TwoPromptLLM(LLMClient):
    """Routes responses to the matching prompt template by content sniffing.

    The initial prompt's ``cascade_events`` shape and the iterative prompt's
    ``layer`` shape are different enough that callers want them dispatched
    independently. The simplest reliable cue is the unique
    ``"BFS anchor rules"`` heading inside ``predict_iterative.txt``.
    """

    def __init__(self, initial: list[str], iterative: list[str]):
        self.initial = list(initial)
        self.iterative = list(iterative)
        self.calls: list[dict] = []

    def call(
        self,
        prompt_template: str,
        variables: dict,
        expert_knowledge: str = "",
        *,
        seed: int | None = None,
    ) -> str:
        self.calls.append({"variables": variables, "knowledge": expert_knowledge})
        if "BFS anchor rules" in prompt_template:
            if not self.iterative:
                # Returning empty layer with explicit saturation keeps tests
                # deterministic when they over-step: BFS is supposed to stop.
                return json.dumps({"layer": [], "stop_reason": "saturation"})
            return self.iterative.pop(0)
        if not self.initial:
            raise RuntimeError("TwoPromptLLM: out of initial responses")
        return self.initial.pop(0)


# ----------------------------------------------------------------------
# Edge corpus helpers
# ----------------------------------------------------------------------
def _edge(
    edge_id: str,
    *,
    source_event_id: str,
    is_first_level: bool,
    parent_text: str,
    parent_domain: str | None = None,
    child_description: str = "child desc",
    child_domain: str = "infrastructure/power",
    child_severity: str = "high",
    delta: float = 6.0,
) -> CascadeEdge:
    return CascadeEdge(
        edge_id=edge_id,
        source_event_id=source_event_id,
        is_first_level=is_first_level,
        parent_text=parent_text,
        parent_domain=parent_domain,
        parent_severity="medium",
        child_description=child_description,
        child_domain=child_domain,
        child_severity=child_severity,
        child_mechanism="m",
        time_offset_hours_delta=delta,
        trigger_country="Italy",
        trigger_iso="ITA",
        trigger_severity="medium",
        trigger_summary="Flood in Italy on 2023-04-30",
    )


def _write_edge_files(edges_dir: Path, edges: list[CascadeEdge]) -> None:
    edges_dir.mkdir(parents=True, exist_ok=True)
    grouped: dict[str, list[CascadeEdge]] = {}
    for e in edges:
        grouped.setdefault(e.source_event_id, []).append(e)
    for event_id, group in grouped.items():
        (edges_dir / f"{event_id}.json").write_text(
            json.dumps([e.model_dump(mode="json") for e in group])
        )


def _make_config(
    tmp_path: Path,
    *,
    similarity_threshold: float = 0.5,
    time_window_hours: float = 336,
    max_layers: int = 8,
    max_total_nodes: int = 200,
    initial_top_k: int = 8,
    iterative_top_k_per_frontier: int = 5,
) -> dict:
    return {
        "embedding": {
            "backend": "sentence-transformers",
            "model": "all-MiniLM-L6-v2",
            "encoding_mode": "trigger_only",
        },
        "rag": {
            "edge_collection_name": "cascade_edges",
            "edge_vectordb_dir": str(tmp_path / "vectordb_v2"),
            "initial_top_k": initial_top_k,
            "iterative_top_k_per_frontier": iterative_top_k_per_frontier,
            "similarity_threshold": similarity_threshold,
            "time_window_hours": time_window_hours,
            "max_layers": max_layers,
            "max_total_nodes": max_total_nodes,
        },
        "paths": {
            "cascade_edges_dir": str(tmp_path / "cascade_edges"),
        },
        "data": {
            "train_years": [2023, 2024],
        },
        "prompts": {
            "predict_initial": "prompts/predict_initial.txt",
            "predict_iterative": "prompts/predict_iterative.txt",
        },
        "knowledge": {
            "predict": "knowledge/expert_predict.md",
        },
    }


@pytest.fixture
def setup_predictor(tmp_path, monkeypatch):
    """Build a CascadePredictor wired to a tmp ChromaDB + FakeEmbedder.

    Returns ``setup(edges, llm, **cfg_overrides) -> (config, predictor)``.
    """
    monkeypatch.setattr("src.rag.ingestion.Embedder", FakeEmbedder)
    monkeypatch.setattr("src.rag.edge_retriever.Embedder", FakeEmbedder)

    # Import here so the monkeypatch is in effect before the module reads
    # the original Embedder reference.
    from src.rag.predictor import CascadePredictor

    def _setup(edges: list[CascadeEdge], llm: LLMClient, **cfg_overrides):
        config = _make_config(tmp_path, **cfg_overrides)
        edges_dir = Path(config["paths"]["cascade_edges_dir"])
        _write_edge_files(edges_dir, edges)

        client = get_edge_chroma_client(config)
        collection = get_or_create_edge_collection(client, config)
        embedder = FakeEmbedder()
        for e in edges:
            add_cascade_edge(e, collection, embedder)

        predictor = CascadePredictor(llm_client=llm, config=config)
        return config, predictor

    return _setup


# ----------------------------------------------------------------------
# Canned LLM payload helpers
# ----------------------------------------------------------------------
def _initial_payload(nodes: list[dict], stop_reason: str | None = None) -> str:
    return json.dumps(
        {
            "event_id": "PRED-2025-05-01-SVN",
            "trigger_summary": "test",
            "trigger_country": "Slovenia",
            "trigger_iso": "SVN",
            "trigger_date": "2025-05-01",
            "trigger_severity": "medium",
            "cascade_events": nodes,
            "stop_reason": stop_reason,
        }
    )


def _iterative_payload(nodes: list[dict], stop_reason: str | None = None) -> str:
    return json.dumps({"layer": nodes, "stop_reason": stop_reason})


def _node(
    nid: str,
    *,
    parent_ids: list[str] | None = None,
    domain: str = "infrastructure/power",
    severity: str = "high",
    time: float = 6.0,
    description: str | None = None,
) -> dict:
    return {
        "id": nid,
        "description": description or f"node {nid}",
        "domain": domain,
        "severity": severity,
        "time_offset_hours": time,
        "mechanism": "m",
        "parent_ids": list(parent_ids or []),
    }


def _predict(predictor) -> PredictionResult:
    return predictor.predict(
        country="Slovenia",
        iso="SVN",
        location="Ljubljana",
        event_date="2025-05-01",
        severity="medium",
        description="Heavy floods in Ljubljana basin tag:0",
    )


# ----------------------------------------------------------------------
# 1. Time window STOP
# ----------------------------------------------------------------------
def test_time_window_stops_bfs_when_layer_exceeds_window(setup_predictor):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="layer1 desc tag:0",
            parent_domain="infrastructure/power",
        ),
    ]
    llm = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", time=6)])],
        iterative=[
            # Layer-2 node placed past the 336h window — should stop BFS.
            _iterative_payload([_node("E2", parent_ids=["E1"], time=400)]),
        ],
    )
    _config, predictor = setup_predictor(edges, llm)

    result = _predict(predictor)
    stop_reasons = [t.get("stop_reason") for t in result.trace]
    assert "time_window_exhausted" in stop_reasons
    # E2 is still emitted into the DAG (it just doesn't get extended).
    ids = [n.id for n in result.predicted_chain.cascade_events]
    assert "E1" in ids and "E2" in ids


# ----------------------------------------------------------------------
# 2. Similarity STOP
# ----------------------------------------------------------------------
def test_similarity_below_threshold_stops_bfs(setup_predictor):
    # Layer-1 frontier embeds tag:5; only edges with tag:5 in parent_text
    # would match. Our corpus has none → max similarity ≈ 0 < 0.5.
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="parent tag:1",  # frontier asks tag:5 — no match
            parent_domain="infrastructure/power",
        ),
    ]
    llm = TwoPromptLLM(
        initial=[
            _initial_payload(
                [_node("E1", description="frontier carries tag:5", time=6)]
            )
        ],
        iterative=[
            _iterative_payload([_node("E2", parent_ids=["E1"], time=12)]),
        ],
    )
    _config, predictor = setup_predictor(edges, llm)

    result = _predict(predictor)
    stop_reasons = [t.get("stop_reason") for t in result.trace]
    assert "similarity_below_threshold" in stop_reasons
    # _llm_iterative was never reached — still 1 LLM call total.
    assert len(llm.calls) == 1
    # Only E1 made it; E2 was never generated because BFS bailed first.
    ids = [n.id for n in result.predicted_chain.cascade_events]
    assert ids == ["E1"]


# ----------------------------------------------------------------------
# 3. Model STOP (saturation)
# ----------------------------------------------------------------------
def test_model_saturation_stops_bfs(setup_predictor):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="parent tag:0",
            parent_domain="infrastructure/power",
        ),
    ]
    llm = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", description="tag:0", time=6)])],
        iterative=[_iterative_payload([], stop_reason="saturation")],
    )
    _config, predictor = setup_predictor(edges, llm)

    result = _predict(predictor)
    stop_reasons = [t.get("stop_reason") for t in result.trace]
    assert "saturation" in stop_reasons
    ids = [n.id for n in result.predicted_chain.cascade_events]
    assert ids == ["E1"]


# ----------------------------------------------------------------------
# 4. Safety net: max_layers
# ----------------------------------------------------------------------
def test_safety_max_layers_stops_runaway_bfs(setup_predictor):
    edges = [
        _edge(
            f"EV1::p{i}->c{i}",
            source_event_id="2023-EV1-ITA",
            is_first_level=(i == 0),
            parent_text=f"parent tag:0",
            parent_domain=None if i == 0 else "infrastructure/power",
        )
        for i in range(3)
    ]
    # Each iterative call returns a single fresh node anchored to the most
    # recent frontier — never stops until the safety net kicks in.
    # Each iterative call returns one fresh node anchored to the previous
    # frontier (E1 for the first iteration, then L1, L2, ...). With
    # max_layers=3 the safety net must trip even though the LLM never stops.
    iterative_responses = [
        _iterative_payload(
            [
                _node(
                    f"L{i}",
                    parent_ids=[f"L{i-1}"] if i > 1 else ["E1"],
                    time=6 + i,
                )
            ]
        )
        for i in range(1, 12)
    ]
    llm = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", description="tag:0", time=6)])],
        iterative=iterative_responses,
    )
    _config, predictor = setup_predictor(edges, llm, max_layers=3)

    result = _predict(predictor)
    stop_reasons = [t.get("stop_reason") for t in result.trace]
    assert any(r and r.startswith("safety:max_layers") for r in stop_reasons)


# ----------------------------------------------------------------------
# 5. Safety net: max_total_nodes
# ----------------------------------------------------------------------
def test_safety_max_total_nodes_stops_bloated_bfs(setup_predictor):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="parent tag:0",
            parent_domain="infrastructure/power",
        ),
    ]
    # Two iterative layers each emit 8 nodes; with max_total_nodes=10 the
    # third iteration must be blocked by the safety net.
    layer1 = [_node(f"E{i}", description="tag:0", time=6) for i in range(8)]
    layer2 = [_node(f"F{i}", parent_ids=["E0"], time=12) for i in range(8)]
    layer3 = [_node(f"G{i}", parent_ids=["F0"], time=18) for i in range(8)]
    llm = TwoPromptLLM(
        initial=[_initial_payload(layer1)],
        iterative=[_iterative_payload(layer2), _iterative_payload(layer3)],
    )
    _config, predictor = setup_predictor(edges, llm, max_total_nodes=10)

    result = _predict(predictor)
    stop_reasons = [t.get("stop_reason") for t in result.trace]
    assert any(r and r.startswith("safety:max_total_nodes") for r in stop_reasons)
    assert len(result.predicted_chain.cascade_events) >= 10


# ----------------------------------------------------------------------
# 6. parent_ids anchor validation
# ----------------------------------------------------------------------
def test_validate_parent_anchors_drops_orphans_and_keeps_valid(setup_predictor, caplog):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="parent tag:0",
            parent_domain="infrastructure/power",
        ),
    ]
    llm = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", description="tag:0", time=6)])],
        iterative=[
            _iterative_payload(
                [
                    _node("E2", parent_ids=["E1"], time=12),
                    _node("E99", parent_ids=["X999"], time=12),  # orphan
                ]
            ),
            # Second iteration to exercise the loop further; saturate so we
            # don't depend on a third response.
            _iterative_payload([], stop_reason="saturation"),
        ],
    )
    with caplog.at_level(logging.WARNING):
        _config, predictor = setup_predictor(edges, llm)
        result = _predict(predictor)

    ids = [n.id for n in result.predicted_chain.cascade_events]
    assert "E2" in ids
    assert "E99" not in ids
    assert any("unknown_parent" in rec.message for rec in caplog.records)


# ----------------------------------------------------------------------
# 7. Flat-DAG warning on accumulated DAG
# ----------------------------------------------------------------------
def test_flat_dag_warning_when_no_inter_node_edges(setup_predictor, caplog):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
    ]
    # Layer 0 emits 3 sibling nodes, all parent_ids=[] (correct for layer 0).
    # Iterative call returns immediate saturation so DAG stays "flat" with
    # zero inter-node edges — that is the v0.1 star-DAG failure mode.
    llm = TwoPromptLLM(
        initial=[
            _initial_payload(
                [_node("E1"), _node("E2"), _node("E3")]
            )
        ],
        iterative=[_iterative_payload([], stop_reason="saturation")],
    )
    with caplog.at_level(logging.WARNING):
        _config, predictor = setup_predictor(edges, llm)
        result = _predict(predictor)

    assert len(result.predicted_chain.cascade_events) == 3
    assert all(not n.parent_ids for n in result.predicted_chain.cascade_events)
    assert any("Flat DAG" in rec.message for rec in caplog.records)


# ----------------------------------------------------------------------
# 8. predict_stream chunk shape contract
# ----------------------------------------------------------------------
def test_predict_stream_emits_per_layer_then_final_chunk(setup_predictor):
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="node E1 tag:0",
            parent_domain="infrastructure/power",
        ),
    ]
    llm = TwoPromptLLM(
        initial=[
            _initial_payload(
                [_node("E1", description="node E1 tag:0", time=6)]
            )
        ],
        iterative=[_iterative_payload([], stop_reason="saturation")],
    )
    _config, predictor = setup_predictor(edges, llm)

    chunks = list(
        predictor.predict_stream(
            country="Slovenia",
            iso="SVN",
            location="Ljubljana",
            event_date="2025-05-01",
            severity="medium",
            description="Heavy floods in Ljubljana basin tag:0",
        )
    )

    # At least one per-layer chunk + exactly one final chunk.
    assert len(chunks) >= 2
    final = chunks[-1]
    assert final["is_final"] is True
    assert isinstance(final["result"], PredictionResult)

    # Per-layer chunks share the contract documented on predict_stream.
    expected_keys = {
        "layer", "trace_record", "produced", "evidence_ids",
        "stop_reason", "partial_dag", "is_final",
    }
    for chunk in chunks[:-1]:
        assert chunk["is_final"] is False
        assert expected_keys.issubset(chunk.keys())
        assert isinstance(chunk["produced"], list)
        assert isinstance(chunk["evidence_ids"], list)
        assert isinstance(chunk["partial_dag"], list)
        # trace_record is the same dict appended to result.trace
        assert chunk["trace_record"]["layer"] == chunk["layer"]

    # The accumulated DAG in the final chunk's result mirrors the last
    # per-layer chunk's partial_dag — predict() returns this same object.
    assert [n.id for n in final["result"].predicted_chain.cascade_events] == [
        n.id for n in chunks[-2]["partial_dag"]
    ]
    assert final["result"].predicted_chain.cascade_events[0].id == "E1"


# ----------------------------------------------------------------------
# 9. dump_full_trace flag — issue #12 diagnostic capture
# ----------------------------------------------------------------------
def test_dump_full_trace_flag_attaches_edges_and_llm_call(setup_predictor):
    """Flag on → trace records carry retrieved_edges + llm_call.
    Flag off (default) → neither key is present."""
    edges = [
        _edge(
            "EV1::TRIGGER->E1",
            source_event_id="2023-EV1-ITA",
            is_first_level=True,
            parent_text="trigger tag:0",
            child_description="downstream cascade child",
        ),
        _edge(
            "EV1::E1->E2",
            source_event_id="2023-EV1-ITA",
            is_first_level=False,
            parent_text="node E1 tag:0",
            parent_domain="infrastructure/power",
        ),
    ]

    # ---- Flag OFF (baseline) ----
    llm_off = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", description="node E1 tag:0", time=6)])],
        iterative=[_iterative_payload([], stop_reason="saturation")],
    )
    _cfg_off, predictor_off = setup_predictor(edges, llm_off)
    assert predictor_off.dump_full_trace is False
    result_off = _predict(predictor_off)
    for record in result_off.trace:
        assert "retrieved_edges" not in record
        assert "llm_call" not in record

    # ---- Flag ON (diagnostic) ----
    llm_on = TwoPromptLLM(
        initial=[_initial_payload([_node("E1", description="node E1 tag:0", time=6)])],
        iterative=[_iterative_payload([], stop_reason="saturation")],
    )
    _cfg_on, predictor_on = setup_predictor(edges, llm_on)
    predictor_on.dump_full_trace = True
    result_on = _predict(predictor_on)

    layer0 = next(t for t in result_on.trace if t["layer"] == 0)
    assert "retrieved_edges" in layer0
    assert "llm_call" in layer0
    # Layer-0 retrieval is keyed by 'trigger' (the seed query bucket).
    assert "trigger" in layer0["retrieved_edges"]
    seed_payload = layer0["retrieved_edges"]["trigger"]
    assert seed_payload, "expected at least one retrieved seed edge"
    first_edge = seed_payload[0]
    # Serialized edge round-trips through Pydantic; identity-bearing fields
    # must survive into the dump.
    assert first_edge["edge"]["edge_id"] == "EV1::TRIGGER->E1"
    assert "similarity" in first_edge

    # Layer-0 ran the predict_initial LLM call → llm_call is populated.
    assert layer0["llm_call"] is not None
    assert layer0["llm_call"]["prompt_template_path"] == "prompts/predict_initial.txt"
    assert "raw_response" in layer0["llm_call"]
    assert "variables" in layer0["llm_call"]
    # Variables must include the seed_edges injected by _format_seed_edges.
    assert "seed_edges" in layer0["llm_call"]["variables"]