GraphRAG Tutorial: From Documents to Knowledge Graph–Powered RAG

[Docs] -> [Chunker] -> (1) [Vector Index]
                 \-> (2) [Triple Extractor] -> [Graph DB] -> [Communities + Summaries]

Query -> [Entity Linking + Seed Nodes] -> [Neighborhood Expand] -> [Local + Global Context] -> [LLM Answer]

python -m venv .venv && source .venv/bin/activate
pip install networkx sentence-transformers faiss-cpu scikit-learn pydantic python-dotenv spacy
python -m spacy download en_core_web_sm

project/
  data/                 # raw docs
  build/
    graph.jsonl         # triples cache
    node_summaries.jsonl
  app/
    ingest.py
    extract_graph.py
    index.py
    retriever.py
    answer.py

# app/ingest.py
from pathlib import Path
import re, json
from typing import List, Dict

def simple_md_split(text: str, max_chars: int = 1800) -> List[str]:
    paras = [p.strip() for p in re.split(r"\n\n+", text) if p.strip()]
    chunks, buf = [], ""
    for p in paras:
        if len(buf) + len(p) + 2 > max_chars:
            if buf: chunks.append(buf); buf = ""
        buf = (buf + "\n\n" + p).strip()
    if buf: chunks.append(buf)
    return chunks

def load_docs(path="data") -> Dict[str, List[str]]:
    docs = {}
    for f in Path(path).glob("**/*.md"):
        text = f.read_text(encoding="utf-8")
        docs[f.stem] = simple_md_split(text)
    return docs

if __name__ == "__main__":
    docs = load_docs()
    Path("build").mkdir(exist_ok=True)
    with open("build/chunks.jsonl", "w", encoding="utf-8") as w:
        for doc_id, chunks in docs.items():
            for i, ch in enumerate(chunks):
                w.write(json.dumps({"doc_id": doc_id, "chunk_id": i, "text": ch})+"\n")

# app/llm.py
import os, json
from typing import List

# Implement this to call your LLM provider (OpenAI, Azure, Anthropic, local, etc.)
# It should return a parsed JSON string that fits the schema we request.

def call_llm(system: str, prompt: str) -> str:
    raise NotImplementedError("Plug in your LLM provider here.")

TRIPLE_SCHEMA = {
  "type": "object",
  "properties": {
    "triples": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "subject": {"type": "string"},
          "relation": {"type": "string"},
          "object": {"type": "string"},
          "evidence": {"type": "string"},
          "confidence": {"type": "number"}
        },
        "required": ["subject","relation","object","evidence","confidence"]
      }
    }
  },
  "required": ["triples"]
}

EXTRACT_SYSTEM = """
You extract knowledge graph triples from text. Output compact JSON only.
Entities should be canonical (merge aliases). Use relation verbs or nouns.
Include a short evidence quote from the text and a confidence in [0,1].
"""

EXTRACT_PROMPT_TMPL = """
Text:\n"""{text}"""\n
Respond with JSON per schema: {schema}
""".strip()

# app/extract_graph.py
import json
from pathlib import Path
from llm import call_llm, EXTRACT_SYSTEM, EXTRACT_PROMPT_TMPL, TRIPLE_SCHEMA

def extract_triples_for_chunks(chunks_path="build/chunks.jsonl", out_path="build/graph.jsonl"):
    with open(chunks_path, "r", encoding="utf-8") as r, open(out_path, "w", encoding="utf-8") as w:
        for line in r:
            rec = json.loads(line)
            prompt = EXTRACT_PROMPT_TMPL.format(text=rec["text"], schema=json.dumps(TRIPLE_SCHEMA))
            try:
                resp = call_llm(EXTRACT_SYSTEM, prompt)
                data = json.loads(resp)
                for t in data.get("triples", []):
                    t.update({"doc_id": rec["doc_id"], "chunk_id": rec["chunk_id"]})
                    w.write(json.dumps(t)+"\n")
            except Exception as e:
                # Optionally log and continue
                pass

if __name__ == "__main__":
    extract_triples_for_chunks()

# app/index.py
import json, networkx as nx
from collections import defaultdict

class GraphIndex:
    def __init__(self):
        self.G = nx.MultiDiGraph()
        self.node_text = defaultdict(list)

    def add_triple(self, s,r,o,evidence,meta):
        self.G.add_node(s)
        self.G.add_node(o)
        self.G.add_edge(s,o,relation=r,evidence=evidence,**meta)
        self.node_text[s].append(evidence)
        self.node_text[o].append(evidence)

    @classmethod
    def from_jsonl(cls, path="build/graph.jsonl"):
        gi = cls()
        with open(path,"r",encoding="utf-8") as f:
            for line in f:
                t = json.loads(line)
                gi.add_triple(t["subject"], t["relation"], t["object"], t["evidence"], {"doc_id":t["doc_id"],"chunk_id":t["chunk_id"]})
        return gi

if __name__ == "__main__":
    gi = GraphIndex.from_jsonl()
    print(gi.G.number_of_nodes(), gi.G.number_of_edges())

# app/summarize_graph.py
import json
import networkx as nx
from llm import call_llm
from collections import defaultdict

COMMUNITY_PROMPT = """
Summarize this set of related entities and relations in 5-8 bullet points.
Be factual; cite 2-4 key entity names.
Input triples:\n{triples}
"""

def louvain_communities(G):
    # lightweight fallback using connected components on undirected view
    # replace with a real community algorithm if desired
    return list(nx.connected_components(G.to_undirected()))

