Skip to content

graphrag.md

Latest commit

 

History

History
234 lines (183 loc) · 7.79 KB

graphrag.md

File metadata and controls

234 lines (183 loc) · 7.79 KB
Raw

rag-graphrag — public reference

GraphRAG layer for AgentContextOS, shipped in Step 2.9. Provides community detection, LLM-backed sub-graph summarisation, and a community-aware graph retriever that plugs into HybridRetriever as the graph source.

Overview

GraphRAG (Edge et al. 2024) is a retrieval style that uses community structure in a knowledge graph to drive retrieval — instead of asking the LLM "which entities are relevant to this query?", we precompute communities and their summaries at index time, then at query time pick the most relevant communities and follow their key entities into the graph.

The rag-graphrag package implements three pieces:

  1. Community detection — partition the graph into dense sub-graphs. Two algorithms ship; both produce identical Community objects.
  2. Sub-graph summarisation — turn each Community into a CommunitySummary via the LLM SPI.
  3. Graph-aware retrieval — score the query against community summaries, pick the top communities, use their key entities as graph-expansion seeds, and emit ChunkRef lists ready for RRF fusion.

Domain types and the GraphRAGError exception live in rag_core so downstream packages can consume them without depending on rag_graphrag.

Installation

# Default — pure-Python Louvain detector.
uv pip install rag-graphrag

# Optional — higher-modularity Leiden via igraph + leidenalg.
uv pip install 'rag-graphrag[leiden]'

Usage

Build the community index (offline / index time)

import networkx as nx
from rag_graphrag import (
    CommunitySummarizer,
    InMemoryCommunityStore,
    LouvainDetector,
    tenant_subgraph,
)

# Pull a tenant-scoped sub-graph out of the multi-tenant graph store.
sub = tenant_subgraph(networkx_store.graph, str(ctx.tenant_id))

# 1) Detect communities.
communities = await LouvainDetector().detect(ctx, sub)

# 2) Summarise each community via the LLM SPI.
store = InMemoryCommunityStore()
summariser = CommunitySummarizer(llm)
for c in communities:
    node_props = {n: dict(sub.nodes[n]) for n in c.node_ids if n in sub}
    summary = await summariser.summarize(ctx, c, node_properties=node_props)
    await store.upsert_community(ctx, c)
    await store.upsert_summary(ctx, summary)

Query (online / retrieval time)

from rag_graphrag import GraphRAGConfig, GraphRAGRetriever

retriever = GraphRAGRetriever(
    store,
    graph_backend,                   # any GraphRetrievalBackend
    config=GraphRAGConfig(community_top_k=3, graph_hops=2),
)
refs = await retriever.retrieve(ctx, "Who manages compliance at Acme?")

refs is a list[ChunkRef] — pass it through HybridRetriever (Step 2.5) for fusion with vector + keyword paths, then reranker + packer.

Wire into HybridRetriever

from rag_retrieval import HybridRetriever
from rag_graphrag import graphrag_adapter

retriever = HybridRetriever(
    vector_backend=vector_store,
    keyword_backend=keyword_store,
    graph_backend=graph_store,
    graph_adapter=graphrag_adapter,   # entity → chunk_ids unwrap
)

API

Detectors

Symbol Type Notes
LouvainDetector(min_community_size=2) class Default; pure-Python NetworkX.
LeidenDetector(min_community_size=2) class [leiden] extra; higher modularity.
tenant_subgraph(graph, tenant_id) -> nx.Graph helper Tenant-scoped subgraph copy.
CommunityDetector Protocol Implement to plug a custom detector.

Both detectors expose:

async def detect(
    ctx: RequestContext,
    graph: nx.Graph,
    *,
    resolution: float = 1.0,
    seed: int | None = 0,
) -> list[Community]

Summariser

CommunitySummarizer(llm: LLM, *, config: SummariserConfig | None = None)

