RAG with Knowledge Graphs

Vector search finds documents that sound similar. A knowledge graph finds the cousin of the founder of the company that acquired the supplier of your competitor. One of these is more useful at the deposition.

RAG, Multi-Hop-Traversing AI Agent

What does that look like concretely? Figure 35.3.2 shows a small worked example around Albert Einstein, with each entity as a coloured node and each labelled edge as a relationship the retriever can traverse.

Figure 35.3.2a: A small knowledge graph centred on Albert Einstein, showing how entities (cylinders, colour-coded by type) are connected by labelled relationships (named edges). A question like "who influenced the developer of General Relativity?" becomes a one-hop traversal: General Relativity ← developed ← Einstein → influenced_by → Newton. This is the structure RAG systems traverse when they have access to a knowledge graph instead of (or in addition to) raw text.

35.3.1 Knowledge Graph Fundamentals

A knowledge graph (KG) is a structured representation of real-world entities and the relationships between them. Unlike documents stored as unstructured text, knowledge graphs encode information as a network of nodes (entities) connected by typed edges (relations). This structured representation complements the dense vector representations used in standard semantic search. This structure enables precise queries about relationships, paths, and patterns that would require complex reasoning over unstructured text.

35.3.1.1 Entities, Relations, and Triples

The fundamental unit of a knowledge graph is the triple, expressed as (subject, predicate, object) or equivalently (head, relation, tail). For example, (Albert Einstein, bornIn, Ulm) encodes a single fact. A knowledge graph is simply a collection of such triples, forming a directed labeled graph where entities are nodes and relations are edge labels.

The diagram below shows how a knowledge graph encodes entities and relationships.

Figure 35.3.4: A knowledge graph encodes entities as nodes and relationships as labeled directed edges, enabling structured queries and multi-hop reasoning.

35.3.1.2 RDF vs. Property Graphs

Two main formalisms exist for representing knowledge graphs. RDF (Resource Description Framework) is a W3C standard that represents all knowledge as triples and uses URIs to identify entities and relations. RDF is queried using SPARQL and is the foundation of the Semantic Web (Wikidata, DBpedia). Property graphs allow nodes and edges to carry arbitrary key-value properties, offering greater flexibility. Property graphs are queried using Cypher (Neo4j) or Gremlin and are more commonly used in industry applications.

35.3.1.3 Ontology Design: Typed Entity and Relation Schemas

Before extracting a single triple, consider what happens when you do not constrain the extractor at all. Schema-free Open IE lets the LLM invent its own vocabulary, so the same fact arrives three different ways across three documents: (Ada, works_at, Acme), (Ada, employed_by, Acme), and (Ada, is_employee_of, Acme). To the graph these are three unrelated edges. A Cypher query that asks "who works at Acme?" matches the first relation and silently misses the other two, and entity nodes fragment the same way when "Acme", "Acme Corp", and "Acme Inc." all become distinct nodes. The graph is technically populated and practically unqueryable.

An ontology fixes this by pinning down the vocabulary before extraction. For our purposes an ontology has three parts. First, a controlled set of entity classes organized in an IS-A hierarchy (an Organization is a kind of Thing; a Company is a kind of Organization), so a query over Organization also reaches every Company. Second, a controlled set of typed relations, each with a domain (the entity class allowed as subject) and a range (the class allowed as object): works_at has domain Person and range Organization, so a person can work at a company but a company cannot "work at" a person. Third, cardinality and other constraints (a person has exactly one born_in country, a company may have many employees) that catch malformed facts. This domain-and-range discipline is the lineage of RDFS and OWL, the W3C ontology languages that formalized class hierarchies and typed properties for the Semantic Web; in practice you do not need a full OWL reasoner to get most of the value, and lightweight encodings such as Pydantic enums or the rdflib and owlready2 libraries are enough to declare and enforce a working schema.

Code Fragment 35.3.1c declares a tiny ontology with Pydantic and enforces it. The EntityType and RelationType enums are the controlled vocabularies, the ALLOWED_RELATIONS set encodes each valid (subject_type, relation, object_type) domain/range triple, and validate_triple rejects anything outside that set.

# A minimal KG ontology: controlled entity classes, controlled relations, and a
# domain/range table saying which (subject_type, relation, object_type) triples are legal.
# validate_triple() is the gate every extracted triple must pass before it enters the graph.
from enum import Enum


class EntityType(str, Enum):
    PERSON = "Person"
    ORGANIZATION = "Organization"
    COUNTRY = "Country"


class RelationType(str, Enum):
    WORKS_AT = "works_at"
    BORN_IN = "born_in"
    HEADQUARTERED_IN = "headquartered_in"

