src/stack3_graphrag/retriever.py

"""Graph retrieval via local search over entities.

1. Vector seeds (FAISS) — semantic base and fallback.
2. Entities of the *query* (spaCy) -> matching graph nodes.
3. Linked chunks: mentioning those entities (MENTIONS) or neighboring entities
   (RELATED_TO, 1 hop, weighted more lightly).
Score = vector similarity + entity overlap weighted by **IDF**
(rare entities count more -> neutralizes "Plant", "Role", etc.),
**normalized by the chunk's entity richness**: without it, a "hub" document
(e.g. the "June" page citing dozens of countries) accumulates a huge additive
boost and displaces the real chunks as the corpus grows. The
sqrt(num entities) normalization — analogous to BM25 length normalization — favors
*focused* chunks and lets vector similarity decide.
Falls back to pure vector if the query has no known entity.
"""

import math

import faiss
import numpy as np

from shared.embeddings import EmbeddingModel
from shared.vector_index import FaissIndexer

from .entity_extractor import extract_entities

_VEC_SEEDS = 20          # vector seeds (semantic base + fallback)
_RELATED_DISCOUNT = 0.5  # weight of entities reached via RELATED_TO (1 hop)
_GRAPH_WEIGHT = 0.3      # weight of the graph signal vs vector similarity

# Normalization of the entity boost by chunk richness: we divide by f(num entities).
# "none" = naive version ("hub" documents rich in entities grab the boost and
# displace focused chunks). sqrt (default) = analogous to BM25 length normalization.
# Single knob, swept on a *held-out* split (cf. eval/sweep_entity_norm.py) — not
# tuned on the test set. All forms are >= 1 for n >= 1 (never any amplification).
_ENTITY_NORMS = {
    "none": lambda n: 1.0,
    "log": lambda n: 1.0 + math.log(n),
    "p25": lambda n: n ** 0.25,
    "sqrt": lambda n: math.sqrt(n),
    "p75": lambda n: n ** 0.75,
    "linear": lambda n: float(n),
}
_DEFAULT_ENTITY_NORM = "p75"  # chosen by held-out sweep (eval/sweep_entity_norm.py)


class GraphRetriever:
    """Entity local search: vector seeds + entity-linked chunks,
    scored by vector similarity + IDF-weighted entity overlap."""

    def __init__(self, indexer: FaissIndexer, embedding_model: EmbeddingModel, graph,
                 entity_norm: str = _DEFAULT_ENTITY_NORM):
        self.indexer = indexer
        self.embedding_model = embedding_model
        self.graph = graph
        self._idf = self._compute_idf()
        self._chunk_n_entities = self._count_chunk_entities()
        self._norm = _ENTITY_NORMS.get(entity_norm, _ENTITY_NORMS[_DEFAULT_ENTITY_NORM])

    def search(self, query: str, k: int = 5) -> list[dict]:
        if self.indexer.size == 0:
            return []

        scored = dict(self._vector_seeds(query, max(_VEC_SEEDS, k)))  # vector base
        shared: dict[int, set] = {}

        for chunk_idx, entities, boost in self._entity_candidates(query):
            scored[chunk_idx] = scored.get(chunk_idx, 0.0) + _GRAPH_WEIGHT * boost
            shared.setdefault(chunk_idx, set()).update(entities)

        top = sorted(scored.items(), key=lambda kv: kv[1], reverse=True)[:k]
        return [self._build_result(idx, score, shared.get(idx)) for idx, score in top]

    def _compute_idf(self) -> dict[str, float]:
        """Per-entity normalized IDF: log(1 + N/df) / log(1 + N), in (0, 1]."""
        n_chunks = sum(1 for _, d in self.graph.nodes(data=True) if d.get("type") == "chunk") or 1
        denom = math.log(1 + n_chunks)
        idf: dict[str, float] = {}
        for node, data in self.graph.nodes(data=True):
            if data.get("type") != "entity":
                continue
            df = sum(1 for nb in self.graph.neighbors(node)
                     if self.graph.nodes[nb].get("type") == "chunk")
            if df:
                idf[node] = math.log(1 + n_chunks / df) / denom
        return idf

    def _count_chunk_entities(self) -> dict[int, int]:
        """Number of distinct entities per chunk (used to normalize the boost: penalizes hubs)."""
        counts: dict[int, int] = {}
        for node, data in self.graph.nodes(data=True):
            if data.get("type") != "chunk":
                continue
            counts[data["index"]] = sum(
                1 for nb in self.graph.neighbors(node)
                if self.graph.nodes[nb].get("type") == "entity"
            )
        return counts

    def _entity_candidates(self, query: str) -> list[tuple[int, set, float]]:
        """Chunks linked to the query entities (direct + RELATED_TO), with IDF boost.

        Returns (chunk_index, {entities}, boost), boost = sum of the IDFs of the
        mentioned query entities (weight 1) and their RELATED_TO neighbors (reduced weight).
        """
        query_ids = [f"entity:{name.lower()}" for name in extract_entities(query)]
        query_ids = [eid for eid in query_ids if eid in self.graph]
        if not query_ids:
            return []

        # Weight per reached entity: 1.0 for the query entities, reduced for their neighbors.
        weight: dict[str, float] = {}
        for eid in query_ids:
            weight[eid] = 1.0
            for nb in self.graph.neighbors(eid):
                if self.graph.nodes[nb].get("type") == "entity":
                    weight[nb] = max(weight.get(nb, 0.0), _RELATED_DISCOUNT)

        contributions: dict[int, list] = {}
        for eid, w in weight.items():
            idf = self._idf.get(eid, 0.0)
            if not idf:
                continue
            name = self.graph.nodes[eid].get("name", eid)
            for nb in self.graph.neighbors(eid):
                node = self.graph.nodes[nb]
                if node.get("type") != "chunk":
                    continue
                entry = contributions.setdefault(node["index"], [0.0, set()])
                entry[0] += w * idf
                entry[1].add(name)
        # Normalization by entity richness: without it, a "hub" chunk (many
        # entities) accumulates an outsized boost and displaces focused chunks.
        return [(idx, ents, boost / self._norm(self._chunk_n_entities.get(idx, 1) or 1))
                for idx, (boost, ents) in contributions.items()]

    def _vector_seeds(self, query: str, n: int) -> list[tuple[int, float]]:
        """(idx, cosine similarity) of the n nearest chunks."""
        query_emb = np.array([self.embedding_model.encode_query(query)], dtype=np.float32)
        faiss.normalize_L2(query_emb)
        n = min(n, self.indexer.size)
        scores, indices = self.indexer.index.search(query_emb, n)
        return [(int(i), float(s)) for s, i in zip(scores[0], indices[0]) if i != -1]

    def _build_result(self, idx: int, score: float, shared_entities: set | None) -> dict:
        metadata = dict(self.indexer.metadata[idx])
        if shared_entities:
            metadata["shared_entities"] = sorted(shared_entities)
        return {"text": self.indexer.chunks[idx], "metadata": metadata, "score": float(score)}