async summarize(
    ctx: RequestContext,
    community: Community,
    *,
    node_properties: dict[str, dict[str, object]] | None = None,
) -> CommunitySummary

SummariserConfig(max_tokens=512, temperature=0.0, max_key_entities=8) tunes per-call cost and determinism.

Community store

CommunityStore is a Protocol; V1 ships InMemoryCommunityStore. Methods:

upsert_community(ctx, community) -> None
upsert_summary(ctx, summary) -> None
get_community(ctx, community_id) -> Community | None
get_summary(ctx, community_id) -> CommunitySummary | None
list_communities(ctx) -> list[Community]
list_summaries(ctx) -> list[CommunitySummary]
community_for_node(ctx, node_id) -> Community | None

Cross-tenant access raises GraphRAGError.

Retriever

GraphRAGRetriever(
    store: CommunityStore,
    graph_backend: GraphRetrievalBackend,
    *,
    config: GraphRAGConfig | None = None,
    adapter: GraphAdapter | None = None,
)

async retrieve(ctx, query: str) -> list[ChunkRef]
async score_communities(ctx, query: str) -> list[tuple[CommunitySummary, float]]

GraphRAGConfig(community_top_k=3, graph_hops=1, seeds_per_community=5, expand_limit=50) tunes the breadth of the search.

Adapter

async def graphrag_adapter(
    ctx: RequestContext,
    neighbors: list[NeighborResult],
) -> list[ChunkRef]

Reads neighbor.properties["chunk_ids"] and emits one ChunkRef per chunk. Nodes without the property are silently skipped (they're legitimate isolated-entity nodes).

CLI smoke

ragctl graphrag "Who manages compliance?"

Builds a 10-node toy graph, runs the full detect → summarise → retrieve flow against NoopLLM + NoopGraphStore, prints communities + chunks. No infra or LLM credentials required.

Internals

  • All public types are frozen Pydantic v2 models — create new instances rather than mutating.
  • Each pipeline stage opens its own OTel span via span_from_trace_context: graphrag.detect, graphrag.summarize, graphrag.retrieve. Attribute keys are namespaced rag.graphrag.* mirroring Step 2.7/2.8 conventions.
  • LouvainDetector collapses directed graphs to undirected before community detection. Multi-graphs are accepted but parallel edges collapse to one for modularity scoring.
  • graphrag_adapter does not deduplicate chunks across entities — that is rrf_fuse's job (Step 2.5).
  • _overlap_score is set-overlap with a 2× weight on key_entities vs. summary text. This is intentionally simple; a real deployment will swap it for embedder-cosine via a constructor argument in Phase 3.

Extension points

Want to Approach
Use Leiden instead of Louvain Construct with LeidenDetector(min_community_size=...) (requires [leiden] extra).
Plug a custom detector Implement the CommunityDetector Protocol; pass it where you'd pass LouvainDetector.
Use a real LLM backend Pass any LLM SPI implementation (OpenAI / Anthropic / vLLM) into CommunitySummarizer.
Replace the community store Implement the CommunityStore Protocol — V1 ships in-memory only; persistent stores arrive with Phase 3 / Step 3.5.
Use embedder cosine for community scoring Override GraphRAGRetriever.score_communities in a subclass or pass a custom adapter (the public scorer signature will become pluggable in Phase 3 alongside the corpus router).

What's deliberately deferred

  • Persistent community store — Phase 3 / Step 3.5 corpus router decides where it lives.
  • Hierarchical communities — LeidenDetector supports hierarchical output via repeated calls with reduced graphs; V1 returns flat partitions to keep the API surface narrow.
  • Embedder-cosine community scoring — V1 uses token overlap; cosine path arrives once Phase 3 wires per-tenant embedder selection.
  • LLM-driven contradiction reconciliation across communities — Step 5.x hallucination guard.
  • /v1/communities REST endpoint — Phase 3 gateway.