# Each tuple is one (domain, relation, range) rule: the only legal shapes for a triple.
ALLOWED_RELATIONS = {
    (EntityType.PERSON, RelationType.WORKS_AT, EntityType.ORGANIZATION),
    (EntityType.PERSON, RelationType.BORN_IN, EntityType.COUNTRY),
    (EntityType.ORGANIZATION, RelationType.HEADQUARTERED_IN, EntityType.COUNTRY),
}


def validate_triple(subj_type: EntityType, rel: RelationType, obj_type: EntityType) -> bool:
    """Return True only if (subj_type, rel, obj_type) is a declared domain/range rule."""
    return (subj_type, rel, obj_type) in ALLOWED_RELATIONS


# Accepted: a Person works at an Organization.
print(validate_triple(EntityType.PERSON, RelationType.WORKS_AT, EntityType.ORGANIZATION))
# Rejected: a Country cannot "work at" a Person (domain/range violation).
print(validate_triple(EntityType.COUNTRY, RelationType.WORKS_AT, EntityType.PERSON))

This ontology plugs directly into the KG-construction pipeline that follows. The EntityType and RelationType enums become the schema you hand to the LLM extractor: instead of free-form Open IE, you ask for constrained, JSON-schema-shaped output whose head type, relation, and tail type are drawn only from these enums, the same role played by the optional entity_types hint passed to extract_triples in Code Fragment 35.3.1b below. The validate_triple gate then runs on every extracted triple before the MERGE step in the Neo4j loader of Code Fragment 35.3.2b, so only ontology-conformant facts ever reach the graph and "works_at" stays a single, queryable relation.

35.3.2 Building Knowledge Graphs with LLMs

Traditionally, knowledge graph construction required extensive manual curation or specialized NLP pipelines for entity extraction and relation classification. LLMs have dramatically simplified this process by extracting structured triples directly from unstructured text through carefully designed prompts.

# Knowledge-graph triple extraction: ask an LLM to turn unstructured text into
# (head, relation, tail) tuples that can be loaded into Neo4j or any KG store.
import json

from openai import OpenAI

client = OpenAI()

TRIPLE_SYSTEM_PROMPT = """Extract knowledge graph triples from the text.
Return a JSON object with a "triples" array of objects with keys: head, relation, tail.
Use concise, canonical entity names. Use lowercase relation names with underscores.
Example:
{"triples": [
  {"head": "OpenAI", "relation": "developed", "tail": "GPT-4"},
  {"head": "GPT-4", "relation": "is_a", "tail": "language model"}
]}"""

def extract_triples(text: str, entity_types: list[str] | None = None) -> list[dict]:
    """Extract KG triples from `text`, optionally restricted to given entity types."""
    # Optional hint nudges the model toward the entity types we care about.
    type_hint = (
        f"\nFocus on these entity types: {', '.join(entity_types)}" if entity_types else ""
    )
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": TRIPLE_SYSTEM_PROMPT + type_hint},
            {"role": "user", "content": text},
        ],
        response_format={"type": "json_object"},
        temperature=0.0,  # Deterministic for reproducible KG ingest
    )
    result = json.loads(response.choices[0].message.content)
    return result.get("triples", [])

35.3.2.1 Storing Triples in Neo4j

This snippet inserts knowledge-graph triples into a Neo4j database using Cypher queries.

# Thin Neo4j wrapper: store extracted (head, relation, tail) triples as nodes
# and relationships, and expose neighborhood and path queries for retrieval.
from neo4j import GraphDatabase
class KnowledgeGraphStore:
    """Store and query knowledge graph triples in Neo4j."""
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
    def add_triples(self, triples):
        """Insert triples as nodes and relationships."""
        with self.driver.session() as session:
            for triple in triples:
                session.run("""
                    MERGE (h:Entity {name: $head})
                    MERGE (t:Entity {name: $tail})
                    MERGE (h)-[r:RELATES {type: $relation}]->(t)
                    """, head=triple["head"],
                    tail=triple["tail"],
                    relation=triple["relation"])
    def find_neighbors(self, entity, max_hops=2):
        """Find all entities within N hops of the given entity."""
        with self.driver.session() as session:
            result = session.run(f"""
                MATCH path = (e:Entity {{name: $entity}})
                -[*1..{max_hops}]-(neighbor)
                RETURN path
                LIMIT 50
                """, entity=entity)
            return [record["path"] for record in result]
            def query_relationship(self, head, tail):
                """Find relationship paths between two entities."""
                with self.driver.session() as session:
                    result = session.run("""
                        MATCH path = shortestPath(
                        (h:Entity {name: $head})-[*]-(t:Entity {name: $tail})
                        )
                        RETURN path
                        """, head=head, tail=tail)
                    return [record["path"] for record in result]

35.3.3 Graph Embeddings

Graph embeddings map entities and relations from a knowledge graph into continuous vector spaces, enabling similarity-based retrieval, link prediction (predicting missing triples), and integration with dense retrieval systems. The key challenge is preserving the structural properties of the graph in the embedding space.

35.3.3.1 TransE

