Adaptive RAG

Ask a research question and get an answer based on papers retrieved specifically for that question, from a library of around 1,000 AI and ML papers published on arXiv.

RAG means the answer is assembled from papers pulled out of that library at the moment you ask, rather than recalled from whatever the model absorbed in training. It can explain itself, too. Ask what RAG is and why it beats a plain LLM call for a task like this.

Plan

Hybrid Retrieval

Cross-Encoder Rerank

Grade

Synthesize

Reformulate

Try one of these

Answer

Sources

How it works

The corpus

1,000 papers across 20 AI/ML topics.

Selection is its own step. It writes a JSON manifest of arXiv metadata without downloading a PDF or spending anything, and every arXiv response is cached, so filters can be tuned and re-run for free. 304 papers are hand-picked and title-verified against arXiv. The other 696 come from per-topic keyword search behind gates: CS/ML categories only, a title blocklist for applied niches (clinical, finance, agriculture), and a reviewed exclude list.

Indexing

Per paper: skip it if it's already in ChromaDB (checked by paper_id, so a crashed run is safe to restart), pull the text out with PyMuPDF, discard anything under 100 characters as a scan, cut everything from the References or Acknowledgments heading onward, then index the abstract as its own chunk and the body as passage chunks.

Chunks target 500 tokens with 20% slack, are assembled sentence by sentence, and are never cut mid-sentence; anything under 20 tokens is dropped. Each chunk carries up to the last two sentences of the one before it, so a fact split across a boundary stays retrievable from either side: 17% overlap on average, measured across the corpus. That produces 24,278 records: 23,278 passages plus one abstract per paper. The sentence splitter is hand-written rather than NLTK because academic text breaks the usual rules, so it protects et al., Fig., i.e., initials like A. Vaswani, and captions like Table 3. Token counts use cl100k_base, matching the embedding model.

Everything is embedded with text-embedding-3-small (1536-dim, cosine), batched 100 per request and truncated at 8000 tokens. Each record keeps the text, the vector, and metadata: paper_id, title, authors, published, categories, field, chunk_index, chunk_type.

Hybrid retrieval

Semantic search alone conflates different uses of the same word. "Coverage" means one thing in RAG, another in facility location, another in conformal prediction, and cosine similarity can't separate them. BM25 can, because IDF weights rare terms heavily, which makes it sharp on proper nouns and coined names. The implementation is rank-bm25's BM25Okapi, held in memory over the same chunks.

So every sub-query runs three searches: 5 abstracts and 15 passages by cosine distance, plus 20 chunks by BM25. Abstracts get a reserved slice so the much larger passage population can't crowd paper-level summaries out. BM25 scores the raw passage text, with the title and abstract prefix that the embeddings rely on removed first.

Reciprocal Rank Fusion merges the lists on rank rather than score (k=60), which sidesteps the fact that cosine distances and BM25 scores aren't comparable quantities. A chunk both searches found outranks one only a single search found. The top 25 survive.

Reranking

RRF fuses rankings; it doesn't judge relevance. So those 25 go to a FlashRank cross-encoder (ms-marco-TinyBERT-L-2-v2, 2-layer) that reads the query and the chunk together and scores the pair. It runs locally: no key, no network, no rate limit.

All 25 are scored, the top 8 are taken, and each of those still has to clear a bar. Often fewer than 8 come back.
The bar depends on where the chunk stood before reranking. Fusion top 10 needs 0.5; ranks 11-25 need 0.8. A high score with no upstream support is usually a keyword coincidence. The score distribution is close to bimodal: across a sample of reranked pools the median is 0.998 and the 25th percentile is 0.952, but roughly a fifth of candidates land under 0.02. 0.5 sits in the empty middle, which is why it separates cleanly instead of slicing into a dense band.
A hand-picked paper already in the fusion top 10 has its score floored at 0.5, so the small model can't drop a canonical source. The floor raises a weak score to the bar; it never lowers a real one.
If nothing clears its bar, only the single best chunk comes back, and only above 0.5. Otherwise the sub-query returns nothing at all.

Agent loop (LangGraph)

