studyhub/core/retrieve.py

from __future__ import annotations

import re
from dataclasses import dataclass

import numpy as np

from ..ingest.embedder import Embedder
from ..ingest.models import Chunk
from .loader import RetrievalIndex


@dataclass
class ScoredChunk:
    chunk: Chunk
    score: float


def _norm(s: str) -> str:
    return re.sub(r"\s+", " ", s.lower()).strip()


def retrieve(query: str, index: RetrievalIndex, embedder: Embedder,
             k: int = 6, candidate_k: int = 30, reranker=None) -> list[ScoredChunk]:
    n = len(index.chunks)
    if n == 0:
        return []
    qv = embedder.encode([query]).astype("float32")  # (1, dim), L2-normalized
    rows: set[int] = set()
    _, ids = index.faiss_index.search(qv, min(candidate_k, n))
    for r in ids[0]:
        if r >= 0:
            rows.add(int(r))
    # lexical concept-expansion (centroid vectors are deferred to Phase 3)
    nq = _norm(query)
    for concept, chunk_ids in index.concept_index.items():
        cn = _norm(concept)
        if cn and cn in nq:
            for cid in chunk_ids:
                r = index.row_by_id.get(cid)
                if r is not None:
                    rows.add(r)
    q = qv[0]
    scored = [
        ScoredChunk(chunk=index.chunks[r],
                    score=float(np.dot(q, np.asarray(index.faiss_index.reconstruct(int(r)), dtype="float32"))))
        for r in rows
    ]
    scored.sort(key=lambda s: s.score, reverse=True)
    # Optional rerank: a cross-encoder re-scores the top cosine candidates for sharper
    # grounding. Cap the pool so CPU latency stays bounded (near-instant on GPU).
    if reranker is not None and scored:
        pool = scored[: max(k, 16)]
        for sc, rscore in zip(pool, reranker.scores(query, [s.chunk.text for s in pool])):
            sc.score = float(rscore)
        pool.sort(key=lambda s: s.score, reverse=True)
        return pool[:k]
    return scored[:k]