kg_retrieval.py

"""KG-aware retrieval over the in-memory fact graph (Phase 5).

Three entry points used by the query-time pipeline:

- :func:`triple_match` (Phase 5.1) — embed the query and ANN-match it
  against the fact-vector index. HippoRAG 2's core retrieval primitive
  ported wholesale: query ↔ (subject, predicate, object) similarity is
  the most discriminating signal for "which facts answer this query".
- :func:`ppr_facts` (Phase 5.2 + 5.4) — Personalized PageRank over the
  pre-built ``MemoryBackend.fact_graph`` with hub-weighted seeding.
  Surfaces facts connected to query entities via multi-hop paths that
  a direct triple match would miss.
- :func:`facts_to_chunk_ids` (Phase 5.6) — collect the source chunk ids
  cited by a set of retrieved facts. The bridge between fact-level
  retrieval and the chunk-level reader.

PPR runs on an undirected view of the fact graph (HippoRAG default —
the probabilistic spread is broader without direction constraints).
Edge weights come from each fact's confidence so high-confidence
relations dominate low-confidence ones.

This module is dormant until Phase 6 wires the new entry points into
``benchmarks/runner.py``.
"""

from __future__ import annotations

import logging
import math
from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING

from engram.core.entities import case_variants, normalize_entity_name
from engram.core.retrieval import TraversalConfig

if TYPE_CHECKING:
    from engram.backends.memory import MemoryBackend
    from engram.core.models import Fact
    from engram.core.protocol import CorpusBackend


logger = logging.getLogger(__name__)


PPR_DAMPING = 0.5
"""Teleport probability for personalized PageRank.

Matches HippoRAG 2's tuning (lower than the classic web-PageRank 0.85
because we want broader propagation in a small graph). See
[[engram-positioning]] — Engram targets MuSiQue-scale corpora where
broad spread beats sharp focus."""

PPR_MAX_ITERATIONS = 100
PPR_TOLERANCE = 1e-6
PPR_MIN_NODE_SCORE = 1e-3
"""Nodes below this PPR score are treated as irrelevant. Anything
below 1e-3 is effectively noise from the random-walk tail."""


async def triple_match(
    query: str,
    backend: CorpusBackend,
    *,
    top_k: int = 30,
    min_confidence: float = 0.7,
) -> list[Fact]:
    """Phase 5.1 — ANN-match the query against stored fact triples.

    Embeds ``query`` with the same encoder used at ingest, asks the
    second hnswlib (fact_vectors) for the top-``top_k`` nearest
    triples, and drops anything below ``min_confidence`` (matches the
    ingest-time floor — keeps quality consistent end-to-end).

    Falls back to an empty list if the backend doesn't expose
    ``neighbors_facts`` or no triples have been embedded yet (legacy
    corpora that pre-date Phase 4.9 fact embedding).
    """
    neighbors_facts = getattr(backend, "neighbors_facts", None)
    if neighbors_facts is None:
        return []
    facts = await neighbors_facts(query, k=top_k)
    return [f for f in facts if f.confidence >= min_confidence]