A single retrieval pass tends to answer a compound question from whichever paper dominates the search. The agent splits the question instead, retrieves for each part separately, checks whether what came back actually covers the question, and searches again with reworded queries if it doesn't.

Nodes

Node	Purpose	LLM Call
`planner_node`	Classify query as simple/compound, generate sub-queries	`gpt-5.6-luna` + `QueryPlan`, effort `none`
`retriever_node`	Run full retrieval pipeline per sub-query, deduplicate by `chunk_id`	None
`grader_node`	Evaluate context sufficiency, identify missing elements	`gpt-5.6-luna` + `InformationCheck`, effort `none`
`reformulator_node`	Generate new sub-queries targeting gaps	`gpt-5.6-luna` + `QueryPlan`, effort `none`
`synthesizer_node`	Build final answer from accumulated context, extract citations. Answers from partial context while staying grounded. Unrelated queries still get "I don't have enough information"	`gpt-5.6-luna`, effort `low`

Graph

START → planner → retriever → grader → [route] → synthesizer → END
                                            │
                              context_sufficient=False
                              AND retry_count < 2
                                            │
                                       reformulator → retriever (loop)

One retry maximum, two retrieval passes total. retriever_node gets ChromaDB and BM25 through a factory closure, which keeps non-serializable objects out of graph state.

State

AgentState has 10 fields. Eight overwrite on write; two use the Annotated[list, add] reducer and concatenate. accumulated_context has to, or each pass would erase the last one and multi-pass retrieval would quietly do nothing. The reducer merges but doesn't dedupe, so retriever_node filters by chunk_id itself. all_sub_queries keeps the full search history, since sub_queries is overwritten on every retry.

Latency

Three things keep a compound query from costing three times a simple one:

All sub-query embeddings go out in one batched call before retrieval starts.
Retrieval runs on a thread pool, one worker per sub-query. Workers only read shared state and return their own lists, so there's nothing to race on. Deduplication happens once, single-threaded, after they all finish.
The candidate pool is wider than precision alone would call for, because these are local lookups and cheap. In one reproduced case "Attention Is All You Need" never entered the pool at all, and no reranker tuning saves a chunk that was never fetched.

What was measured

Medians across trials after a warmup, from bench_latency.py, re-run against the current 24,278-chunk corpus. The fan-out row is the mean of two separate runs over the same three sub-queries.

Change	Before → after	What the number is
Reranker model: MiniLM-L-12 → TinyBERT-L-2	1.309s → 0.067s 19.4x	Time to score one 25-candidate pool on CPU. Reranking runs once per sub-query per pass, so this is the single largest cost in a compound query. Speed is measured; the ranking precision given up is not, and there is no evaluation here that quantifies it.
BM25 index: rebuild at boot → prebuilt pickle	2.694s → 0.358s 7.5x	Startup cost of getting a usable keyword index. Rebuilding means reading all 24,278 documents out of ChromaDB and tokenizing them. Originally a workaround for a 512MB memory ceiling; kept because a cold Space should not spend three seconds on work already done.
Sub-query embedding: one call each → one batched call	0.709s → 0.263s 2.70x	Time to embed the sub-queries of one compound question. Nothing about the computation changed; this is two fewer network round trips to OpenAI on a three-part question.
Retrieval fan-out: sequential → thread pool	480ms → 366ms 1.31x	Three sub-queries retrieved concurrently instead of one after another, embeddings prepared beforehand so this isolates the fan-out itself. The work genuinely overlaps: instrumenting the workers shows all three running at once, 1,205ms of retrieval finishing in 423ms of wall clock. The wall-clock win is smaller than that overlap because reranking is CPU-bound ONNX inference and the workers share one core, so each retrieval slows while the others run. The gain peaks at three sub-queries and falls to 1.19x at five, and one to three is the range the planner emits. The benchmark also asserts both modes return the same chunk IDs in the same order (0 mismatches), so parallelism cannot silently change results.

The reranker swap carries a real tradeoff: a 2-layer cross-encoder is a weaker judge than a 12-layer one, and nothing here measures how much weaker. It made compound queries usable on a free CPU tier, and retrieval recall across the 304 hand-picked papers held at 99.7% afterwards.

