Literature Retrieval

A keyword search over abstracts returns the papers that share vocabulary with the query. That covers most of the surface, and it misses two cases. In the first, the foundational paper uses different vocabulary: a query for "sparse attention" misses a paper titled "block selection in transformer routing." In the second, the relevant paper is a follow-up that cites a known anchor, and finding the anchor then walking forward is more efficient than brute-forcing the abstract pool.

So this lesson builds two passes. BM25 over titles and abstracts catches the lexical hits. A citation graph traversal takes those hits as a seed set and expands it forward and backward by one or two hops. The union is deduplicated by paper id and ranked by a small combined score.

Paper
  id          : str           (stable identifier, "p001" for the mock corpus)
  title       : str
  abstract    : str
  year        : int
  authors     : list[str]
  references  : list[str]     (paper ids this paper cites)
  citations   : list[str]     (paper ids that cite this paper)
  source      : str           (which mock api supplied it, "arxiv" or "s2"; every result from RetrievalClient.search carries "merged")

The references and citations fields together form the directed citation graph. The two mock APIs return overlapping but not identical fields, and id is the key that identifies the same paper across them.

The retrieval client owns both passes and the merge. It builds the BM25 index and the citation graph directly from the corpus the two mock clients wrap. You hand it a query and get back a ranked list. Each entry carries per-paper score fields (bm25_score, graph_distance, recency_score, final_score) that explain where it landed.

BM25 is a lexical ranking function: it scores a document against a query by summing, over the query terms, how rare each term is across the corpus and how often it appears in that document, adjusted for document length. The implementation here is standard Okapi BM25 with default parameters k1=1.5 and b=0.75.

The index is three structures: _term_freq (term -> {doc_id: term_count}), _doc_freq (a Counter of how many documents contain each term), and _doc_lens (token count per document). Each document is its title and abstract joined by a space, so document length is the token count of that text, and the average document length is computed once, when the index is built. Scoring a query sums idf * tf_norm over the query terms, where tf_norm is the standard length-normalised BM25 term frequency.

The tokeniser lowercases, then splits on non-alphanumeric characters. It does no stemming. A production system would swap in a small stemmer, and the interface would stay the same.

idf(t)      = log((N - df + 0.5) / (df + 0.5) + 1.0)
tf_norm(t)  = (f * (k1 + 1)) / (f + k1 * (1 - b + b * dl / avgdl))
score(d, q) = sum over t in q of idf(t) * tf_norm(t)

The citation graph is a directed graph with papers as nodes and citations as edges. It is built once from the corpus. Forward edges go from a paper to its references, and backward edges go from a paper to its citations. Traversal is a breadth-first search seeded by the top BM25 hits and capped at two hops.

Two hops is a deliberate ceiling. One hop is too shallow, since the agent often wants the immediate ancestor or descendant. Three hops blows up the result size on a connected graph and tends to drift off topic. The hop limit is exposed as a config knob so a downstream loop can tighten it. RetrievalConfig clamps any max_hops above two back to two, so the knob only goes down.

Deduplication here means collapsing the lexical and graph results into one entry per paper. The two passes return overlapping sets, so the merge keys on paper id. Each paper's final score is a weighted blend.

final_score = w_bm25 * bm25_score_norm
            + w_graph * graph_score
            + w_recency * recency_score

bm25_score_norm is the BM25 score divided by the maximum BM25 score among the lexical top-k hits, which keeps the field between zero and one. graph_score is 1.0 at distance zero, which is the seed set of lexical top-k hits (top_k_lexical, default 10, and 5 in the demo), then 0.6 for one hop, 0.3 for two hops, and zero otherwise. Those constants are hardcoded in _graph_score. recency_score is a linear ramp from zero at the corpus minimum year to one at the maximum.

The default weights are 0.5, 0.3, and 0.2. Only the three weights live in config (w_bm25, w_graph, w_recency): a stale topic might tune recency down, while a fast-moving topic raises it.

The corpus is one hundred papers, generated by build_corpus(). Each paper belongs to one of five topics (attention sparsity, retrieval augmentation, low rank adapters, dataset distillation, and evaluation harnesses), twenty papers per topic. Titles and abstracts are templated from the topic label (a title looks like "Attention Sparsity Study 1"), not hand-written. Every paper from the third onward in a topic cites that topic's first two papers, and each topic has exactly one cross-topic edge: its sixth paper also cites the first paper of the next topic.

The two mock API clients (ArxivMockClient, SemanticScholarMockClient) read from the same corpus but expose different fields. Arxiv returns title, abstract, year, and authors. Semantic Scholar adds references and citations. RetrievalClient.search does not call either mock client. It works from the corpus directly, and every paper it returns has source="merged". Reconciling fields where the two clients disagree is deferred to a follow-up lesson.

Two later stages consume this output. The runner in lesson fifty-two reads paper.id, paper.title, and the top three sentences of the abstract as context for the experiment. The evaluator in lesson fifty-three reads paper.year and paper.references to attribute a baseline to a specific paper.

The retrieval client returns a RetrievalResult with the ranked list and the per-query metrics: hit count, average score, top score, and total wall time. The runner logs these, so a downstream observability pass can plot quality over time.

code/main.py defines Paper, ArxivMockClient, SemanticScholarMockClient, BM25Index, CitationGraph, RetrievalClient, and a deterministic demo. RetrievalClient takes the two mock clients in its constructor, but the search path uses only the corpus. The mock clients and the corpus live in the same file to keep it portable. The BM25 implementation is one class of about sixty lines, and the graph traversal is one method.

code/tests/test_retrieval.py covers the lexical path, the graph path, the merge, the dedup, and the empty query.

Lesson fifty produces a hypothesis. Lesson fifty-one searches the literature to see whether that hypothesis is already settled. Lesson fifty-two runs the experiment if it is not, and lesson fifty-three reads both the retrieval result and the experiment metrics to write the verdict. The retrieval client is the cheapest of the four stages and runs first in the orchestrator.