def ppr_facts(
    query_entities: Sequence[str],
    backend: MemoryBackend,
    *,
    damping: float = PPR_DAMPING,
    max_iterations: int = PPR_MAX_ITERATIONS,
    tolerance: float = PPR_TOLERANCE,
    min_node_score: float = PPR_MIN_NODE_SCORE,
    hub_weighted: bool = True,
    top_k: int = 100,
) -> list[Fact]:
    """Phase 5.2 + 5.4 — Personalized PageRank over the fact graph.

    Resolves ``query_entities`` to nodes in the in-memory fact graph
    (with normalization + case-variant fallback to bridge minor
    surface-form differences), builds a personalization vector that
    hub-weights the seeds, then runs networkx PageRank on the
    undirected view. Facts are ranked by the max PPR score of their
    two endpoint nodes — a fact connected to a high-PPR entity wins
    even if its other endpoint is in the tail.

    Hub weighting: with multiple seeds, high-degree entities get
    more teleport probability via ``1 + log(1 + d / max_deg)``,
    normalized. Single-seed runs skip the hub weighting and just put
    all teleport mass on the one seed.

    Returns up to ``top_k`` facts sorted by PPR score descending. Each
    fact is the canonical instance for its ``(s, p, o)`` tuple (parallel
    edges with the same triple are deduped — best score wins).
    """
    import networkx as nx

    graph = backend.fact_graph
    if graph.number_of_nodes() < 2:
        return []

    seed_nodes = _resolve_entities_to_nodes(query_entities, graph)
    if not seed_nodes:
        logger.debug("PPR: no query entities matched any graph node")
        return []

    if hub_weighted and len(seed_nodes) > 1:
        degrees = {n: graph.degree(n) for n in seed_nodes}
        max_deg = max(degrees.values()) or 1
        weights = {n: 1.0 + math.log(1.0 + degrees[n] / max_deg) for n in seed_nodes}
        total = sum(weights.values()) or 1.0
        personalization = {n: w / total for n, w in weights.items()}
    else:
        seed_weight = 1.0 / len(seed_nodes)
        personalization = dict.fromkeys(seed_nodes, seed_weight)

    undirected = graph.to_undirected(as_view=False)
    ppr_scores = nx.pagerank(
        undirected,
        alpha=damping,
        personalization=personalization,
        max_iter=max_iterations,
        tol=tolerance,
        weight="confidence",
    )

    relevant_nodes = {n: s for n, s in ppr_scores.items() if s >= min_node_score}
    if not relevant_nodes:
        return []

    # Walk only edges incident to relevant nodes; dedup parallel edges
    # with the same (s, p, o), keep highest-scoring instance.
    best_by_triple: dict[str, tuple[float, str]] = {}
    visited_keys: set[str] = set()
    for node in relevant_nodes:
        if node not in undirected:
            continue
        for u, v, key in undirected.edges(node, keys=True):
            if key in visited_keys:
                continue
            visited_keys.add(key)
            score = max(
                relevant_nodes.get(u, 0.0),
                relevant_nodes.get(v, 0.0),
            )
            triple_key = _triple_dedup_key(u, v, key, graph)
            if triple_key is None:
                continue
            existing = best_by_triple.get(triple_key)
            if existing is None or score > existing[0]:
                best_by_triple[triple_key] = (score, key)

    sorted_keys = sorted(best_by_triple.values(), key=lambda x: x[0], reverse=True)
    out: list[Fact] = []
    for _score, fact_id in sorted_keys:
        fact = backend.get_fact(fact_id)
        if fact is not None:
            out.append(fact)
        if len(out) >= top_k:
            break
    return out


