backend/tests/test_retrieve.py

import numpy as np
from astroparse_api.retrieve import Corpus, retrieve_for_paragraphs
from astroparse_api.schemas import Paragraph


class FakeIndex:
    """Mimics the faiss index API Corpus uses: .search(queries_2d, k) -> (D, I)."""

    def __init__(self, vectors):
        self.vectors = np.asarray(vectors, dtype=np.float32)

    def search(self, queries, k):
        d = np.linalg.norm(self.vectors[None, :, :] - queries[:, None, :], axis=2)
        d = d**2 + 1e-6  # squared L2, like faiss METRIC_L2
        order = np.argsort(d, axis=1)[:, :k]
        return np.take_along_axis(d, order, axis=1), order


class FakeDataset:
    """Mimics the HF dataset column-slicing API Corpus uses."""

    def __init__(self, rows):
        self.rows = rows

    def __getitem__(self, idx):
        if isinstance(idx, list):
            return {k: [self.rows[i][k] for i in idx] for k in self.rows[0]}
        return self.rows[idx]


def _rows():
    import datetime

    return [
        {
            "embed": [1.0, 0.0],
            "title": "Galaxy quenching",
            "abstract": "a1",
            "authors": ["Mowla, L.", "Iyer, K."],
            "date": datetime.date(2022, 6, 1),
            "keywords": ["galaxy", "quenching"],
            "bibcode": "2022ApJ...933L...9M",
            # Real corpus: ads_id is a numeric ADS internal id; arxiv_id is the arXiv identifier
            "ads_id": "12345678",
            "arxiv_id": "2206.11862",
        },
        {
            "embed": [0.0, 1.0],
            "title": "Cosmic dawn",
            "abstract": "a2",
            "authors": ["Smith, A."],
            "date": datetime.date(2020, 1, 1),
            "keywords": ["reionization"],
            "bibcode": "2020ApJ...000....1S",
            "ads_id": "99991111",
            "arxiv_id": "2001.00001",
        },
        # no arXiv id: arxiv_id is in old pre-2007 format — common in the corpus for older papers
        {
            "embed": [0.9, 0.1],
            "title": "Galaxy SFHs",
            "abstract": "a3",
            "authors": ["Iyer, K."],
            "date": datetime.date(2019, 3, 1),
            "keywords": ["galaxy", "star formation"],
            "bibcode": "2019ApJ...000....2I",
            "ads_id": "77778888",
            "arxiv_id": "astro-ph9912345_arXiv.txt",  # old-style id → arxiv="" in LitPaper
        },
    ]


def _corpus():
    rows = _rows()
    return Corpus(FakeDataset(rows), FakeIndex([r["embed"] for r in rows]))


def test_retrieve_ranks_by_similarity_times_keywords():
    corpus = _corpus()
    paras = [
        Paragraph(
            id="p1",
            section="Intro",
            text="We study galaxy quenching and star formation." + "x" * 200,
        )
    ]
    embs = [np.array([1.0, 0.0], dtype=np.float32)]
    lit_papers, lit_by_para = retrieve_for_paragraphs(
        corpus, paras, embs, top_k=2, pool=3
    )
    ids = lit_by_para["p1"]
    assert len(ids) == 2
    top = lit_papers[ids[0]]
    assert top.bibcode == "2022ApJ...933L...9M"  # closest embed AND keyword match
    assert (
        top.scores.rank == 1
        and 0 < top.scores.embed <= 1.0
        and 0 < top.scores.keywords <= 1.0
    )
    assert top.short == "Mowla22" and top.year == 2022 and top.arxiv == "2206.11862"
    second = lit_papers[ids[1]]
    assert second.bibcode == "2019ApJ...000....2I" and second.id == second.bibcode
    assert second.arxiv == ""  # no arXiv id — UI hides the arXiv link, ADS link still works


def test_search_handles_zero_distance():
    class ZeroIndex:
        def search(self, queries, k):
            return np.array([[0.0]]), np.array([[0]])

    corpus = Corpus(FakeDataset(_rows()), ZeroIndex())
    _, sims = corpus.search(np.zeros(2, dtype=np.float32), k=1)
    assert np.isfinite(sims).all()


# ---------------------------------------------------------------------------
# Feature: top-10 default retrieval pool
# ---------------------------------------------------------------------------

def _rows_12():
    """12-row dataset for testing default top_k=10."""
    import datetime
    rows = []
    for i in range(12):
        rows.append({
            "embed": [float(i % 3), float(i % 2)],
            "title": f"Paper {i}",
            "abstract": f"Abstract {i}",
            "authors": [f"Author{i}, A."],
            "date": datetime.date(2020 + i % 5, 1, 1),
            "keywords": ["galaxy", f"kw{i}"],
            "bibcode": f"2020ApJ...{i:03d}....{i}X",
            "ads_id": str(10000 + i),
            "arxiv_id": f"20{i:02d}.{i:05d}",
        })
    return rows


def _corpus_12():
    rows = _rows_12()
    return Corpus(FakeDataset(rows), FakeIndex([r["embed"] for r in rows]))


def test_retrieve_default_top_k_is_10():
    """Default top_k=10: a 12-row corpus should yield 10 results for a paragraph."""
    corpus = _corpus_12()
    paras = [
        Paragraph(
            id="p1",
            section="Intro",
            text="We study galaxy formation and evolution in detail." + "x" * 200,
        )
    ]
    embs = [np.array([1.0, 0.0], dtype=np.float32)]
    # Call with default top_k (must be 10 now)
    lit_papers, lit_by_para = retrieve_for_paragraphs(corpus, paras, embs)
    ids = lit_by_para["p1"]
    assert len(ids) == 10, f"Expected 10 results with default top_k, got {len(ids)}"
    # Ranks should be 1..10 in ascending order of assignment
    ranks = [lit_papers[i].scores.rank for i in ids]
    assert ranks == list(range(1, 11)), f"Expected ranks 1..10, got {ranks}"