TransE (Bordes et al., 2013) is the foundational graph embedding model. It represents relations as translations in embedding space: if (h, r, t) is a true triple, then the embedding of the head entity plus the embedding of the relation should approximately equal the embedding of the tail entity: h + r ≈ t. The model is trained by minimizing a margin-based loss that pushes correct triples closer than corrupted triples.

35.3.3.2 DistMult

DistMult (Yang et al., 2015) models relations as diagonal matrices (represented as vectors) and scores triples using a bilinear product: score(h, r, t) = h · diag(r) · t. This is equivalent to an element-wise product followed by a sum. DistMult is simple and effective but cannot model asymmetric relations (it gives the same score to (h, r, t) and (t, r, h)), which limits its applicability.

import torch
import torch.nn as nn
class TransE(nn.Module):
    """TransE knowledge graph embedding model."""
    def __init__(self, num_entities, num_relations, dim=128):
        super().__init__()
        self.entity_emb = nn.Embedding(num_entities, dim)
        self.relation_emb = nn.Embedding(num_relations, dim)
        # Initialize with uniform distribution
        nn.init.uniform_(self.entity_emb.weight, -6/dim**0.5, 6/dim**0.5)
        nn.init.uniform_(self.relation_emb.weight, -6/dim**0.5, 6/dim**0.5)
    def forward(self, heads, relations, tails):
        """Score triples: lower distance = more plausible."""
        h = self.entity_emb(heads)
        r = self.relation_emb(relations)
        t = self.entity_emb(tails)
        # TransE scoring: ||h + r - t||
        score = torch.norm(h + r - t, p=2, dim=-1)
        return score
    def margin_loss(self, pos_score, neg_score, margin=1.0):
        """Margin-based ranking loss."""
        return torch.relu(margin + pos_score - neg_score).mean()

35.3.4 Cypher-Based Multi-Hop Retrieval

Once entities live in a property graph, Cypher queries enable precise, multi-hop retrieval that vector search cannot replicate. A two-hop query from a drug, to the gene it targets, to the disease that gene is associated with, is a single MATCH path. The retriever side of a production system usually wraps Cypher behind a natural-language translator, so the agent or LLM only sees a tool with a query parameter and the right answers come back grounded in graph paths.

from neo4j_graphrag.retrievers import CypherRetriever
retriever = CypherRetriever(
    driver=driver,
    llm=llm,
    neo4j_schema="""
    Node labels: Drug, Disease, SideEffect, Gene
    Relationships: TREATS, CAUSES, INHIBITS, TARGETS
    """,
    )
# The retriever translates natural language to Cypher.
# Behind the scenes it generates Cypher like:
# MATCH (d:Drug {name: 'Metformin'})-[:TARGETS]->(g:Gene)
#       -[:ASSOCIATED_WITH]->(dis:Disease)
# RETURN d.name, g.name, dis.name
result = retriever.search(
    query_text="What genes does metformin target, and what diseases "
    "are those genes associated with?"
    )
for item in result.items:
    print(item.content)

The same shape generalizes to RDF + SPARQL on Wikidata or DBpedia, to Amazon Neptune's openCypher, and to TigerGraph's GSQL. The pattern is always the same: declare the schema, let an LLM-backed retriever generate the path query, and return the resulting subgraph as retrieved evidence. This is the foundation that any community-summarization or graph-summarization technique (covered in Section 35.4) sits on top of.

35.3.5 Hybrid KG + Vector Retrieval

The most effective production systems combine knowledge graph traversal with vector search. The knowledge graph provides structured multi-hop reasoning and entity-aware retrieval, while vector search provides semantic similarity for unstructured content that the KG does not cover.

# Hybrid KG + vector retrieval: ask the KG for structured multi-hop neighborhoods
# around the query's entities, then stitch them with semantic search hits.
def hybrid_kg_vector_retrieve(query: str, kg_store, vector_collection, k: int = 5) -> str:
    """Merge KG traversal (structured) with vector search (semantic) into one context."""
    # 1) Pull entity mentions out of the query so we know where to anchor the KG walk.
    entities = extract_entities_from_query(query)

    # 2) KG traversal: 2-hop neighborhoods around each query entity.
    kg_context: list[str] = []
    for entity in entities:
        neighbors = kg_store.find_neighbors(entity, max_hops=2)
        kg_context.extend(format_paths(neighbors))

    # 3) Vector search picks up unstructured text the KG does not encode.
    vector_results = vector_collection.query(query_texts=[query], n_results=k)
    vector_context = vector_results["documents"][0]

    # 4) Concatenate the two into a single prompt-ready string.
    return (
        "Structured Knowledge (from Knowledge Graph):\n"
        + "\n".join(kg_context)
        + "\n\nRelevant Documents (from semantic search):\n"
        + "\n".join(vector_context)
    )

35.3.6 When to Use Knowledge Graph RAG