Adaptive RAG

Corpus Sourcing

A curated 1,000-paper set spanning 20 AI/ML topics. This replaced an earlier 3-tier keyword/category scrape of ~480 papers whose keyword search was badly polluted (e.g. "retrieval augmented generation" pulling in augmented-reality papers).

Selection produces a JSON manifest of chosen papers with arXiv metadata, touching no PDFs and spending no money. ~304 hand-picked anchors and curated papers (verified against arXiv titles) carry the corpus quality. The remaining ~696 slots are filled by per-topic keyword search with quality gates: CS/ML categories only, a title blocklist filtering applied niches (clinical, finance, agriculture, etc.), and a reviewed exclude set of specific off-topic strays. All arXiv responses are cached, so re-running while tuning filters is instant.

Per-Paper Indexing Pipeline

Per-paper sequence

• Skip if indexed — paper_exists_in_db() checks ChromaDB by paper_id. Idempotent and safe to re-run after interruption.
• Extract text via PyMuPDF. Papers under 100 characters (scanned/malformed) are skipped.
• Strip garbage — truncate at standalone References, Bibliography, or Acknowledgments headings. Drop floating page numbers, collapse excess blank lines.
• Extract and embed the abstract as its own chunk.
• Chunk the body, enrich each chunk with context, embed in a batch, and store.

Chunking rules

• Token counts use cl100k_base (tiktoken), matching the embedding model's encoding.
• Target size 500 tokens with 20% tolerance (up to 600 tokens). Chunks are assembled sentence by sentence and never split mid-sentence.
• Custom sentence splitter (no NLTK) protects academic abbreviations (et al., Fig., i.e., e.g.), single initials (A. Vaswani), and numbered references (Table 3.).
• Fragments under 20 tokens are discarded.

Embedding & storage

• text-embedding-3-small, 1536-dim vectors, cosine space. Batched 100 per request. Texts over 8000 tokens truncated before sending.
• Each chunk stores: id, document text, embedding, and metadata (paper_id, title, authors, published, categories, chunk_index, chunk_type, field).

Architecture Overview

A multi-strategy RAG pipeline where each feature solves a specific failure mode of the previous iteration:

Base:      embed → retrieve (cosine top-k) → generate
+ Chunking: sentence-boundary splitting, abstract indexing
+ Hybrid:   BM25 keyword search + RRF fusion
+ Rerank:   FlashRank cross-encoder (local ONNX)
+ Agent:    LangGraph (decompose → retrieve → grade → retry → synthesize)

Hybrid Retrieval

Pure semantic search conflates different uses of the same word ("coverage" in RAG systems vs. facility location vs. conformal prediction). Adding BM25 keyword search fixes this — BM25 weights rare terms via IDF, making it naturally discriminative on proper nouns and coined terms.

Three parallel searches

• Semantic abstracts (ABSTRACT_K=6) — guarantees abstract representation.
• Semantic passages (TOP_K=14) — enriched chunk embeddings.
• BM25 all chunks (BM25_K=20) — keyword scoring on raw passage text (enrichment prefix stripped).

Reciprocal Rank Fusion (RRF) merges the ranked lists with k=60, deduplicates by chunk_id. A chunk ranked #1 by BM25 and #3 semantically scores higher than one ranked #1 semantically only — RRF rewards agreement between signals. Returns top 24 candidates.

Cross-encoder reranking

RRF is a ranking fusion, not a relevance judgment. The 24 candidates are reranked by a FlashRank cross-encoder (ms-marco-TinyBERT-L-2-v2, ~4MB ONNX model, 2-layer) that reads query and chunk text together, producing a true relevance score.

• Runs entirely locally — no API, no key, no network, no rate limits. Lazy-loaded once into a module-level singleton.
• Two-step selection: score all 24, take the top 8 by score (top_n=8), then each survivor must clear a relevance bar — so the result is at most 8 chunks, often fewer.
• Tiered bar by fusion rank: a chunk in the post-fusion top 10 (RRF_TRUST_CUTOFF, both searches agreed) needs only RELEVANCE_THRESHOLD = 0.5; one that barely made the pool (rank 11–24) must clear the stricter STRICT_RELEVANCE_THRESHOLD = 0.8. A high score with no upstream support is treated as a likely keyword coincidence. (TinyBERT scores cluster 0.94–0.99, so the base floor is deliberately low.)
• Anchor/curated floor: a hand-picked paper already in the top 10 has its reranker score floored at 0.5 (TRUSTED_TIERS) before the top-8 cut, so the small model can't drop a known-canonical source. The floor only ever raises a weak score to the base bar — never lowers a real one.
• Gated fallback: if none of the 8 clear their bar, only the single highest-scoring chunk is returned, and only if it still scores ≥0.5 (FALLBACK_MIN_SCORE) — otherwise the sub-query legitimately returns nothing rather than force a low-confidence citation.

Reranking was first moved from a hosted Jina API to local FlashRank to drop per-call cost and rate-limit handling — the agent reranks once per sub-query per retrieval pass, so compound queries were generating many billed API calls. The 12-layer MiniLM model then became the dominant latency cost (~7s/call, ~67% of total request time on the CPU-only Space), so it was switched to the 2-layer TinyBERT — ~20× faster, trading some ranking precision, which the synthesizer absorbs by answering from partial context.

