GraphRAG with Leiden communities

Open in Colab ▶️ Run this notebook in Colab

GraphRAG with Leiden communities#

Track 02 - RAG · Notebook 08 · Runtime: ≈1 min on CPU

Prerequisites: 02_rag/07 (RAPTOR).

Paper: Edge et al. 2024, From Local to Global: A GraphRAG Approach to Query-Focused Summarization (2404.16130).

What#

RAG struggles with global-sensing queries like “what are the main themes across this corpus?” - chunk retrieval surfaces local text but can’t compose a global summary.

GraphRAG solves this with three components:

  1. Entity + relation extraction. LLM pass over every chunk produces (entity_a, relation, entity_b) triples.

  2. Graph construction. Triples become a weighted undirected graph; edge weight = co-occurrence count.

  3. Community detection + summarisation. Cluster the graph into communities (Leiden algorithm) and summarise each community with an LLM. At query time, ask each community summary to contribute a partial answer; map-reduce the partials into a final answer.

We stub the LLM calls (deterministic keyword extractor for entities; nearest-member summaries for community summaries) and use networkx’s greedy modularity communities as a Leiden stand-in. The pipeline shape is identical to Microsoft’s production GraphRAG.

from llm_systems_cookbook.nb import bootstrap

import re
from collections import Counter, defaultdict

import networkx as nx
import numpy as np

s = bootstrap("02_rag_08_graphrag_leiden")

Mini corpus - synthetic illustrative snippets#

15 short, made-up snippets about well-known tech companies. They exist only to populate a co-occurrence graph with recognisable names so the GraphRAG pipeline is easy to follow - they are not factual claims about any company’s products, partnerships, offices, or roadmap. Two entities share an edge if they appear in the same snippet.

# Illustrative entities - real company/product names used only because
# readers will recognise them. The alias lists are not meant to be
# authoritative taxonomies.
ENTITIES = {
    "anthropic": ["anthropic", "claude"],
    "openai": ["openai", "gpt"],
    "google": ["google", "gemini", "deepmind"],
    "meta": ["meta", "llama"],
    "microsoft": ["microsoft", "copilot"],
    "nvidia": ["nvidia", "h100", "blackwell"],
    "amd": ["amd", "mi300"],
    "sf": ["san francisco"],
    "nyc": ["new york"],
    "amazon": ["amazon", "bedrock"],
}

# Synthetic snippets: plausible-sounding sentences engineered to make the
# co-occurrence graph have clean community structure. Do not treat as news.
SNIPPETS = [
    "anthropic's claude 4 benchmarks against gpt-4 and gemini.",
    "openai and microsoft extended their partnership, embedding gpt models in copilot.",
    "google's deepmind released gemini 2.0 focusing on long context.",
    "meta llama 4 will be trained on nvidia h100 clusters.",
    "nvidia h100 remains the dominant training accelerator at major labs.",
    "anthropic moved its offices to san francisco near openai's headquarters.",
    "meta opened a new research lab in new york.",
    "amazon bedrock hosts claude, llama, and titan models side by side.",
    "amd mi300 aims to compete with nvidia's blackwell chips.",
    "microsoft copilot runs on gpt-4 turbo.",
    "deepmind and anthropic researchers co-authored a safety paper.",
    "openai's valuation rose; anthropic's also climbed.",
    "gemini 2.0 was trained on google's tpu pods, not nvidia blackwell.",
    "llama 4 will be distributed via amazon bedrock and meta's own servers.",
    "nvidia blackwell is expected to dominate 2026 data-centre upgrades.",
]


def extract_entities(text: str) -> set[str]:
    text_l = text.lower()
    ents = set()
    for canonical, aliases in ENTITIES.items():
        if any(alias in text_l for alias in aliases):
            ents.add(canonical)
    return ents


doc_entities = [extract_entities(s) for s in SNIPPETS]
entity_counts = Counter()
for es in doc_entities:
    entity_counts.update(es)
print(f"extracted entities across {len(SNIPPETS)} snippets:")
for e, c in entity_counts.most_common():
    print(f"  {e:>10}  {c}")

Build the graph#

Nodes = entities. Edges = co-occurrence in the same snippet, weighted by count.

G = nx.Graph()
for es in doc_entities:
    for a in es:
        G.add_node(a)
    es_list = sorted(es)
    for i in range(len(es_list)):
        for j in range(i + 1, len(es_list)):
            a, b = es_list[i], es_list[j]
            if G.has_edge(a, b):
                G[a][b]["weight"] += 1
            else:
                G.add_edge(a, b, weight=1)

print(f"graph: |V|={G.number_of_nodes()}  |E|={G.number_of_edges()}")
print(f"top-weight edges:")
for a, b, w in sorted(G.edges(data="weight"), key=lambda x: -x[2])[:5]:
    print(f"  {a} - {b}  weight={w}")

Communities (Leiden-style)#

Real GraphRAG uses the Leiden algorithm (Traag 2019). We use networkx’s greedy modularity communities - a well-behaved approximation. Output: a partition of nodes into communities plus per-community modularity.

from networkx.algorithms.community import greedy_modularity_communities, modularity

communities = list(greedy_modularity_communities(G, weight="weight"))
mod = modularity(G, communities, weight="weight")
print(f"{len(communities)} communities, modularity = {mod:.3f}")
for i, c in enumerate(communities):
    print(f"  community {i}: {sorted(c)}")

Community summaries (stub)#

For each community, collect every snippet that mentions at least one member entity, then build a deterministic “summary” by joining the first sentences. A real pipeline would call an LLM here.