def two_stage_ppr_facts(
    query_entities: Sequence[str],
    backend: MemoryBackend,
    *,
    stage1_damping: float = 0.75,
    stage2_damping: float = 0.45,
    stage1_top_n_entities: int = 20,
    max_iterations: int = PPR_MAX_ITERATIONS,
    tolerance: float = PPR_TOLERANCE,
    min_node_score: float = PPR_MIN_NODE_SCORE,
    hub_weighted: bool = True,
    top_k: int = 100,
) -> list[Fact]:
    """PropRAG-style two-stage Personalized PageRank.

    **Stage 1** (broad spread): hub-weighted PPR over the fact graph
    with the relatively-high damping ``stage1_damping`` (default 0.75
    matches PropRAG). The random walk's high persistence widens the
    explored neighborhood — surfaces the candidate entity pool that
    plausibly relates to the query.

    **Stage 2** (tight focus): the top-N entities from Stage 1 become a
    new personalization vector — the seeds are now refined by Stage 1's
    spread, not just the literal query mentions. PPR runs again with
    the lower damping ``stage2_damping`` (default 0.45) — more teleport
    back to refined seeds, less probability mass leaking to the tail.

    The two stages together separate "which entities matter" (Stage 1)
    from "rank them tightly" (Stage 2). Single-stage PPR at 0.5 (the
    HippoRAG 2 default we use in :func:`ppr_facts`) tries to do both at
    once and trades off either precision or recall.

    Returns up to ``top_k`` facts ranked by the Stage 2 score of the
    higher-scoring endpoint, deduped by ``(s, p, o)``.
    """
    import networkx as nx

    graph = backend.fact_graph
    if graph.number_of_nodes() < 2:
        return []

    seed_nodes = _resolve_entities_to_nodes(query_entities, graph)
    if not seed_nodes:
        return []

    undirected = graph.to_undirected(as_view=False)

    # Stage 1: broad PPR with hub-weighted seeds (same seeding as ppr_facts).
    if hub_weighted and len(seed_nodes) > 1:
        degrees = {n: graph.degree(n) for n in seed_nodes}
        max_deg = max(degrees.values()) or 1
        weights = {n: 1.0 + math.log(1.0 + degrees[n] / max_deg) for n in seed_nodes}
        total = sum(weights.values()) or 1.0
        stage1_personalization = {n: w / total for n, w in weights.items()}
    else:
        seed_weight = 1.0 / len(seed_nodes)
        stage1_personalization = dict.fromkeys(seed_nodes, seed_weight)

    stage1_scores = nx.pagerank(
        undirected,
        alpha=stage1_damping,
        personalization=stage1_personalization,
        max_iter=max_iterations,
        tol=tolerance,
        weight="confidence",
    )

    # Stage 2: re-seed from top-N entities of Stage 1.
    top_stage1 = sorted(stage1_scores.items(), key=lambda kv: kv[1], reverse=True)[
        :stage1_top_n_entities
    ]
    if not top_stage1:
        return []

    stage2_total = sum(s for _, s in top_stage1) or 1.0
    stage2_personalization = {n: s / stage2_total for n, s in top_stage1}

    stage2_scores = nx.pagerank(
        undirected,
        alpha=stage2_damping,
        personalization=stage2_personalization,
        max_iter=max_iterations,
        tol=tolerance,
        weight="confidence",
    )

    relevant_nodes = {n: s for n, s in stage2_scores.items() if s >= min_node_score}
    if not relevant_nodes:
        return []

    # Same fact-ranking logic as ppr_facts: walk edges incident to
    # relevant nodes, dedup parallel edges by (s,p,o), keep best score.
    best_by_triple: dict[str, tuple[float, str]] = {}
    visited_keys: set[str] = set()
    for node in relevant_nodes:
        if node not in undirected:
            continue
        for u, v, key in undirected.edges(node, keys=True):
            if key in visited_keys:
                continue
            visited_keys.add(key)
            score = max(
                relevant_nodes.get(u, 0.0),
                relevant_nodes.get(v, 0.0),
            )
            triple_key = _triple_dedup_key(u, v, key, graph)
            if triple_key is None:
                continue
            existing = best_by_triple.get(triple_key)
            if existing is None or score > existing[0]:
                best_by_triple[triple_key] = (score, key)

    sorted_items = sorted(best_by_triple.values(), key=lambda x: x[0], reverse=True)
    out: list[Fact] = []
    for _score, fact_id in sorted_items:
        fact = backend.get_fact(fact_id)
        if fact is not None:
            out.append(fact)
        if len(out) >= top_k:
            break
    return out


DEFAULT_PREDICATE_BOOST_MULTIPLIER = 1.5
"""Score multiplier applied to beam edges whose predicate matches one
the retrieval planner flagged as a priority. 1.5x keeps the boosted
edges competitive against high-confidence non-priority edges without
swamping the score distribution. Tune in Phase 2 ablation."""


