retrieval.py

"""retrieval.py - R3: vector search (+ optional NVIDIA rerank) + graph traversal,
merged into one cited evidence list with [Source ID: N] handles.

Two retrieval modes combine:
  - vector : semantic/lexical similarity over the chunk store (the 'find passages')
  - graph  : connected_facts() over boroughs named in the question (the 'connect
             facts across sources for the same borough' = the differentiator)

rerank() is a seam: on the Spark it calls llama-3.2-nv-rerankqa-1b-v2; offline it
is identity (vector order). Keeps today runnable, Spark a config flip.
"""
from __future__ import annotations

import config
import embeddings
import store as store_mod
from graphstore import GraphStore

_STORE = None
_GRAPH = None


def _ensure_index():
    """Build the index for the CURRENT embed mode if it is missing. Because the store
    path is dimension-specific, switching EMBED_MODE just builds (once) its own index
    instead of crashing on a dim mismatch."""
    import os
    if not (os.path.exists(config.STORE_PATH + ".json") and os.path.exists(config.GRAPH_PATH)):
        print(f"[retrieval] no index for EMBED_MODE={config.EMBED_MODE} ({config.EMBED_DIM}D) - building it")
        import ingest
        ingest.run(None)


def _load():
    global _STORE, _GRAPH
    if _STORE is None:
        _ensure_index()
        _STORE = (store_mod.LocalStore.load(config.STORE_PATH)
                  if config.STORE_MODE == "local" else store_mod.get_store(config.EMBED_DIM))
        _GRAPH = GraphStore.load(config.GRAPH_PATH)
    return _STORE, _GRAPH


def detect_boroughs(question: str, graph: GraphStore) -> list[str]:
    ql = question.lower()
    return [b["label"] for b in graph.boroughs() if b["label"].lower() in ql]


def rerank(question: str, hits: list[dict]) -> list[dict]:
    # Spark: NVIDIA cross-encoder rerank. Offline: identity (already similarity-sorted).
    return hits


# ranking intent -> answer it from the graph (exact), not from lossy vector recall
_RANK_DESC = ("highest", "most", "top", "richest", "wealthiest", "greatest", "largest", "longest", "best", "biggest")
_RANK_ASC = ("lowest", "least", "bottom", "poorest", "smallest", "shortest", "worst")
# (indicator substring, trigger words, invert). invert=True where a HIGHER value is "worse":
# a higher house-price-to-earnings ratio means LESS affordable, so "least affordable" = highest.
_IND_HINTS = [
    ("income", ("income", "taxpayer", "earn", "wage", "salary", "rich", "poor", "money", "wealth", "pay"), False),
    ("life_expectancy", ("life expectancy", "expectancy", "longev", "lifespan", "live long"), False),
    ("house_price_to_earnings", ("afford", "housing", "house price", "rent", "priced out", "buy a home"), True),
]


def _ranking(question: str):
    ql = question.lower()
    desc = any(w in ql for w in _RANK_DESC)
    asc = any(w in ql for w in _RANK_ASC)
    if not (desc or asc):
        return None
    for ind, words, invert in _IND_HINTS:
        if any(w in ql for w in words):
            descending = desc or not asc
            return ind, (not descending if invert else descending)
    return None  # ranking words but no indicator we can sort exactly -> vector handles it


def retrieve(question: str, k: int = 6) -> dict:
    store, graph = _load()
    evidence, seen = [], set()

    # 0. ranking questions: exact, sorted answer straight from the graph
    rk = _ranking(question)
    if rk:
        hint, desc = rk
        for r in graph.rank_by_indicator(hint, descending=desc, n=8):
            if r["content"] in seen:
                continue
            seen.add(r["content"])
            evidence.append({"content": r["content"], "via": "rank",
                             "borough": r["borough"], "indicator": r["indicator"]})

    # 1. vector search
    qvec = embeddings.embed_one(question, "query")
    for h in rerank(question, store.search(qvec, k)):
        if h["content"] in seen:
            continue
        seen.add(h["content"])
        evidence.append({"content": h["content"], "via": "vector",
                         "score": round(h["score"], 3), "meta": h["metadata"]})

    # 2. graph join for any borough named in the question (the cross-source move)
    boroughs = detect_boroughs(question, graph)
    if boroughs:
        for r in graph.connected_facts(boroughs=boroughs, limit=k):
            if r["content"] in seen:
                continue
            seen.add(r["content"])
            evidence.append({"content": r["content"], "via": "graph",
                             "borough": r["borough"], "indicator": r["indicator"],
                             "meta": {"source": r["source"], "period": r["period"]}})

    for i, e in enumerate(evidence, 1):
        e["sid"] = i
    return {"question": question, "boroughs": boroughs, "ranking": bool(rk), "evidence": evidence}