"""Iterative BFS cascade predictor (v0.2 issue #6). Replaces the v0.1 single-shot ``CascadePredictor`` with a breadth-first search over edge-fragment retrieval: * **Layer 0** — embed the new event's trigger and pull top-K *first-level* edges from the historical ``cascade_edges`` collection; ask the LLM (``predict_initial.txt``) for the layer-1 cascades anchored at the trigger. * **Layer ≥1** — for each frontier node, retrieve top-K node-to-node edges whose parent description matches the frontier node; ask the LLM (``predict_iterative.txt``) to produce the next layer rooted at the frontier; repeat until any of the three termination conditions trip: 1. cumulative ``time_offset_hours`` exceeds ``rag.time_window_hours``; 2. every per-frontier max similarity falls below ``rag.similarity_threshold``; 3. the LLM returns ``{"layer": [], "stop_reason": "saturation"}`` (or ``"out_of_domain"``). Two non-semantic safety nets — ``rag.max_layers`` and ``rag.max_total_nodes`` — bound the loop independently of the model. Public surface (``predict()`` signature, ``PredictionResult`` shape) is preserved so that ``Evaluator`` / ``GoldEvaluator`` / the Streamlit Predict tab keep working without changes; the only additive field is ``PredictionResult.trace`` (per-layer observability record). """ from __future__ import annotations import json import logging from collections import Counter from collections.abc import Iterator from datetime import date from typing import Any, Optional from src.llm.client import ( LLMClient, load_config, load_expert_knowledge, load_prompt_template, ) from src.models.schemas import CascadeChain, CascadeEdge, CascadeNode, PredictionResult from src.rag.chain_index import enrich_retrieval_info, load_chain_index from src.rag.edge_retriever import EdgeRetriever from src.rag.embedder import trigger_to_embedding_text logger = logging.getLogger(__name__) # Closed taxonomy mirrored from prompts/expert_predict.md — used only for # light-touch validation logging; the LLM is the source of truth. _VALID_DOMAINS: set[str] = { "infrastructure/power", "infrastructure/water", "infrastructure/transport", "infrastructure/communication", "health/casualties", "health/hospital_service", "health/disease_outbreak", "social/evacuation", "social/supply_shortage", "economy/business_damage", "economy/agriculture", "environment/contamination", } _VALID_STOP_REASONS: set[str] = {"saturation", "out_of_domain"} # v0.4 issue #63 #4 — domain-budget cap. # D5 audit identified three domains with high over-prediction ratio AND # zero or near-zero matches against gold. Each event's emits in these # domains are capped at gold-avg + small margin; further LLM emits in the # same domain are dropped (still recorded in trace.domain_budget_dropped). # Domains not in this map are unlimited. _DOMAIN_BUDGET: dict[str, int] = { "social/supply_shortage": 1, # gold avg 0.5/event (3/6); 13 pred / 0 matched "infrastructure/communication": 2, # gold avg 1/event (6/6); 13 pred / 0 matched "economy/business_damage": 2, # gold avg 0.83/event (5/6); 17 pred / 1 matched } def _check_domain_budget(domain: str, counts: dict[str, int]) -> bool: """Return True if this domain has remaining budget for one more emit. Domains absent from ``_DOMAIN_BUDGET`` are unlimited (always returns True). """ budget = _DOMAIN_BUDGET.get(domain) if budget is None: return True return counts.get(domain, 0) < budget def _apply_domain_budget( nodes: list, counts: dict[str, int] ) -> tuple[list, list[dict]]: """Filter ``nodes`` against per-domain budgets, mutating ``counts``. Returns (kept, dropped_records). Each dropped_record: ``{"node_id": str, "domain": str, "reason": "domain_budget_exceeded"}``. The function increments ``counts`` only for kept nodes — so calling again with the same dict carries running budget across BFS layers. """ kept = [] dropped: list[dict] = [] for n in nodes: d = n.domain if hasattr(n, "domain") else n.get("domain", "") if not _check_domain_budget(d, counts): dropped.append( { "node_id": n.id if hasattr(n, "id") else n.get("id", "?"), "domain": d, "reason": "domain_budget_exceeded", } ) continue kept.append(n) counts[d] = counts.get(d, 0) + 1 return kept, dropped class CascadePredictor: """Predict cascade risks for a new flood event using BFS over edge RAG.""" def __init__( self, llm_client: LLMClient, config: dict | None = None, *, dump_full_trace: bool = False, ): self.config = config or load_config() self.llm_client = llm_client self.retriever = EdgeRetriever(self.config) self._chain_index = load_chain_index(self.config) # When True, every layer's trace record additionally carries the full # retrieved-edge payload + the per-layer LLM call (template path, # variables, raw response). Off by default; flipped on by issue #12 # diagnostic runs through ``scripts/05_evaluate.py --dump-bfs-full``. self.dump_full_trace: bool = dump_full_trace rag_cfg = self.config.get("rag", {}) self.initial_top_k: int = rag_cfg.get("initial_top_k", 8) self.iterative_top_k: int = rag_cfg.get("iterative_top_k_per_frontier", 5) self.similarity_threshold: float = rag_cfg.get("similarity_threshold", 0.5) self.time_window_hours: float = rag_cfg.get("time_window_hours", 336) self.max_layers: int = rag_cfg.get("max_layers", 8) self.max_total_nodes: int = rag_cfg.get("max_total_nodes", 200) # Per-layer cap on newly-emitted nodes (semantic safety, complements # the global max_total_nodes ceiling). Hitting the cap truncates the # layer but does NOT terminate BFS — the loop keeps going on the # surviving subset. See technical_report/v0.2/rag/issue11_predictor_tuning.md. self.max_new_nodes_per_layer: int = rag_cfg.get( "max_new_nodes_per_layer", 10 ) # Pre-load templates / knowledge once. The judge evaluator does the # same (see src/eval/evaluator.py), so this keeps the call hot path # free of disk I/O across the BFS loop. self._initial_prompt = load_prompt_template( self.config["prompts"]["predict_initial"] ) self._iterative_prompt = load_prompt_template( self.config["prompts"]["predict_iterative"] ) self._predict_knowledge = load_expert_knowledge( self.config["knowledge"]["predict"] ) # ------------------------------------------------------------------ # Public entry point # ------------------------------------------------------------------ def predict( self, country: str, iso: str, location: str, event_date: str, severity: str, description: str, seed: int | None = None, ) -> PredictionResult: """Run BFS to completion and return the assembled :class:`PredictionResult`. Internally consumes :meth:`predict_stream`; callers that need layer-by-layer progress (e.g. Streamlit Predict tab) should call :meth:`predict_stream` directly. Args: country: Human-readable country name (e.g. ``"Italy"``). iso: ISO 3166-1 alpha-3 country code (e.g. ``"ITA"``). location: Sub-national location string used for RAG context. event_date: ISO-format date string (``"YYYY-MM-DD"``). severity: Flood severity label (e.g. ``"major"``). description: Free-text description of the flood event. seed: Optional integer forwarded to every :meth:`LLMClient.call ` invocation inside the BFS loop (both ``_llm_initial`` and ``_llm_iterative``). Honored by local backends (greedy decode when ``temperature=0``, sampling with fixed seed otherwise); cloud backends (``claude`` / ``claude-cli``) raise ``NotImplementedError`` when *seed* is not ``None``. Default ``None`` preserves backward-compatible non-deterministic behavior. """ final: dict | None = None for chunk in self.predict_stream( country=country, iso=iso, location=location, event_date=event_date, severity=severity, description=description, seed=seed, ): if chunk.get("is_final"): final = chunk if final is None: # defensive — predict_stream always yields a final chunk raise RuntimeError("predict_stream did not yield a final chunk") return final["result"] def predict_stream( self, country: str, iso: str, location: str, event_date: str, severity: str, description: str, seed: int | None = None, ) -> Iterator[dict]: """Yield one chunk per BFS layer for streaming UIs. Per-layer chunk shape (``is_final=False``):: { "layer": int, # 0 = seed; safety:max_layers may exceed max_layers "trace_record": dict, # the same record appended to result.trace "produced": list[CascadeNode], # nodes newly added to the DAG this layer "evidence_ids": list[str], # edge_ids retrieved for this layer "stop_reason": str | None, # mirrors trace_record["stop_reason"] "partial_dag": list[CascadeNode], # snapshot of DAG after this layer commits "is_final": False, } Final chunk (``is_final=True``):: {"is_final": True, "result": PredictionResult} ``predict_stream`` always emits at least one per-layer chunk and exactly one final chunk (even when BFS short-circuits at layer 0). Args: country: Human-readable country name (e.g. ``"Italy"``). iso: ISO 3166-1 alpha-3 country code (e.g. ``"ITA"``). location: Sub-national location string used for RAG context. event_date: ISO-format date string (``"YYYY-MM-DD"``). severity: Flood severity label (e.g. ``"major"``). description: Free-text description of the flood event. seed: Optional integer forwarded to every :meth:`LLMClient.call ` invocation inside the BFS loop (both ``_llm_initial`` and ``_llm_iterative``). Honored by local backends (greedy decode when ``temperature=0``, sampling with fixed seed otherwise); cloud backends (``claude`` / ``claude-cli``) raise ``NotImplementedError`` when *seed* is not ``None``. Default ``None`` preserves backward-compatible non-deterministic behavior. """ event_id = f"PRED-{event_date}-{iso}" event_inputs = { "country": country, "iso": iso, "location": location, "date": event_date, "severity": severity, "description": description, "event_id": event_id, } dag: list[CascadeNode] = [] trace: list[dict] = [] # v0.4 issue #63 #4 — running per-domain emit counts (cross-layer) # plus a flat list of dropped-by-budget node records for trace. domain_emit_counts: dict[str, int] = {} domain_budget_dropped: list[dict] = [] # Aggregate retrieval provenance across all layers; layer-0 events # are the analogue of v0.1 reference_event_ids, but we union in # layer-≥1 evidence too so the Predict tab can show what BFS leaned # on at every step. seen_event_ids: list[str] = [] event_id_to_max_sim: dict[str, float] = {} # ── Layer 0: seed ──────────────────────────────────────────── q0 = trigger_to_embedding_text( country=country, event_date=event_date, severity=severity, summary=description, location=location, ) layer1_nodes, seed_results, stop_reason, layer0_confidence, layer0_llm_call = ( self._seed_layer(q0, event_inputs, seed=seed) ) self._update_provenance(seed_results, seen_event_ids, event_id_to_max_sim) layer1_nodes, layer0_truncated = self._apply_per_layer_cap( layer1_nodes, layer0_confidence ) if layer0_truncated: logger.info( "Layer 0 truncated by per-layer cap (%d → %d, dropped %d)", layer0_truncated + len(layer1_nodes), len(layer1_nodes), layer0_truncated, ) # v0.4 issue #63 #4 — domain-budget cap (after per-layer cap) layer1_nodes, layer0_budget_drops = _apply_domain_budget( layer1_nodes, domain_emit_counts ) if layer0_budget_drops: logger.info( "Layer 0 dropped %d node(s) by domain-budget cap: %s", len(layer0_budget_drops), [(d["node_id"], d["domain"]) for d in layer0_budget_drops], ) domain_budget_dropped.extend(layer0_budget_drops) layer0_record = self._record_layer( layer_idx=0, frontier_ids=[], evidence={"trigger": seed_results}, stop_reason=stop_reason, produced_ids=[n.id for n in layer1_nodes], truncated_by_per_layer_cap=layer0_truncated, llm_call=layer0_llm_call, ) trace.append(layer0_record) dag.extend(layer1_nodes) yield self._stream_chunk( layer_idx=0, trace_record=layer0_record, produced=layer1_nodes, evidence_ids=_collect_edge_ids(seed_results), stop_reason=stop_reason, dag=dag, ) if stop_reason or not layer1_nodes: yield self._final_chunk( event_id, country, iso, event_date, severity, dag, trace, seen_event_ids, event_id_to_max_sim, ) return frontier = layer1_nodes # ── Layer ≥1 ───────────────────────────────────────────────── for layer_idx in range(1, self.max_layers + 1): if len(dag) >= self.max_total_nodes: rec = self._record_layer( layer_idx=layer_idx, frontier_ids=[n.id for n in frontier], evidence={}, stop_reason="safety:max_total_nodes", produced_ids=[], ) trace.append(rec) yield self._stream_chunk( layer_idx=layer_idx, trace_record=rec, produced=[], evidence_ids=[], stop_reason="safety:max_total_nodes", dag=dag, ) break evidence = { node.id: self.retriever.retrieve_node_to_node( query_text=node.description, top_k=self.iterative_top_k, parent_domain=node.domain, frontier_time_offset_hours=node.time_offset_hours, ) for node in frontier } flat_ev = self._flatten_evidence(evidence) stop_pre = self._should_stop(evidence) if stop_pre is not None: self._update_provenance(flat_ev, seen_event_ids, event_id_to_max_sim) rec = self._record_layer( layer_idx=layer_idx, frontier_ids=[n.id for n in frontier], evidence=evidence, stop_reason=stop_pre, produced_ids=[], ) trace.append(rec) yield self._stream_chunk( layer_idx=layer_idx, trace_record=rec, produced=[], evidence_ids=_collect_edge_ids(flat_ev), stop_reason=stop_pre, dag=dag, ) break next_layer, stop_reason, iter_confidence, iter_llm_call = ( self._llm_iterative(event_inputs, dag, frontier, evidence, seed=seed) ) next_layer = self._validate_parent_anchors(next_layer, dag, frontier) next_layer, iter_truncated = self._apply_per_layer_cap( next_layer, iter_confidence ) if iter_truncated: logger.info( "Layer %d truncated by per-layer cap (kept %d, dropped %d)", layer_idx, len(next_layer), iter_truncated, ) # v0.4 issue #63 #4 — domain-budget cap (after per-layer cap) next_layer, iter_budget_drops = _apply_domain_budget( next_layer, domain_emit_counts ) if iter_budget_drops: logger.info( "Layer %d dropped %d node(s) by domain-budget cap: %s", layer_idx, len(iter_budget_drops), [(d["node_id"], d["domain"]) for d in iter_budget_drops], ) domain_budget_dropped.extend(iter_budget_drops) self._update_provenance(flat_ev, seen_event_ids, event_id_to_max_sim) if stop_reason in _VALID_STOP_REASONS or not next_layer: effective_stop = stop_reason or "empty_layer" rec = self._record_layer( layer_idx=layer_idx, frontier_ids=[n.id for n in frontier], evidence=evidence, stop_reason=effective_stop, produced_ids=[], llm_call=iter_llm_call, ) trace.append(rec) yield self._stream_chunk( layer_idx=layer_idx, trace_record=rec, produced=[], evidence_ids=_collect_edge_ids(flat_ev), stop_reason=effective_stop, dag=dag, ) break dag.extend(next_layer) rec = self._record_layer( layer_idx=layer_idx, frontier_ids=[n.id for n in frontier], evidence=evidence, stop_reason=None, produced_ids=[n.id for n in next_layer], truncated_by_per_layer_cap=iter_truncated, llm_call=iter_llm_call, ) trace.append(rec) yield self._stream_chunk( layer_idx=layer_idx, trace_record=rec, produced=next_layer, evidence_ids=_collect_edge_ids(flat_ev), stop_reason=None, dag=dag, ) # Termination ①: time window — only nodes still inside the window # propagate to the next frontier. If none, stop. live = [ n for n in next_layer if (n.time_offset_hours or 0.0) <= self.time_window_hours ] if not live: rec = self._record_layer( layer_idx=layer_idx, frontier_ids=[n.id for n in next_layer], evidence={}, stop_reason="time_window_exhausted", produced_ids=[], ) trace.append(rec) yield self._stream_chunk( layer_idx=layer_idx, trace_record=rec, produced=[], evidence_ids=[], stop_reason="time_window_exhausted", dag=dag, ) break frontier = live else: # Loop fell through without break ⇒ hit max_layers cap. rec = self._record_layer( layer_idx=self.max_layers + 1, frontier_ids=[n.id for n in frontier], evidence={}, stop_reason="safety:max_layers", produced_ids=[], ) trace.append(rec) yield self._stream_chunk( layer_idx=self.max_layers + 1, trace_record=rec, produced=[], evidence_ids=[], stop_reason="safety:max_layers", dag=dag, ) yield self._final_chunk( event_id, country, iso, event_date, severity, dag, trace, seen_event_ids, event_id_to_max_sim, ) # ------------------------------------------------------------------ # BFS sub-steps # ------------------------------------------------------------------ def _seed_layer( self, q0: str, event_inputs: dict, *, seed: int | None = None ) -> tuple[list[CascadeNode], list[dict], Optional[str], dict[str, float], dict | None]: seed_results = self.retriever.retrieve_first_level( q0, top_k=self.initial_top_k ) nodes, stop_reason, confidence, llm_call = self._llm_initial( event_inputs, seed_results, seed=seed ) return nodes, seed_results, stop_reason, confidence, llm_call def _llm_initial( self, event_inputs: dict, seed_edges: list[dict], *, seed: int | None = None ) -> tuple[list[CascadeNode], Optional[str], dict[str, float], dict | None]: seed_text = _format_seed_edges(seed_edges) variables = {**event_inputs, "seed_edges": seed_text} raw = self.llm_client.call( self._initial_prompt, variables, self._predict_knowledge, seed=seed ) nodes, stop_reason, confidence = self._parse_initial_layer(raw) llm_call = self._build_llm_call_record("predict_initial", variables, raw) return nodes, stop_reason, confidence, llm_call def _llm_iterative( self, event_inputs: dict, dag: list[CascadeNode], frontier: list[CascadeNode], evidence: dict[str, list[dict]], *, seed: int | None = None, ) -> tuple[list[CascadeNode], Optional[str], dict[str, float], dict | None]: variables = { "event_inputs": _format_event_inputs(event_inputs), "dag_snapshot": _format_dag_snapshot(dag), "frontier_nodes": _format_frontier_nodes(frontier), "evidence_per_frontier": _format_evidence_per_frontier(frontier, evidence), } raw = self.llm_client.call( self._iterative_prompt, variables, self._predict_knowledge, seed=seed ) nodes, stop_reason, confidence = self._parse_iterative_layer(raw) llm_call = self._build_llm_call_record("predict_iterative", variables, raw) return nodes, stop_reason, confidence, llm_call def _build_llm_call_record( self, prompt_key: str, variables: dict, raw_response: str ) -> dict | None: """Snapshot one LLM call for diagnostic dumps. Returns ``None`` when ``dump_full_trace`` is off, so the BFS hot path pays only a tuple-wrap cost when diagnostics aren't requested. """ if not self.dump_full_trace: return None return { "prompt_template_path": self.config["prompts"][prompt_key], "variables": dict(variables), "raw_response": raw_response, } def _apply_per_layer_cap( self, nodes: list[CascadeNode], confidence: dict[str, float] | None = None, ) -> tuple[list[CascadeNode], int]: """Truncate ``nodes`` to ``max_new_nodes_per_layer``. Selection strategy: - If a non-empty ``confidence`` map is provided, keep the top-N by confidence (descending). Output preserves the original LLM emission order for the kept subset (stable display in the trace + UI). - Otherwise fall back to LLM emission order — most prompts produce the higher-importance cascades first, so head-truncation is a reasonable proxy. Returns ``(kept_nodes, truncated_count)``. Never terminates BFS; the BFS loop keeps going on the surviving subset. """ if len(nodes) <= self.max_new_nodes_per_layer: return nodes, 0 cap = self.max_new_nodes_per_layer if confidence: # Stable sort: confidence desc, original index asc as tiebreaker. indexed = list(enumerate(nodes)) indexed.sort( key=lambda pair: (-confidence.get(pair[1].id, 0.0), pair[0]) ) kept_ids = {pair[1].id for pair in indexed[:cap]} kept = [n for n in nodes if n.id in kept_ids] else: kept = list(nodes[:cap]) return kept, len(nodes) - len(kept) def _should_stop(self, evidence: dict[str, list[dict]]) -> Optional[str]: """Termination ② — every frontier node's max similarity below threshold.""" if not evidence: return None max_sims = [_max_similarity(es) for es in evidence.values()] if all(sim < self.similarity_threshold for sim in max_sims): return "similarity_below_threshold" return None def _validate_parent_anchors( self, candidates: list[CascadeNode], dag: list[CascadeNode], frontier: list[CascadeNode], ) -> list[CascadeNode]: """Validate parent anchors and enforce v0.2 issue #9 / B' edge rules. BFS layer ≥1 nodes MUST anchor to either a frontier node or an earlier DAG node. This function drops malformed nodes and prunes ancestors from multi-parent declarations: 1. **Empty parent_ids** — layer ≥1 nodes need at least one parent; drop silently with a warning (the iterative prompt is explicit). 2. **Self-loop** — a node listing its own id as a parent is dropped. 3. **Unknown id** — any parent id not in ``frontier ∪ dag`` is treated as a hallucination; the whole node is dropped. 4. **No-grandparent pruning** — for surviving multi-parent nodes, transitively walk ``dag.parent_ids`` upward; if a listed parent P is an ancestor of another listed parent Q, P is dropped (most- direct cause wins). The node is kept with its pruned parent set. """ known_by_id = {n.id: n for n in dag} for n in frontier: known_by_id.setdefault(n.id, n) known_ids = set(known_by_id.keys()) all_known_nodes = list(known_by_id.values()) survivors: list[CascadeNode] = [] dropped: list[str] = [] for node in candidates: if not node.parent_ids: dropped.append(f"{node.id}(no parent_ids)") continue if node.id in node.parent_ids: dropped.append(f"{node.id}(self-loop)") continue unknown = [pid for pid in node.parent_ids if pid not in known_ids] if unknown: dropped.append(f"{node.id}(unknown_parent={unknown})") continue pruned = _prune_ancestor_parents(node.parent_ids, all_known_nodes) if len(pruned) < len(node.parent_ids): logger.info( "Pruned ancestor parents for %s: %s -> %s", node.id, node.parent_ids, pruned, ) node.parent_ids = pruned survivors.append(node) if dropped: logger.warning( "Dropped %d malformed BFS nodes with bad parent anchors: %s", len(dropped), "; ".join(dropped), ) return survivors # ------------------------------------------------------------------ # Parsing # ------------------------------------------------------------------ def _parse_initial_layer( self, response: str ) -> tuple[list[CascadeNode], Optional[str], dict[str, float]]: data = _extract_json(response) if data is None: logger.error("Layer-0 LLM response unparseable; treating as empty") return [], None, {} nodes_data = data.get("cascade_events") or [] nodes, confidence = _build_nodes(nodes_data) # Layer 0: parent_ids MUST be empty per prompt — coerce defensively. for n in nodes: if n.parent_ids: logger.warning( "Layer-0 node %s had non-empty parent_ids %s — coercing to []", n.id, n.parent_ids, ) n.parent_ids = [] stop_reason = data.get("stop_reason") return nodes, stop_reason, confidence def _parse_iterative_layer( self, response: str ) -> tuple[list[CascadeNode], Optional[str], dict[str, float]]: data = _extract_json(response) if data is None: logger.error("Iterative LLM response unparseable; treating as empty") return [], None, {} layer = data.get("layer") or [] nodes, confidence = _build_nodes(layer) stop_reason = data.get("stop_reason") return nodes, stop_reason, confidence # ------------------------------------------------------------------ # Result assembly # ------------------------------------------------------------------ def _build_result( self, event_id: str, country: str, iso: str, event_date: str, severity: str, dag: list[CascadeNode], trace: list[dict], seen_event_ids: list[str], event_id_to_max_sim: dict[str, float], ) -> PredictionResult: chain = CascadeChain( event_id=event_id, trigger_summary=f"Flood prediction for {country}", trigger_country=country, trigger_iso=iso, trigger_date=date.fromisoformat(event_date), trigger_severity=severity, cascade_events=dag, extraction_date=date.today(), ) # Flat-DAG observability — final accumulated DAG with ≥3 nodes and # zero inter-node edges almost always means BFS collapsed to the # layer-0 seed only. Logged here (not per-layer) so callers see one # warning per prediction. if len(dag) >= 3 and not any(n.parent_ids for n in dag): logger.warning( "Flat DAG: %d cascade nodes emitted with zero inter-node edges " "for event %s. Every node will fall back to TRIGGER in the viewer.", len(dag), event_id, ) confidence_scores: dict[str, float] = {} for layer_record in trace: for nid in layer_record.get("produced_ids", []): # _build_nodes captured per-node confidence into the trace # via the optional 'confidence' field; surface it here so # downstream UI / eval can read it the same way as v0.1. pass # explicit no-op — see _build_nodes return path # Re-derive confidence_scores from the dag itself by looking up the # per-node 'confidence' attached during parsing — we stored them # transiently via a parallel dict per parse, but the per-node # confidence is not part of CascadeNode. To keep v0.1 wire format # we just leave confidence_scores empty unless the parsers carry it # forward; the in-memory dicts they returned are layer-local. For # now (Predict tab still uses similarity_scores), an empty dict is # acceptable. # reference_event_ids / similarity_scores: union of all retrieved # source events, sorted by descending max similarity for stable UI. ranked = sorted( seen_event_ids, key=lambda eid: event_id_to_max_sim.get(eid, 0.0), reverse=True, ) similarity_scores = [ round(event_id_to_max_sim.get(eid, 0.0), 4) for eid in ranked ] reference_info = enrich_retrieval_info( ranked, similarity_scores, self._chain_index ) return PredictionResult( predicted_chain=chain, confidence_scores=confidence_scores, reference_event_ids=ranked, similarity_scores=similarity_scores, reference_info=reference_info, trace=trace, ) # ------------------------------------------------------------------ # Trace + provenance helpers # ------------------------------------------------------------------ def _record_layer( self, layer_idx: int, frontier_ids: list[str], evidence: dict[str, list[dict]], stop_reason: Optional[str], produced_ids: list[str], truncated_by_per_layer_cap: int = 0, llm_call: dict | None = None, ) -> dict[str, Any]: flat_evidence = [r for es in evidence.values() for r in es] evidence_summary = { key: { "n_results": len(es), "max_similarity": _max_similarity(es), "source_event_ids": _unique_event_ids(es), } for key, es in evidence.items() } record: dict[str, Any] = { "layer": layer_idx, "frontier_ids": list(frontier_ids), "evidence_summary": evidence_summary, "filter_path_distribution": dict(Counter( r.get("filter_path", "?") for r in flat_evidence )), "stop_reason": stop_reason, "produced_ids": list(produced_ids), "truncated_by_per_layer_cap": truncated_by_per_layer_cap, } if self.dump_full_trace: # Diagnostic-only fields. Keep them under explicit keys so a # consumer can ``record.get("retrieved_edges")`` without breaking # when the flag is off. record["retrieved_edges"] = _serialize_evidence(evidence) record["llm_call"] = llm_call # may be None for non-LLM safety stops return record @staticmethod def _stream_chunk( layer_idx: int, trace_record: dict, produced: list[CascadeNode], evidence_ids: list[str], stop_reason: Optional[str], dag: list[CascadeNode], ) -> dict[str, Any]: return { "layer": layer_idx, "trace_record": trace_record, "produced": list(produced), "evidence_ids": list(evidence_ids), "stop_reason": stop_reason, "partial_dag": list(dag), "is_final": False, } def _final_chunk( self, event_id: str, country: str, iso: str, event_date: str, severity: str, dag: list[CascadeNode], trace: list[dict], seen_event_ids: list[str], event_id_to_max_sim: dict[str, float], ) -> dict[str, Any]: return { "is_final": True, "result": self._build_result( event_id=event_id, country=country, iso=iso, event_date=event_date, severity=severity, dag=dag, trace=trace, seen_event_ids=seen_event_ids, event_id_to_max_sim=event_id_to_max_sim, ), } @staticmethod def _flatten_evidence(evidence: dict[str, list[dict]]) -> list[dict]: out: list[dict] = [] for es in evidence.values(): out.extend(es) return out @staticmethod def _update_provenance( results: list[dict], seen_event_ids: list[str], event_id_to_max_sim: dict[str, float], ) -> None: for r in results: eid = r.get("source_event_id") or "" if not eid: continue sim = float(r.get("similarity") or 0.0) if eid not in event_id_to_max_sim: seen_event_ids.append(eid) event_id_to_max_sim[eid] = sim else: event_id_to_max_sim[eid] = max(event_id_to_max_sim[eid], sim) # ---------------------------------------------------------------------- # Module-level helpers (no state) # ---------------------------------------------------------------------- def _extract_json(raw: str) -> dict | None: text = raw.strip() if "```json" in text: text = text.split("```json", 1)[1].split("```", 1)[0] elif "```" in text: text = text.split("```", 1)[1].split("```", 1)[0] try: return json.loads(text.strip()) except (json.JSONDecodeError, ValueError): return None def _build_nodes( nodes_data: list[dict], ) -> tuple[list[CascadeNode], dict[str, float]]: nodes: list[CascadeNode] = [] confidence: dict[str, float] = {} for raw in nodes_data: nid = raw.get("id") or f"E{len(nodes) + 1}" try: node = CascadeNode( id=nid, description=raw.get("description", ""), domain=raw.get("domain", "unknown"), severity=raw.get("severity", "medium"), time_offset_hours=raw.get("time_offset_hours"), mechanism=raw.get("mechanism", ""), parent_ids=list(raw.get("parent_ids") or []), ) except Exception as exc: logger.warning("Skipping malformed cascade node %r: %s", raw, exc) continue nodes.append(node) if "confidence" in raw: try: confidence[node.id] = float(raw["confidence"]) except (TypeError, ValueError): pass if node.domain not in _VALID_DOMAINS: logger.warning("Node %s emitted with unknown domain %r", node.id, node.domain) return nodes, confidence def _prune_ancestor_parents( parents: list[str], known_nodes: list[CascadeNode], ) -> list[str]: """Drop parents that are ancestors of another listed parent. Implements the v0.2 issue #9 / B' no-grandparent rule on the predictor side: when the LLM emits a multi-parent node, transitively walk the known DAG's ``parent_ids`` upward to find each listed parent's ancestor set. Any listed parent that appears in another's ancestor set is removed so only the most-direct causes remain. Order is preserved for the survivors so the prompt-side ordering (which typically puts the frontier expander first) is retained. """ if len(parents) <= 1: return list(parents) parent_map = {n.id: list(n.parent_ids) for n in known_nodes} def ancestors_of(nid: str) -> set[str]: seen: set[str] = set() stack = list(parent_map.get(nid, [])) while stack: cur = stack.pop() if cur in seen: continue seen.add(cur) stack.extend(parent_map.get(cur, [])) return seen candidate_set = set(parents) drop: set[str] = set() for p in parents: anc = ancestors_of(p) for q in candidate_set: if q != p and q in anc: drop.add(q) return [p for p in parents if p not in drop] def _serialize_evidence(evidence: dict[str, list[dict]]) -> dict[str, list[dict]]: """JSON-serializable snapshot of retriever output (issue #12 diagnostics). The raw retriever returns dicts that already contain a ``CascadeEdge`` Pydantic object under ``edge``. We dump it via ``.model_dump(mode="json")`` so dates / enums round-trip into JSON, preserve the rest of the record, and return a plain ``{frontier_id: [...]}`` mapping suitable for inclusion in ``PredictionResult.trace``. """ out: dict[str, list[dict]] = {} for frontier_id, results in evidence.items(): items: list[dict] = [] for r in results: ser: dict = { "similarity": float(r.get("similarity") or 0.0), "source_event_id": r.get("source_event_id"), "metadata": r.get("metadata"), "document": r.get("document"), } edge = r.get("edge") if edge is not None: if hasattr(edge, "model_dump"): ser["edge"] = edge.model_dump(mode="json") else: # defensive: edge already a dict in test fixtures ser["edge"] = dict(edge) items.append(ser) out[frontier_id] = items return out def _max_similarity(results: list[dict]) -> float: if not results: return 0.0 return max(float(r.get("similarity") or 0.0) for r in results) def _collect_edge_ids(results: list[dict]) -> list[str]: """Pull edge_id strings out of retriever results for stream chunks. Falls back to ``{source_event_id}::?`` when the on-disk edge could not be reverse-loaded (test fixtures sometimes index without the edges/*.json file). Order is preserved so the UI can show edges in retrieval order. """ out: list[str] = [] for r in results: edge = r.get("edge") if edge is not None and getattr(edge, "edge_id", None): out.append(edge.edge_id) continue sid = (r.get("metadata") or {}).get("source_event_id") or r.get("source_event_id") or "unknown" out.append(f"{sid}::?") return out def _unique_event_ids(results: list[dict]) -> list[str]: seen: list[str] = [] for r in results: eid = r.get("source_event_id") if eid and eid not in seen: seen.append(eid) return seen def _format_event_inputs(event_inputs: dict) -> str: return ( f"- Country: {event_inputs.get('country', '')}\n" f"- ISO: {event_inputs.get('iso', '')}\n" f"- Location: {event_inputs.get('location', '')}\n" f"- Date: {event_inputs.get('date', '')}\n" f"- Severity: {event_inputs.get('severity', '')}\n" f"- Description: {event_inputs.get('description', '')}\n" f"- Event ID: {event_inputs.get('event_id', '')}" ) def _format_dag_snapshot(dag: list[CascadeNode]) -> str: if not dag: return "(empty)" parts = [] for n in dag: parts.append( f"- [{n.id}] {n.description}\n" f" domain={n.domain}, severity={n.severity}, " f"time_offset_hours={n.time_offset_hours}, " f"parent_ids={n.parent_ids}" ) return "\n".join(parts) def _format_frontier_nodes(frontier: list[CascadeNode]) -> str: if not frontier: return "(empty)" parts = [] for n in frontier: parts.append( f"- id={n.id}, description={n.description}, " f"domain={n.domain}, severity={n.severity}, " f"time_offset_hours={n.time_offset_hours}" ) return "\n".join(parts) def _render_edge_evidence(idx: int, result: dict) -> str: """Render one retrieved edge result for LLM prompt consumption. v0.5 issue C: use ``evidence_template`` (redacted, with ```` / ```` placeholders) when present so that raw numbers and place-names from historical events do not reach the LLM as in-context examples. The template is built by ``EdgeRetriever`` (Task 6) and stored per result. Fallback order when ``evidence_template`` is absent (legacy fixtures / edges that pre-date Task 6): 1. ``edge.parent_text`` / ``edge.child_description`` if edge is loaded. 2. ChromaDB ``metadata`` fields when edge object is unavailable. """ edge: CascadeEdge | None = result.get("edge") sim = result.get("similarity", 0.0) src = result.get("source_event_id", "") # --- Primary path: use pre-built evidence_template (Task 6) --- evidence_template = result.get("evidence_template") if evidence_template: return ( f" [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n" f" {evidence_template}" ) # --- Fallback: raw fields (legacy / test fixtures without template) --- if edge is None: # Fall back to ChromaDB metadata when the edge file could not be # reverse-loaded (e.g. test fixtures without on-disk edges). md = result.get("metadata") or {} return ( f" [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n" f" parent_text: {result.get('document', '')}\n" f" child_description: {md.get('child_description', '(unavailable)')}\n" f" child_domain: {md.get('child_domain', '')}\n" f" child_severity: {md.get('child_severity', '')}\n" f" time_offset_hours_delta: {md.get('time_offset_hours_delta', '')}" ) return ( f" [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n" f" parent_text: {edge.parent_text}\n" f" child_description: {edge.child_description}\n" f" child_domain: {edge.child_domain}\n" f" child_severity: {edge.child_severity}\n" f" time_offset_hours_delta: {edge.time_offset_hours_delta}" ) def _format_seed_edges(results: list[dict]) -> str: if not results: return "(no historical first-level edges retrieved — proceed from physical reasoning alone.)" parts = [] for i, r in enumerate(results, start=1): parts.append(_render_edge_evidence(i, r)) return "\n".join(parts) def _format_evidence_per_frontier( frontier: list[CascadeNode], evidence: dict[str, list[dict]] ) -> str: if not frontier: return "(no frontier)" blocks = [] for n in frontier: es = evidence.get(n.id, []) header = f"### Frontier node {n.id}: {n.description}" if not es: blocks.append(f"{header}\n (no historical edges retrieved)") continue lines = [header] for i, r in enumerate(es, start=1): lines.append(_render_edge_evidence(i, r)) blocks.append("\n".join(lines)) return "\n\n".join(blocks)