def beam_search_facts(
    seed_entities: Sequence[str],
    backend: MemoryBackend,
    *,
    config: TraversalConfig | None = None,
    max_hops: int = 10,
    top_k: int = 250,
    predicate_boost: Sequence[str] | None = None,
    predicate_boost_multiplier: float = DEFAULT_PREDICATE_BOOST_MULTIPLIER,
) -> list[Fact]:
    """Phase 5.3 — confidence-decayed multi-hop beam search over the fact graph.

    Port of Vrin's ``find_facts_multi_hop`` (enhanced_query_handler.py:1997)
    swapping the Neptune Gremlin walk for a networkx in-memory walk.

    Algorithm:

    - Resolve seed entities via the same normalize → case-variant fallback
      chain as :func:`ppr_facts`.
    - **Hop 0**: every edge incident to a seed (no confidence filter — seed
      neighborhood is presumed relevant by the question's framing).
    - **Hops 1..max**:
       - Extract the frontier (new entities discovered at the prior hop,
         scored by their best incoming hop_score).
       - Cap the frontier to ``max_fan_out_per_hop`` to prevent explosion.
       - Walk every edge incident to each frontier entity.
       - Drop edges below ``min_edge_confidence``.
       - Score: ``edge_confidence * decay_factor ** hop``.
       - Apply hub-aware fan-out: facts touching a high-degree node get
         the stricter ``hub_fan_out`` cap; non-hub facts get the regular
         ``max_fan_out_per_hop`` cap.
       - Only entities whose best hop_score >= ``path_confidence_floor``
         advance to the next hop.
    - Dedupe by fact id (parallel paths to the same fact collapse to the
      highest-scoring instance).

    Returns up to ``top_k`` Facts sorted by hop_score descending.

    Complementary to :func:`ppr_facts`: PPR is high-recall probabilistic
    spread, beam is high-precision path selection. RRF-fuse them in
    :func:`benchmarks.retrieval.kg_hybrid_neighbors`.
    """
    cfg = config or TraversalConfig()
    graph = backend.fact_graph
    if graph.number_of_nodes() < 2 or not seed_entities:
        return []

    resolved_seeds = _resolve_entities_to_nodes(seed_entities, graph)
    if not resolved_seeds:
        return []

    undirected = graph.to_undirected(as_view=False)

    # Build a single lowercase-name → original-name lookup so each
    # frontier expansion is O(1) per node instead of O(N).
    nodes_by_lower: dict[str, str] = {}
    for n in undirected.nodes():
        nodes_by_lower.setdefault(n.lower(), n)

    visited: set[str] = {n.lower() for n in resolved_seeds}

    # Hub set: degree > threshold gets the stricter fan-out cap downstream.
    hubs: set[str] = {
        n.lower() for n in undirected.nodes() if undirected.degree(n) > cfg.hub_node_threshold
    }

    all_scored: list[tuple[float, int, str]] = []  # (hop_score, hop, fact_id)
    fact_ids_seen: set[str] = set()
    # Normalize the planner's priority predicates for case-insensitive
    # matching against edge data. None / empty set bypasses the boost.
    boost_set: set[str] = (
        {p.strip().lower() for p in predicate_boost if p and p.strip()}
        if predicate_boost
        else set()
    )

    def _process_hop(
        edge_tuples: list[tuple[str, str, str]],
        hop: int,
        apply_confidence_filter: bool,
    ) -> dict[str, float]:
        scored: list[tuple[float, str, str, str]] = []
        for u, v, key in edge_tuples:
            if key in fact_ids_seen:
                continue
            edge_data = undirected.get_edge_data(u, v, key=key)
            if edge_data is None:
                continue
            edge_conf = float(edge_data.get("confidence", 0.0))
            if apply_confidence_filter and edge_conf < cfg.min_edge_confidence:
                continue
            hop_score = edge_conf * (cfg.decay_factor**hop)
            if boost_set:
                predicate = str(edge_data.get("predicate", "")).strip().lower()
                if predicate in boost_set:
                    hop_score *= predicate_boost_multiplier
            scored.append((hop_score, key, u, v))

        scored.sort(key=lambda x: x[0], reverse=True)

        if hubs:
            hub_scored = [s for s in scored if s[2].lower() in hubs or s[3].lower() in hubs]
            non_hub_scored = [
                s for s in scored if s[2].lower() not in hubs and s[3].lower() not in hubs
            ]
            kept = non_hub_scored[: cfg.max_fan_out_per_hop] + hub_scored[: cfg.hub_fan_out]
        else:
            kept = scored[: cfg.max_fan_out_per_hop]

        next_frontier: dict[str, float] = {}
        for hop_score, fact_id, u, v in kept:
            fact_ids_seen.add(fact_id)
            all_scored.append((hop_score, hop, fact_id))
            for endpoint in (u, v):
                key_lower = endpoint.lower()
                if key_lower in visited:
                    continue
                if hop_score > next_frontier.get(key_lower, 0.0):
                    next_frontier[key_lower] = hop_score
        return next_frontier

    # Hop 0: every edge incident to a seed, no confidence filter.
    hop0_edges: list[tuple[str, str, str]] = []
    for seed in resolved_seeds:
        if seed not in undirected:
            continue
        for u, v, key in undirected.edges(seed, keys=True):
            hop0_edges.append((u, v, key))
    frontier = _process_hop(hop0_edges, hop=0, apply_confidence_filter=False)
    visited.update(frontier.keys())

    # Hops 1..max with cap + filter.
    max_frontier = cfg.max_fan_out_per_hop
    for hop in range(1, max_hops):
        if not frontier:
            break
        if len(frontier) > max_frontier:
            top = sorted(frontier.items(), key=lambda x: x[1], reverse=True)
            frontier = dict(top[:max_frontier])

        next_edges: list[tuple[str, str, str]] = []
        for node_lower in frontier:
            actual = nodes_by_lower.get(node_lower)
            if actual is None or actual not in undirected:
                continue
            for u, v, key in undirected.edges(actual, keys=True):
                next_edges.append((u, v, key))

        frontier = _process_hop(next_edges, hop=hop, apply_confidence_filter=True)
        # Path confidence floor: only entities whose best inbound score
        # clears the floor advance to the next hop.
        frontier = {k: s for k, s in frontier.items() if s >= cfg.path_confidence_floor}
        visited.update(frontier.keys())

    all_scored.sort(key=lambda x: x[0], reverse=True)

    out: list[Fact] = []
    for _score, _hop, fact_id in all_scored[:top_k]:
        fact = backend.get_fact(fact_id)
        if fact is not None:
            out.append(fact)
    return out