community_summaries = []
for c in communities:
    members = set(c)
    relevant_snippets = [s for s, es in zip(SNIPPETS, doc_entities, strict=True) if es & members]
    summary = " ".join(relevant_snippets[:3]) or "(no coverage)"
    community_summaries.append({"entities": sorted(members), "summary": summary})

for i, cs in enumerate(community_summaries):
    print(f"community {i} ({', '.join(cs['entities'])})")
    print(f"  {cs['summary'][:160]}...")

Global-sensing vs factoid queries#

Two query types:

  • Global: “what are the main groupings of AI companies in this corpus?” - GraphRAG’s summaries should surface entire communities.

  • Factoid: “which chip does Meta use for Llama 4?” - standard snippet retrieval wins.

We score GraphRAG on the global query (looking for at least 2 communities’ summaries to mention entities relevant to the query) and compare against flat snippet retrieval.

def global_query_coverage(query_entities: set[str]) -> float:
    '''Fraction of query entities covered by the community summaries.'''
    covered = set()
    for cs in community_summaries:
        if set(cs["entities"]) & query_entities:
            covered.update(set(cs["entities"]) & query_entities)
    return len(covered) / max(len(query_entities), 1)


def flat_global_retrieve(query: str, k: int = 3) -> set[str]:
    '''Return the entities present in the top-k most relevant snippets.'''
    q_tokens = set(tokenize(query))
    scores = [(i, len(q_tokens & set(tokenize(s)))) for i, s in enumerate(SNIPPETS)]
    top_ids = [i for i, _ in sorted(scores, key=lambda x: -x[1])[:k]]
    ents: set[str] = set()
    for i in top_ids:
        ents.update(doc_entities[i])
    return ents


def tokenize(text: str) -> list[str]:
    return re.findall(r"[a-z]+", text.lower())


# Global query: what AI companies and hardware appear together?
query_entities = {"anthropic", "openai", "google", "meta", "microsoft", "nvidia", "amd", "amazon"}
graphrag_cov = global_query_coverage(query_entities)
flat_ents = flat_global_retrieve("AI companies and hardware providers", k=3)
flat_cov = len(flat_ents & query_entities) / len(query_entities)

print(f"global query coverage:")
print(f"  GraphRAG community summaries: {graphrag_cov:.1%} of query entities surfaced")
print(f"  flat top-3 retrieval:           {flat_cov:.1%}")

s.check(
    "graph_has_substantial_edges",
    lambda: G.number_of_edges() >= 5,
    msg=f"|E| = {G.number_of_edges()}",
)
s.check(
    "communities_nonempty",
    lambda: all(len(c) >= 1 for c in communities),
)
s.check(
    "modularity_positive",
    lambda: mod > 0.05,
    msg=f"modularity = {mod:.3f}",
)
s.check(
    "graphrag_global_coverage_beats_flat",
    lambda: graphrag_cov > flat_cov,
    msg=f"graphrag={graphrag_cov:.1%}  flat={flat_cov:.1%}",
)
s.check(
    "at_least_two_communities",
    lambda: len(communities) >= 2,
    msg=f"n_communities = {len(communities)}",
)

Communities and what they unlock#

Two panels. Left: community sizes, showing how greedy modularity carved the entity graph into a small number of coherent groups. Right: coverage of the “AI companies + hardware” global query - GraphRAG’s community summaries surface the full set of relevant entities, while flat top-k snippet retrieval only reaches a fraction.

import matplotlib.pyplot as plt

sizes  = [len(c) for c in communities]
labels = [", ".join(sorted(c)[:2]) + ("+" if len(c) > 2 else "")
          for c in communities]
order  = np.argsort(sizes)[::-1]
sizes  = [sizes[i]  for i in order]
labels = [labels[i] for i in order]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.6))
ax1.barh(range(len(sizes)), sizes, color="tab:blue")
ax1.set_yticks(range(len(sizes)))
ax1.set_yticklabels(labels, fontsize=8)
ax1.invert_yaxis()
ax1.set_xlabel("|community| (entities)")
ax1.set_title(f"{len(communities)} communities, mod={mod:.2f}")
for i, sz in enumerate(sizes):
    ax1.text(sz + 0.05, i, str(sz), va="center", fontsize=8)

methods = ["flat top-3", "GraphRAG summaries"]
covs    = [flat_cov, graphrag_cov]
bars = ax2.bar(methods, covs, color=["tab:gray", "tab:orange"])
for bar, v in zip(bars, covs):
    ax2.text(bar.get_x() + bar.get_width() / 2, v + 0.02,
             f"{v:.0%}", ha="center", fontsize=9)
ax2.set_ylabel("query-entity coverage")
ax2.set_ylim(0, 1.05)
ax2.set_title("global query coverage")
fig.tight_layout()
plt.show()

Exercises#

  1. Real Leiden. pip install leidenalg python-igraph and swap the community detector. Leiden strictly dominates Louvain (Traag 2019); compare modularity on this graph.

  2. Map-reduce final answer. Pass each community’s summary into an LLM with the global query and ask for a partial answer; then pass all partial answers into a second LLM call to produce one final synthesis.

  3. Weighted edges by sentiment. Some co-occurrences are positive (partnership), some negative (competition). Add signed edge weights; the graph’s community structure changes.

References#

  • Edge et al. 2024, Microsoft GraphRAG paper.

  • Traag, Waltman, van Eck 2019, From Louvain to Leiden.

  • Microsoft’s open-source graphrag repo for the production pipeline.

s.summary()
s.save()