Add arxiv-complete retrieval layer (semantic + keyword)

retrieval/query.py

@@ -0,0 +1,207 @@
+"""
+Query interface for the retrieval layer.
+
+``query()`` is the single entry point the CLI (and any caller) uses.  It:
+
+  1. embeds the natural-language query,
+  2. runs an ANN search against the FAISS index,
+  3. attaches title / year / category / URL / abstract snippet to each hit,
+  4. applies the requested filters (year range, primary category, survey-only),
+  5. returns a ranked list of ``Result`` objects.
+
+If no index exists on disk yet (or building it failed), ``query()`` falls back
+to a **pure-keyword** search that streams the metadata and matches on title/
+abstract with a TF-style score.  This guarantees the CLI always returns
+something, even before an embedding index has been built.
+"""
+
+from __future__ import annotations
+
+import datetime
+import re
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from data import DataStore, SURVEY_KEYWORDS
+
+# Cap the working set for the keyword fallback so it stays fast on the 3.15M-row
+# metadata file (term-overlap scoring doesn't need the full corpus).
+_KEYWORD_WORKING_SET = 200_000
+
+
+@dataclass
+class Result:
+    paper_id: str
+    title: str
+    year: int
+    category: str
+    url: str
+    snippet: str
+    score: float
+    method: str  # "semantic" | "keyword"
+
+
+@dataclass
+class QueryOptions:
+    top_k: int = 10
+    year_min: int | None = None
+    year_max: int | None = None
+    primary_category: str | None = None
+    survey_only: bool = False
+    embedder = None  # set by caller; optional
+
+
+def _snippet(abstract: str, maxlen: int = 200) -> str:
+    text = re.sub(r"\s+", " ", abstract or "").strip()
+    if len(text) <= maxlen:
+        return text
+    return text[:maxlen].rsplit(" ", 1)[0] + " …"
+
+
+def _normalize(text: str) -> str:
+    return re.sub(r"[^a-z0-9]+", " ", (text or "").lower()).strip()
+
+
+def _keyword_score(text: str, terms: list[str]) -> float:
+    """TF-style overlap score between a doc and the query terms."""
+    words = set(text.split())
+    score = sum(1 for t in terms if t in words)
+    # small bonus for title matches
+    return score
+
+
+def _keyword_search(
+    query: str,
+    store: DataStore,
+    opts: QueryOptions,
+) -> list[Result]:
+    """
+    Pure-keyword fallback: stream matching rows, score by term overlap, filter,
+    and return the top-k.  Used when no ANN index is available.
+    """
+    terms = [t for t in _normalize(query).split() if len(t) > 1]
+    if not terms:
+        return []
+
+    # Pull a working set of rows (push the cheap filters into the scan).
+    scan = store.scan_filtered(
+        year_min=opts.year_min,
+        year_max=opts.year_max,
+        primary_category=opts.primary_category,
+        survey_only=opts.survey_only,
+    )
+    rows = (
+        scan.select(
+            ["paper_id", "title", "primary_category", "first_version_date", "arxiv_abs_url", "abstract"]
+        )
+        .limit(_KEYWORD_WORKING_SET)
+        .collect()
+        .to_dicts()
+    )
+    if not rows:
+        return []
+
+    scored: list[Result] = []
+    for r in rows:
+        rec = dict(r)
+        title = str(rec["title"] or "")
+        abstract = str(rec["abstract"] or "")
+        body = f"{title} {abstract}"
+        score = _keyword_score(body.lower(), terms)
+        if score <= 0:
+            continue
+        fvd = rec["first_version_date"]
+        year = fvd.year if isinstance(fvd, __import__("datetime").datetime) else (int(fvd) if fvd is not None else 0)
+        scored.append(
+            Result(
+                paper_id=str(rec["paper_id"]),
+                title=title,
+                year=year,
+                category=str(rec["primary_category"] or ""),
+                url=str(rec["arxiv_abs_url"] or ""),
+                snippet=_snippet(abstract),
+                score=score,
+                method="keyword",
+            )
+        )
+    scored.sort(key=lambda x: x.score, reverse=True)
+    return scored[: opts.top_k]
+
+
+def query(
+    query_text: str,
+    store: DataStore,
+    ann=None,
+    embedder=None,
+    opts: QueryOptions | None = None,
+) -> list[Result]:
+    """
+    Run a search.  Tries the ANN index first; falls back to keyword search.
+
+    Returns a list of ranked ``Result`` objects.
+    """
+    opts = opts or QueryOptions()
+    query_text = query_text or ""
+
+    # --- semantic path ---
+    if ann is not None and embedder is not None:
+        try:
+            return _semantic_query(query_text, store, ann, embedder, opts)
+        except Exception as exc:  # pragma: no cover - defensive
+            print(f"[query] ANN search failed ({exc}); falling back to keyword search.")
+
+    # --- keyword fallback ---
+    return _keyword_search(query_text, store, opts)
+
+
+def _semantic_query(
+    query_text: str,
+    store: DataStore,
+    ann,
+    embedder,
+    opts: QueryOptions,
+) -> list[Result]:
+    import faiss
+
+    qvec = np.ascontiguousarray(embedder.embed_one(query_text), dtype=np.float32)
+    k = min(opts.top_k * 8, len(ann), ann.index.ntotal)
+    k = max(k, 1)
+    dists, idx = ann.index.search(qvec.reshape(1, -1), k)
+    dists = dists[0]
+    idx = idx[0]
+
+    # Gather candidate hits and apply filters.
+    candidates: list[Result] = []
+    for d, i in zip(dists.tolist(), idx.tolist()):
+        if i < 0 or i >= len(ann._ids):
+            continue
+        pid = ann._ids[i]
+        year = ann._years[i]
+        cat = ann._cats[i]
+        if opts.year_min is not None and year < opts.year_min:
+            continue
+        if opts.year_max is not None and year > opts.year_max:
+            continue
+        if opts.primary_category and cat != opts.primary_category:
+            continue
+        if opts.survey_only:
+            if not any(kw in ann._abstracts[i].lower() for kw in SURVEY_KEYWORDS):
+                continue
+        # FAISS L2 distance -> similarity in [0,1]; closer => higher.
+        sim = 1.0 / (1.0 + float(d))
+        candidates.append(
+            Result(
+                paper_id=pid,
+                title=ann._titles[i],
+                year=year,
+                category=cat,
+                url=ann._urls[i],
+                snippet=_snippet(ann._abstracts[i]),
+                score=sim,
+                method="semantic",
+            )
+        )
+
+    candidates.sort(key=lambda x: x.score, reverse=True)
+    return candidates[: opts.top_k]