Retrieval has its own check: eval_anchor_recall.py asks every hand-picked paper to find itself through the full hybrid and rerank stack.

API and deployment

It ships as a Docker container on Hugging Face Spaces (free CPU tier, 16GB), running Uvicorn on port 7860. FastAPI serves the same agent the CLI uses, compiled once in a lifespan handler at startup and shared across requests.

BM25 is built offline and pickled; the server loads that file instead of rebuilding the index at boot.

The indexer rebuilds that pickle itself when it finishes writing ChromaDB, so the keyword index can't be left behind by a corpus rebuild. Startup then compares the pickle's stored chunk count against the live collection count (a metadata lookup, not a read) and rebuilds in memory with a warning if they disagree.

Endpoints

Endpoint	Purpose
`POST /query`	Run a full agent query. Returns answer, sub-queries, retry count, retrieved chunks with scores, and citations.
`POST /query/stream`	The same query as Server-Sent Events, emitting a progress event as each graph node returns, then the full result. This is what drives the pipeline diagram on the search page.
`GET /papers`	Returns the full corpus as a browsable list: paper ID, title, authors, field, year, and arXiv URL. Loaded once from the manifest at startup.
`GET /health`	Returns `{"status": "ok"}`. Used for health checks.

Known limitations

A cross-encoder this small can be confidently wrong. Both TinyBERT and MiniLM will occasionally score an unrelated chunk very high. A Video Diffusion passage scored 0.96 against a transformer-architecture query from fusion rank 14, clearing even the strict bar. The tiered threshold is designed to catch false positives with weak upstream support; it cannot catch one the model is simply wrong about. Fixing that means a larger reranker, which the latency budget on a free CPU tier does not have room for.
Citations are provenance, not attribution. The list is every paper that was in the synthesis context, not the ones the answer provably drew on, and there are no inline references. Claim-level attribution needs a verification pass over the generated text, which this system does not do.
Retrieval is measured; answer quality is not. Anchor recall covers whether the right papers come back, and the latency harness backs every performance number on this page. There is no faithfulness scorer or labelled QA set, so the grader's sufficiency calls go unchecked: a false positive there ends retrieval a pass early.
Consistency between the two indexes is eventual, not atomic. The BM25 rebuild runs after ChromaDB is written, with no transaction spanning both. The chunk-count check at startup narrows the window rather than closing it; a divergence that preserves the count would pass unnoticed.
Answers are not reproducible. The model accepts only its default temperature, so sampling is always on. The same question can return different wording across runs.
Scoped out for a public demo. No authentication or rate limiting, and CORS is open to every origin. FlashRank also re-downloads its model on a cold start because it caches to /tmp.

How the design got here:

Baseline

A basic pipeline working end to end - embed, retrieve by cosine top-k, and generate a response.

Chunking, dedicated slot for abstracts

Modified chunking strategy, and added dedicated chunks for abstract sections of papers. Initially added a fixed sliding window which often split sentences down the middle, which produced fragments of passages and confused the model. Replaced it with improved chunking logic to guarantee assembling whole sentences, or pushing the sentence to the next chunk if too long. The splitter had to account for academic text to work around non sentence ending periods, ex: et al., Fig. contain periods without ending sentences. Abstracts became their own chunk type, with reserved slots at the retrieval stage.

Hybrid retrieval: adding a lexical signal

Cosine similarity cannot separate different uses of a single word. "Coverage" might mean one thing in RAG evaluation and another in the context of location or prediction models, and one query might yield all three in the same returned set of chunks. The BM25 ranking function helps with this. Inverse document frequency weights a term by how rare it is across the corpus, so a name or title carries more weight than a common word. Both now run per sub-query and merge through RRF (reciprocal rank fusion), allowing the merge to bypass score normalization between BM25 and cosine similarity.

Cross-encoder reranking