def facts_to_chunk_ids(facts: Iterable[Fact]) -> list[str]:
    """Phase 5.6 — collect chunk ids cited by ``facts`` preserving order.

    Each fact carries ``source_chunk_ids`` (the chunks it was extracted
    from). Aggregating those across the fact set gives the "KG chunks"
    candidate list — the chunks that the KG retrieval path surfaces,
    independent of dense + BM25.

    Order is fact-input order; duplicates are dropped on first sight
    so the rank a chunk would get in an RRF fusion reflects when it
    first appeared in the fact stream.
    """
    seen: set[str] = set()
    out: list[str] = []
    for fact in facts:
        for cid in fact.source_chunk_ids:
            if not cid or cid in seen:
                continue
            seen.add(cid)
            out.append(cid)
    return out


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _resolve_entities_to_nodes(
    entities: Sequence[str],
    graph,
) -> list[str]:
    """Map query entity strings to actual graph node names.

    Tries direct membership first, then normalization (article/punct
    stripping + lowercase), then case variants (CamelCase/no-space).
    Skipped facts whose subject or object didn't make it into the
    graph (rare — happens if the ingest path canonicalized them).
    """
    if not entities:
        return []
    nodes = list(graph.nodes())
    node_by_norm: dict[str, str] = {}
    for n in nodes:
        norm = normalize_entity_name(n)
        node_by_norm.setdefault(norm, n)

    resolved: list[str] = []
    seen: set[str] = set()

    def _add(node: str) -> None:
        if node not in seen:
            seen.add(node)
            resolved.append(node)

    for entity in entities:
        if not entity or not entity.strip():
            continue
        if entity in graph:
            _add(entity)
            continue
        norm = normalize_entity_name(entity)
        if norm in node_by_norm:
            _add(node_by_norm[norm])
            continue
        matched = False
        for variant in case_variants(entity):
            if variant in graph:
                _add(variant)
                matched = True
                break
        if matched:
            continue
        # Last-ditch substring match against normalized node names.
        for node_norm, node in node_by_norm.items():
            if not node_norm:
                continue
            if norm and (norm in node_norm or node_norm in norm):
                _add(node)
                break
    return resolved


def _triple_dedup_key(u: str, v: str, key: str, graph) -> str | None:
    """Build a lowercased ``s||p||o`` dedup key for an edge.

    MultiDiGraph stores facts as edges keyed by fact id. Multiple facts
    may share the same ``(s, p, o)`` triple — typical when the same
    relation is extracted from multiple chunks. PPR ranking dedupes by
    this key so the surfaced fact list is a clean unique-triple ranking.
    """
    edge = graph.get_edge_data(u, v, key=key) or graph.get_edge_data(v, u, key=key)
    if edge is None:
        return None
    predicate = str(edge.get("predicate", "")).lower().strip()
    # We don't always know which of (u, v) is the original subject vs
    # object — the undirected view collapses direction. Recover via the
    # underlying directed graph if possible.
    if graph.has_edge(u, v, key=key):
        subj, obj = u, v
    else:
        subj, obj = v, u
    return f"{subj.lower().strip()}||{predicate}||{obj.lower().strip()}"


__all__ = [
    "PPR_DAMPING",
    "PPR_MAX_ITERATIONS",
    "PPR_MIN_NODE_SCORE",
    "PPR_TOLERANCE",
    "beam_search_facts",
    "facts_to_chunk_ids",
    "ppr_facts",
    "triple_match",
    "two_stage_ppr_facts",
]