def summarize_communities(G, out_path="build/node_summaries.jsonl"):
    comms = louvain_communities(G)
    with open(out_path, "w", encoding="utf-8") as w:
        for i, nodes in enumerate(comms):
            sub = G.subgraph(nodes)
            triples = []
            for u,v,k,d in sub.edges(keys=True, data=True):
                triples.append(f"({u}) -[{d.get('relation','related')}]-> ({v})")
            prompt = COMMUNITY_PROMPT.format(triples="\n".join(triples[:60]))
            try:
                summary = call_llm("You write precise technical summaries.", prompt)
            except Exception:
                summary = "- Related entities: " + ", ".join(list(nodes)[:8])
            w.write(json.dumps({"community_id": i, "nodes": list(nodes), "summary": summary})+"\n")

# app/embed.py
import json
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss

class DualIndex:
    def __init__(self, model_name="all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.vec_dim = self.model.get_sentence_embedding_dimension()
        self.faiss_chunks = faiss.IndexFlatIP(self.vec_dim)
        self.faiss_nodes  = faiss.IndexFlatIP(self.vec_dim)
        self.chunk_meta = []
        self.node_meta  = []

    def _emb(self, texts):
        X = self.model.encode(texts, normalize_embeddings=True)
        return np.asarray(X, dtype="float32")

    def add_chunks(self, chunks_jsonl="build/chunks.jsonl"):
        texts, metas = [], []
        with open(chunks_jsonl,"r",encoding="utf-8") as f:
            for line in f:
                rec = json.loads(line)
                texts.append(rec["text"])
                metas.append({k:rec[k] for k in ("doc_id","chunk_id")})
        X = self._emb(texts)
        self.faiss_chunks.add(X)
        self.chunk_meta = metas

    def add_nodes(self, graph_index, min_snips=3):
        texts, metas = [], []
        for node, snips in graph_index.node_text.items():
            if len(snips) < min_snips: continue
            txt = f"Entity: {node}\nEvidence:\n- " + "\n- ".join(snips[:10])
            texts.append(txt)
            metas.append({"entity": node})
        X = self._emb(texts)
        self.faiss_nodes.add(X)
        self.node_meta = metas

# app/retriever.py
import json, re
import networkx as nx
import numpy as np
from typing import List, Dict, Any

MENTION_RX = re.compile(r"[A-Z][A-Za-z0-9_\-]{2,}")

def detect_entities(q: str) -> List[str]:
    # very naive: use spaCy NER in production
    return list(set(MENTION_RX.findall(q)))

def k_hop_neighborhood(G, seeds: List[str], k=2, max_nodes=60):
    visited = set(seeds)
    frontier = set(seeds)
    for _ in range(k):
        nxt = set()
        for u in frontier:
            for _, v in G.out_edges(u): nxt.add(v)
            for v, _ in G.in_edges(u): nxt.add(v)
        frontier = nxt - visited
        visited |= frontier
        if len(visited) > max_nodes: break
    return G.subgraph(visited)

def paths_as_text(SG):
    lines = []
    for u,v,k,d in SG.edges(keys=True, data=True):
        rel = d.get('relation','related')
        ev  = d.get('evidence','')
        lines.append(f"({u}) -[{rel}]-> ({v}); evidence: {ev[:120]}")
    return "\n".join(lines[:80])

def search_faiss(index, query_vec, topk=5):
    D, I = index.search(query_vec, topk)
    return I[0], D[0]

def embed_query(model, q):
    x = model.encode([q], normalize_embeddings=True).astype('float32')
    return x

# app/answer.py
import json
from llm import call_llm
from embed import DualIndex
from index import GraphIndex
from retriever import detect_entities, k_hop_neighborhood, paths_as_text, embed_query, search_faiss

SYNTH_PROMPT = """
You are a careful assistant. Use only the provided context. Cite entities or doc_ids.
Question: {q}

Local evidence (top chunks):
{local_blocks}

Graph paths (k-hop neighborhood):
{graph_paths}

Global summaries (communities):
{global_summaries}

Instructions:
- First list 3-6 key grounded facts with citations.
- Then produce a concise answer.
- Finally, show 1-3 critical paths as bullet points: (A) -[rel]-> (B) -[rel]-> (C).
"""

def build_and_answer(q:str):
    gi = GraphIndex.from_jsonl()
    di = DualIndex(); di.add_chunks(); di.add_nodes(gi)

    qv = embed_query(di.model, q)
    # Local text search
    I, D = search_faiss(di.faiss_chunks, qv, topk=6)
    local_blocks = []
    for idx in I:
        meta = di.chunk_meta[idx];
        # In production, also keep the text body for each chunk
        local_blocks.append(f"- doc={meta['doc_id']} chunk={meta['chunk_id']}")

    # Graph neighborhood
    seeds = detect_entities(q)
    if not seeds:
        # try to seed from best-matching nodes
        NI, _ = search_faiss(di.faiss_nodes, qv, topk=3)
        seeds = [di.node_meta[i]['entity'] for i in NI]
    SG = k_hop_neighborhood(gi.G, seeds, k=2, max_nodes=80)
    graph_paths = paths_as_text(SG)

    # Global summaries: collect any communities overlapping SG nodes
    # Here we just load precomputed summaries
    comm_summaries = []
    try:
        with open("build/node_summaries.jsonl","r",encoding="utf-8") as f:
            for line in f:
                rec = json.loads(line)
                if any(n in rec["nodes"] for n in SG.nodes):
                    comm_summaries.append(f"- c{rec['community_id']}: {rec['summary']}")
    except FileNotFoundError:
        pass

    prompt = SYNTH_PROMPT.format(q=q,
        local_blocks="\n".join(local_blocks),
        graph_paths=graph_paths,
        global_summaries="\n".join(comm_summaries[:4]))

    ans = call_llm("Grounded answering with rigorous citations.", prompt)
    return ans