Agent Loop (LangGraph)

Single-pass retrieval favors the dominant paper on compound queries. The agent decomposes queries into sub-queries, retrieves independently for each, grades coverage, and retries with reformulated queries when context is insufficient.

State

AgentState is a TypedDict with 10 fields. Two use the Annotated[list, add] reducer, which concatenates onto the existing list instead of overwriting. On accumulated_context — without it, each retrieval pass would overwrite the previous context rather than accumulate it, silently breaking multi-pass retrieval (duplicates are still removed inside retriever_node, since the reducer merges but can't dedupe). The second, all_sub_queries, is the full history of every sub-query searched across all passes, since the plain sub_queries field is overwritten by the reformulator on each retry.

Nodes

Graph

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

Max 1 retry (2 total retrieval passes). The retriever_node receives ChromaDB and BM25 via a factory closure (make_retriever_node(collection, bm25_index)), keeping non-serializable infrastructure out of the graph state.

Latency & Parallel Retrieval

A compound query thatonce took ~60s because every sub-query's work (embedding, ChromaDB queries, BM25, rerank) ran one after another. Three changes cut that specific to ~15s with no reranker model change:

• All sub-query embeddings are computed in one batched call before retrieval starts.
• Per-sub-query retrieval runs in parallel on a thread pool, one worker per sub-query. Each worker only reads shared resources (the collection, the BM25 index, the reranker singleton) and returns its own list, so there are no race conditions; deduplication runs once, single-threaded, after all workers finish.
• The candidate pool was widened (ABSTRACT_K 4→6, TOP_K 10→14, BM25_K 15→20, pool 20→24). These are local lookups, not network calls, so the cost is small — the point is to give a genuinely relevant canonical paper a chance to reach the reranker at all. In one reproduced case "Attention Is All You Need" never entered the pre-rerank pool, and no reranker tweak can rescue a chunk that was never fetched.

Swapping the larger MiniLM model back in under the new design was tested as a follow-up: it doubled latency and did not fix the residual miss noted under Known Limitations, so it was not adopted.

API & Deployment

FastAPI serves the same LangGraph agent as the CLI over HTTP. The compiled agent is initialized once at startup via FastAPI's @asynccontextmanager lifespan pattern and shared across all requests.

AOT BM25 serialization

Building the BM25 index from scratch at server boot originally exceeded Render's free-tier memory limits. A separate offline script (build_bm25_cache.py) pre-builds and pickles the index; the server loads the serialized file on startup instead of reconstructing it. The app now runs on Hugging Face's 16GB CPU tier, but the AOT pattern is kept for fast startup.

Endpoints

Hugging Face Spaces deployment

Docker container on Hugging Face Spaces (free CPU tier, 16GB RAM). Uvicorn serves the ASGI app on port 7860. CORS middleware with allow_origins=["*"].

Tech Stack

• Language: Python 3.12
• Embedding: OpenAI text-embedding-3-small (1536-dim)
• LLM: GPT-4o-mini (planning, grading, reformulation), gpt-5.4-nano (synthesis)
• Vector store: ChromaDB persistent client
• Reranker: FlashRank (ms-marco-TinyBERT-L-2-v2, ~4MB local ONNX)
• BM25: rank-bm25 BM25Okapi
• Agent: LangGraph stateful graph
• API: FastAPI + Uvicorn
• PDF parsing: PyMuPDF (fitz)
• Tokenization: tiktoken (cl100k_base)
• Deployment: Hugging Face Spaces (Docker, free CPU tier)

Dependencies

openai>=1.30.0        # Embeddings + GPT-4o-mini + gpt-5.4-nano
chromadb>=0.5.0       # Vector store
PyMuPDF>=1.24.0       # PDF parsing
feedparser>=6.0.0     # arXiv Atom feed
requests>=2.31.0      # arXiv PDF downloads
python-dotenv>=1.0.0  # .env loading
tiktoken>=0.7.0       # Token counting
rank-bm25>=0.2.2      # BM25Okapi implementation
flashrank>=0.2.10     # Local cross-encoder reranker
langgraph>=0.2.0      # Stateful agent graph
fastapi>=0.115.0      # HTTP API framework
uvicorn>=0.34.0       # ASGI server

Known Limitations

• A reranker can be confidently wrong. Both TinyBERT and the larger MiniLM can score a topically unrelated chunk very high on some phrasings (e.g. a Video Diffusion / UNet chunk scoring 0.96 against a transformer-architecture query, clearing even the strict 0.8 bar at fusion rank 14). The tiered threshold catches false positives that also had weak upstream support, but not cases where the model is simply wrong regardless of support. Not yet root-caused.
• The API serializes concurrent requests. async def query() calls the synchronous agent.invoke() directly (no run_in_threadpool), so it blocks the event loop for a whole request. Retrieval is parallel within one request, but two overlapping HTTP requests are handled one after the other.
• The reranker re-downloads on cold starts. FlashRank caches its model to /tmp by default, so an ephemeral host that resets between runs re-downloads it each cold start. A persistent cache directory or bundling the model would remove that cost.
• No streaming. The graph runs to completion before any output; the API returns the full response at once.
• No evaluation framework. The grader's sufficiency judgments are not measured systematically.
• Wide-open CORS. allow_origins=["*"] is fine for development but should be scoped for production.