RRF does not judge relevance on its own. Instead, it rewards agreement between two searches that score a query and document independently. However, the query and document are never read together. A cross encoder does this by taking the query and candidate as one input and re-ranking relevance of selected chunks given the query. Unfortunately, it's expensive to run across the corpus, leading to the split of retrieving first and reranking second. Previous searches narrow ~25,000 chunks to the top 25, and the expensive reranking model judges the 25.

The agent loop, with a bounded retry

A single retrieval pass answers a compound question from whichever paper the combined query string sits closest to in similarity. However, complex queries might require multiple sources that don't necessarily appear more relevant to the given input. Using a LangGraph agent solves this by decomposing the question into subqueries, retrieving for each part separately, and grading whether what came back covers the question or not, searching again with reworked queries if not. Currently capped to one retry (two total passes) for speed, putting a firm ceiling on worst case latency and API spend. During testing, a third pass rarely contributed meaningfully better chunks.

Corpus: keyword scrape to a curated manifest

The first iteration of the corpus was roughly 480 papers pulled via simple keyword and category search. It demonstrated the vocabulary problem mentioned above. For example, searching "retrieval augmented generation" returned papers on augmented reality. The corpus also did not adequately cover the intended topic range to always provide a good answer, sometimes yielding no answer at all. To remedy this, paper selection got its own dedicated phase that writes a reviewable JSON manifest from cached arXiv metadata without downloading a PDF, so the filters could be fine tuned and rerun for free. Final (current) split ended up at 1,000 papers across 20 topics in AI/ML, including 304 manually picked and title verified papers. The remaining 696 papers were filled on a per topic basis based on categories, passing through a title blocklist and a reviewed exclusion list.

Reranking inside a latency and memory budget

Reranking runs once per subquery per pass, so a compound query with a second retry run would get expensive. Initial versions of this project used a hosted reranker via a free usage-limited API. It was later moved to a local model, removing token and rate limits, but introducing a memory and time constraint of its own. Running a model in process blew past the 512MB ceiling on the original host, which moved the service to a Hugging Face Space and pushed the BM25 index out of startup and into a build artifact. It is pickled when the corpus is built and loaded straight off disk, so the server never tokenizes 24,278 documents at boot: 0.36s instead of 2.7s before it can answer anything. Initial model choice was ms-marco-MiniLM-L-12-v2 via FlashRank, which on measurement was still two thirds of request time per query. It was later replaced by a 2 layer model (ms-marco-TinyBERT-L-2-v2), resulting in reranking being almost 20x faster.

Threshold and context tuning

A few of the tuning numbers turned out to be doing less than expected. The reranker cutoff was set at 0.95. TinyBERT scores turn out to be close to bimodal, mostly landing either near zero or near one, so a cutoff at 0.95 sat inside the upper cluster and was rejecting real matches over hairline differences. It was lowered to 0.5, which sits in the empty middle, and chunks passed to the synthesizer went from 3-5 up to 8. The synthesizer prompt was also refusing to answer whenever coverage was partial, so it now answers with what it has and states what is missing. Retrieval moved onto a thread pool with subquery embeddings batched into a single call. The candidate pool was widened after a test case where "Attention Is All You Need" never made it into the pool at all - no amount of reranker tuning fixes a chunk that was never fetched.

The grader was reading headers, not evidence

Passage chunks are stored with the paper title and abstract prefixed onto them, which helps a short fragment embed closer to the right neighborhood. The grader built its context summary from the first 200 characters of each chunk, which on those passages is the prefix and nothing else. Measured over 200 chunks, the median amount of actual body text the grader could see was zero, and so was the maximum. Every sufficiency check, and therefore every retry decision, was being made off paper titles. The prefix is now stripped once retrieval hands off, so the grader and synthesizer both read the passage instead of the header. The prefix does a real job at embedding time, which is what kept it from looking wrong for so long.

Concurrency, overlap, and a corpus rebuild

The query route was declared async but called the graph synchronously, so it held the event loop for a whole request and two users would queue behind each other. It now runs on a worker thread. A second route streams progress as each graph node finishes, which is what drives the pipeline diagram on the search page. Chunk overlap went from 5.8% to 17%, and the corpus was reindexed and repickled to apply it, taking it from 20,845 to 24,278 chunks.