Lab 01 — RAPTOR Tree & Mini-GraphRAG

Phase 06 · Lab 01 · Phase README · Warmup

The problem

Phase 05's hybrid retriever is the right default, and it fails on a whole class of questions — the ones where no single chunk contains the answer:

1. Synthesis — "Summarize how the three Mars missions each advanced exploration." The answer is the combination of several chunks; top-k over leaves returns fragments.
2. Global / thematic — "What are the main themes in this corpus?" The answer is a property of the whole corpus; no chunk says "the themes are …".
3. Multi-hop — "Who works with the person who owns the fraud model?" You must traverse entities and relations, not match a paragraph.

You build the two advanced-retrieval architectures that answer them — the three acronyms Citi's JD names explicitly (RAPTOR, GraphRAG, LightRAG) — with the embedder, summarizer, and triple-extractor injected so the whole pipeline is deterministic and testable.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # the RAPTOR + GraphRAG worked example

Success criteria

• Your RAPTOR tree has strictly decreasing node counts per level up to a single root, and each summary node lists valid children one level down.
• greedy_cluster is deterministic and groups the same-theme chunks together.
• On the synthesis query, collapsed_tree_search ranks a summary node (level > 0) above any single leaf — and you can explain, from the bag-of-words math, why the summary that fused the chunks wins.
• detect_communities groups the connected entities deterministically; global_query returns the relevant community summary; local_query returns the right neighborhood.
• All 33 tests pass under both lab and solution.

How this maps to the real stack

• RAPTOR (build_raptor_tree + collapsed_tree_search) is the recursive-summary tree from Sarthi et al., 2024. Production builds cluster with UMAP + Gaussian-mixture soft clustering and summarize with a real LLM; retrieval runs the same collapsed-tree pool you built, over a normal vector store (pgvector / Pinecone / Weaviate) that just happens to also hold the summary nodes. LangChain and LlamaIndex both ship a RAPTOR pack.
• GraphRAG (extract_graph → detect_communities → summarize_communities → global_query) is Microsoft GraphRAG (Edge et al., 2024): an LLM extracts entities/relations, Leiden finds hierarchical communities, an LLM writes community reports, and a map-reduce over community summaries answers global queries. Your connected-components pass is the deterministic stand-in for Leiden; your global_query is the reduce step. The graph lives in Neo4j (or networkx / a property graph); Neo4j's neo4j-graphrag package wires this to an LLM.
• LightRAG (Guo et al., 2024) is the cheaper cousin: dual-level retrieval — local_query (low-level, entity-anchored) and global_query (high-level, theme-anchored) — plus incremental graph updates (no full re-index when a document arrives). Your two query functions are exactly that local/global split.
• pgvector / Neo4j are the two stores this phase maps to Citi's JD: the vector index holds chunk and summary embeddings (RAPTOR); the graph DB holds entities, relations, and community assignments (GraphRAG/LightRAG). Real systems use both (hybrid), which is why the JD names them together.

Limits of the miniature. A hashing bag-of-words embedder has no semantics (it matches tokens, not meaning), so the corpus is stylized to make the mechanism visible; a real embedder would cluster on meaning. Real clustering is soft (a node can belong to several parents) and GMM-based; real community detection is Leiden over a weighted graph, not connected components; real extraction and summarization are LLM calls whose quality varies and whose indexing cost is the headline tradeoff (GraphRAG can be thousands of LLM calls to index a corpus). The control flow and the retrieval idea — collapse the tree, summarize the community, split local from global — are exactly what you built.

Extensions (your own machine)

• Tree-traversal retrieval: implement the other RAPTOR mode — start at the root, descend into the highest-scoring child at each level — and compare its results and node-visits to the collapsed-tree search you built.
• Real embeddings: swap the injected embed for a sentence-transformer (one function change) and re-run; watch the clustering group on meaning instead of shared tokens.
• Leiden: replace detect_communities with networkx's Louvain/Leiden and detect sub-communities within a connected component; summarize the community hierarchy.
• Incremental GraphRAG (LightRAG): add a add_document(graph, doc, extractor) that merges new triples and only re-summarizes the touched communities — the update path GraphRAG lacks.
• Neo4j: persist the graph with the Neo4j Python driver and answer local_query with a Cypher MATCH (e)-[r]-(n) query.

Interview / resume signal

"Built RAPTOR (recursive clustering + summarization tree with collapsed-tree retrieval) and a mini-GraphRAG (entity/relation extraction → community detection → community summaries → global vs local query) from scratch, with injected embedder/summarizer/extractor for deterministic tests — proving, in the retrieval trace, that a summary node beats any single leaf on a synthesis query and that community summaries answer global/thematic questions no chunk contains."