Graph RAG: Knowledge-Graph-Based Retrieval

Vector-similarity retrieval, the mechanism behind naive RAG and most of the architectures covered so far, is fundamentally good at finding text that's semantically similar to a query. But it's not particularly good at answering questions that hinge on RELATIONSHIPS between entities — 'which suppliers does Company X share with Company Y?' isn't really a similarity-matching question, it's a graph-traversal question.

Graph RAG addresses this by converting retrieved content into a knowledge graph that explicitly captures entities (people, organizations, products, concepts) and the relationships between them (works-for, supplies-to, depends-on, is-a), rather than leaving everything as unstructured chunks of prose. Once information is represented this way, retrieval can traverse the graph's actual structure — following edges between related entities — rather than relying purely on semantic similarity between the query and stored text.

This structured relational context gets passed to the LLM alongside the raw text it was extracted from, which meaningfully enhances the model's reasoning on relationship-heavy or multi-entity questions: instead of the LLM having to infer connections purely from whatever prose happened to be retrieved, it's handed the relationships explicitly, already extracted and structured. This is particularly valuable for domains that are naturally graph-shaped to begin with — organizational structures, supply chains, citation networks, biomedical entity relationships, or any dataset where 'how are these things connected' is a more common question than 'find me text about this one thing.'

The trade-off is real upfront and ongoing cost: building and maintaining a knowledge graph from unstructured source documents requires an entity-and-relationship extraction step (often itself LLM-powered) that plain vector-store RAG doesn't need at all, plus graph database infrastructure and keeping that graph in sync as source documents change over time. For datasets and query patterns that are genuinely relationship-heavy, this cost is well worth paying; for simple fact-lookup queries, it's substantial overhead for little additional benefit over naive RAG.

# Simplified: extract entities/relationships, then answer a
# relationship-style query by traversing the resulting graph.
import networkx as nx
from openai import OpenAI

client = OpenAI()

def extract_entities_and_relations(text: str) -> list[tuple[str, str, str]]:
    """Returns (entity_a, relation, entity_b) triples — in production
    this is usually its own dedicated LLM extraction prompt."""
    resp = client.chat.completions.create(model="gpt-4.1", messages=[
        {"role": "user", "content":
            f"Extract (entity, relation, entity) triples from: {text}\n"
            f"Return one triple per line as entity_a | relation | entity_b"}
    ])
    lines = resp.choices[0].message.content.strip().split("\n")
    return [tuple(p.strip() for p in line.split("|")) for line in lines if "|" in line]

def build_graph(documents: list[str]) -> nx.DiGraph:
    graph = nx.DiGraph()
    for doc in documents:
        for a, relation, b in extract_entities_and_relations(doc):
            graph.add_edge(a, b, relation=relation)
    return graph

def graph_rag_query(graph: nx.DiGraph, entity: str, hops: int = 2) -> str:
    # Traverse relationships within `hops` steps of the entity —
    # something plain vector similarity search can't do at all
    subgraph_nodes = nx.single_source_shortest_path_length(graph, entity, cutoff=hops)
    edges = [(u, graph[u][v]["relation"], v) for u, v in graph.edges(subgraph_nodes)]
    return "\n".join(f"{u} --{r}--> {v}" for u, r, v in edges)