Standard RAG is good at finding relevant passages. It's terrible at following a chain of reasoning across multiple documents. Ask a vector search system "Which pharmaceutical compounds developed by researchers at MIT have been tested in clinical trials funded by organizations that also sponsor climate research?" and watch it confidently return the most semantically similar chunk — which almost certainly won't actually answer your question. The retrieval was fine. The reasoning capability was never there to begin with.
This is the fundamental limitation of pure vector search: it retrieves by similarity, not by relationship. When your question requires traversing a web of connected facts — entity A relates to entity B which connects to entity C — cosine similarity between embeddings cannot save you. You need something that understands structure.
Knowledge Graph-Augmented RAG (KG-RAG) solves this by combining the fuzzy semantic matching power of vector search with the precise relational traversal of a graph database. By the end of this lesson, you'll have built a working system that extracts entities and relationships from documents, stores them in a queryable graph, and uses multi-hop graph traversal to augment the context passed to an LLM — enabling the kind of connected reasoning that pure RAG systems simply cannot do.
What you'll learn:
You should be comfortable with:
You'll need running instances of Neo4j (community edition works) and a vector store. We'll use Qdrant for vector storage, but the patterns apply to Pinecone, Weaviate, or pgvector equally well.
Before building anything, let's be precise about the failure mode. Consider a corpus of biomedical research papers. You embed everything, chunk it, stuff it into Qdrant, and fire queries at it.
For a question like "What are the side effects of metformin?" — great. Vector search finds chunks mentioning metformin and side effects. The LLM synthesizes a competent answer.
Now try: "Are there any researchers who collaborated with the team that published the ACCORD trial, and did any of them later work on SGLT2 inhibitor studies?"
The problem isn't that the answer isn't in your corpus. It might be. The problem is that answering it requires:
No individual chunk contains all of this. Vector similarity will happily return chunks about the ACCORD trial and separate chunks about SGLT2 inhibitors, but it has no mechanism to follow the chain of people connecting them. You'd need to hope the LLM hallucinates a bridge — which it often will, confidently and incorrectly.
A knowledge graph makes these connections explicit and queryable. Instead of hoping the right context lands in your retrieval window by semantic proximity, you traverse edges: Person-[:AUTHORED]->Paper-[:CITES]->Paper-[:STUDIES]->Drug.
The system we're building has three major components:
Ingestion Pipeline: Extracts entities and relationships from documents and populates both a vector store (for semantic search) and a knowledge graph (for relational traversal).
Hybrid Retrieval Engine: Takes a query, uses vector search to identify seed entities, then traverses the graph outward to gather relational context unavailable to pure vector retrieval.
Context Assembly & Generation: Combines vector-retrieved chunks with graph-traversed paths into a structured prompt that gives the LLM the connected context it needs.
Install the core dependencies:
pip install openai neo4j qdrant-client sentence-transformers spacy tiktoken pydantic tenacity
python -m spacy download en_core_web_trf
Start Neo4j and Qdrant locally with Docker:
# docker-compose.yml
version: '3.8'
services:
neo4j:
image: neo4j:5.15-community
ports:
- "7474:7474" # Browser UI
- "7687:7687" # Bolt protocol
environment:
- NEO4J_AUTH=neo4j/your_password
- NEO4J_PLUGINS=["apoc"]
volumes:
- neo4j_data:/data
qdrant:
image: qdrant/qdrant:v1.7.4
ports:
- "6333:6333"
volumes:
- qdrant_data:/qdrant/storage
volumes:
neo4j_data:
qdrant_data:
docker-compose up -d
Now set up your configuration:
# config.py
import os
from dataclasses import dataclass
@dataclass
class Config:
openai_api_key: str = os.getenv("OPENAI_API_KEY")
neo4j_uri: str = "bolt://localhost:7687"
neo4j_user: str = "neo4j"
neo4j_password: str = "your_password"
qdrant_host: str = "localhost"
qdrant_port: int = 6333
collection_name: str = "biomedical_papers"
embedding_model: str = "text-embedding-3-small"
llm_model: str = "gpt-4o"
embedding_dim: int = 1536
entity_extraction_model: str = "gpt-4o-mini" # Cheaper for extraction
config = Config()
Using gpt-4o-mini for entity extraction is a deliberate cost decision. You're going to run this against every chunk in your corpus during ingestion. The extraction task is structured and well-defined — it doesn't need the most powerful model. Save your GPT-4o budget for final answer generation.
This is where the magic starts — and where most tutorials gloss over the hard parts. Entity extraction from scientific text is not trivial. Named entity recognition (NER) with spaCy will catch person names and organizations, but it won't understand domain-specific entities like drug compounds, clinical trial identifiers, or research methodologies.
We'll use a hybrid approach: spaCy for coarse entity detection as a pre-filter, and an LLM with a structured extraction prompt for precise entity and relationship extraction.
Before writing any extraction code, define what your graph will look like. This is your most important design decision — get it wrong and you'll end up with an unmaintainable hairball.
# schema.py
from pydantic import BaseModel, Field
from typing import Optional
from enum import Enum
class EntityType(str, Enum):
PERSON = "Person"
ORGANIZATION = "Organization"
DRUG = "Drug"
DISEASE = "Disease"
CLINICAL_TRIAL = "ClinicalTrial"
RESEARCH_PAPER = "ResearchPaper"
INSTITUTION = "Institution"
BIOMARKER = "Biomarker"
class RelationshipType(str, Enum):
AUTHORED = "AUTHORED"
AFFILIATED_WITH = "AFFILIATED_WITH"
STUDIED = "STUDIED"
TREATS = "TREATS"
CAUSES = "CAUSES"
PARTICIPATED_IN = "PARTICIPATED_IN"
FUNDED_BY = "FUNDED_BY"
CITES = "CITES"
COLLABORATED_WITH = "COLLABORATED_WITH"
REGULATES = "REGULATES"
class Entity(BaseModel):
name: str
entity_type: EntityType
description: Optional[str] = None
canonical_name: Optional[str] = None # For disambiguation
class Relationship(BaseModel):
source_entity: str
source_type: EntityType
relationship_type: RelationshipType
target_entity: str
target_type: EntityType
context: Optional[str] = Field(None, description="Brief context for this relationship")
confidence: float = Field(default=1.0, ge=0.0, le=1.0)
class ExtractionResult(BaseModel):
entities: list[Entity]
relationships: list[Relationship]
chunk_id: str
The canonical_name field on Entity is critical and often neglected. Your corpus will mention "Dr. Sarah Chen," "S. Chen," "Chen et al.," and "Sarah Chen, PhD" — all the same person. Without canonicalization, you build a graph with four disconnected nodes where there should be one. We'll handle this with a disambiguation pass.
# extraction.py
import json
from openai import AsyncOpenAI
from tenacity import retry, stop_after_attempt, wait_exponential
from schema import ExtractionResult, EntityType, RelationshipType
client = AsyncOpenAI(api_key=config.openai_api_key)
EXTRACTION_SYSTEM_PROMPT = """You are a biomedical knowledge graph extraction system.
Your task is to extract entities and relationships from scientific text with high precision.
ENTITY TYPES you should extract:
- Person: Researchers, authors, clinicians
- Organization: Companies, hospitals, funding agencies
- Drug: Pharmaceutical compounds, treatments, interventions
- Disease: Medical conditions, syndromes, disorders
- ClinicalTrial: Named trials (e.g., "ACCORD trial", "UKPDS")
- ResearchPaper: Referenced papers (use title or identifier)
- Institution: Universities, research institutes
- Biomarker: Biological markers, lab values
RELATIONSHIP TYPES:
- AUTHORED: Person -> ResearchPaper
- AFFILIATED_WITH: Person -> Institution/Organization
- STUDIED: ResearchPaper/ClinicalTrial -> Drug/Disease/Biomarker
- TREATS: Drug -> Disease
- CAUSES: Drug/Disease -> Disease/Biomarker (side effects, complications)
- PARTICIPATED_IN: Person/Organization -> ClinicalTrial
- FUNDED_BY: ResearchPaper/ClinicalTrial -> Organization
- CITES: ResearchPaper -> ResearchPaper
- COLLABORATED_WITH: Person -> Person (inferred from co-authorship)
- REGULATES: Biomarker/Drug -> Biomarker/Drug
RULES:
1. Only extract relationships explicitly stated or strongly implied in the text
2. Normalize entity names to their most complete form
3. Assign confidence scores: 1.0 for explicit statements, 0.7-0.9 for strong implications
4. For COLLABORATED_WITH, infer from co-authorship when multiple authors are listed
Return valid JSON matching the ExtractionResult schema."""
EXTRACTION_USER_PROMPT = """Extract all entities and relationships from this text chunk.
Chunk ID: {chunk_id}
TEXT:
{text}
Return JSON with this structure:
{{
"entities": [
{{"name": "...", "entity_type": "...", "description": "...", "canonical_name": "..."}}
],
"relationships": [
{{
"source_entity": "...", "source_type": "...",
"relationship_type": "...",
"target_entity": "...", "target_type": "...",
"context": "...", "confidence": 0.0
}}
],
"chunk_id": "{chunk_id}"
}}"""
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
async def extract_entities_and_relationships(
text: str,
chunk_id: str
) -> ExtractionResult:
response = await client.chat.completions.create(
model=config.entity_extraction_model,
messages=[
{"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
{"role": "user", "content": EXTRACTION_USER_PROMPT.format(
chunk_id=chunk_id,
text=text
)}
],
response_format={"type": "json_object"},
temperature=0.1, # Low temperature for consistent extraction
max_tokens=2000
)
raw = json.loads(response.choices[0].message.content)
return ExtractionResult(**raw)
Warning:
temperature=0.1is intentional here. Higher temperatures introduce variability in entity names — "Metformin" vs "metformin" vs "METFORMIN" — which destroys your graph structure when the same entity gets multiple nodes. You want the model to be boring and consistent.
After extraction, you need to resolve variations of the same entity before writing to the graph. Here's a practical approach using embedding similarity:
# disambiguation.py
import numpy as np
from openai import AsyncOpenAI
from schema import Entity
client = AsyncOpenAI(api_key=config.openai_api_key)
async def get_embedding(text: str) -> list[float]:
response = await client.embeddings.create(
model=config.embedding_model,
input=text
)
return response.data[0].embedding
class EntityDisambiguator:
def __init__(self, similarity_threshold: float = 0.92):
self.known_entities: dict[str, tuple[Entity, list[float]]] = {}
self.threshold = similarity_threshold
async def resolve(self, entity: Entity) -> Entity:
"""
Find the canonical form of an entity, or register it as new.
Returns the canonical entity.
"""
if not self.known_entities:
embedding = await get_embedding(
f"{entity.entity_type}: {entity.name}"
)
self.known_entities[entity.name] = (entity, embedding)
return entity
# Embed the new entity
new_embedding = await get_embedding(
f"{entity.entity_type}: {entity.name}"
)
best_match = None
best_similarity = 0.0
for known_name, (known_entity, known_embedding) in self.known_entities.items():
# Only compare same entity types
if known_entity.entity_type != entity.entity_type:
continue
similarity = cosine_similarity(new_embedding, known_embedding)
if similarity > best_similarity:
best_similarity = similarity
best_match = known_entity
if best_match and best_similarity >= self.threshold:
# Return the canonical entity (the first one we saw)
return best_match
else:
# New entity - register it
self.known_entities[entity.name] = (entity, new_embedding)
return entity
def cosine_similarity(a: list[float], b: list[float]) -> float:
a_arr = np.array(a)
b_arr = np.array(b)
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
The 0.92 threshold is a starting point, not a rule. In a medical corpus, "ACE inhibitors" and "ACE inhibitor therapy" should resolve to the same entity (very high similarity), but "metformin" and "metformin hydrochloride" might need expert judgment. Tune this threshold empirically on a sample of your data.
Now we write to the graph. The design of your Cypher queries here has enormous performance implications at scale.
# graph_store.py
from neo4j import AsyncGraphDatabase
from schema import Entity, Relationship, ExtractionResult
class KnowledgeGraphStore:
def __init__(self):
self.driver = AsyncGraphDatabase.driver(
config.neo4j_uri,
auth=(config.neo4j_user, config.neo4j_password)
)
async def initialize_schema(self):
"""Create indexes for performance. Run once."""
async with self.driver.session() as session:
# Uniqueness constraints prevent duplicate nodes
await session.run(
"CREATE CONSTRAINT entity_name_type IF NOT EXISTS "
"FOR (e:Entity) REQUIRE (e.name, e.entity_type) IS UNIQUE"
)
# Full-text index for entity name search
await session.run(
"CREATE FULLTEXT INDEX entity_name_index IF NOT EXISTS "
"FOR (e:Entity) ON EACH [e.name, e.canonical_name]"
)
# Composite index for relationship traversal performance
await session.run(
"CREATE INDEX chunk_entity_idx IF NOT EXISTS "
"FOR ()-[r:MENTIONED_IN]-() ON (r.chunk_id)"
)
async def upsert_entity(self, entity: Entity) -> None:
"""
MERGE ensures we don't create duplicate nodes.
ON CREATE sets initial properties; ON MATCH updates them.
"""
async with self.driver.session() as session:
await session.run(
"""
MERGE (e:Entity {name: $name, entity_type: $entity_type})
ON CREATE SET
e.description = $description,
e.canonical_name = $canonical_name,
e.created_at = timestamp(),
e.mention_count = 1
ON MATCH SET
e.mention_count = e.mention_count + 1,
e.description = CASE
WHEN e.description IS NULL THEN $description
ELSE e.description
END
WITH e
CALL apoc.create.addLabels(e, [$entity_type]) YIELD node
RETURN node
""",
name=entity.name,
entity_type=entity.entity_type.value,
description=entity.description or "",
canonical_name=entity.canonical_name or entity.name
)
async def upsert_relationship(
self,
relationship: Relationship,
chunk_id: str
) -> None:
"""
Create a typed relationship between two entities.
We also store the chunk_id so we can trace back to source text.
"""
async with self.driver.session() as session:
# Dynamic relationship types require APOC in Neo4j
await session.run(
"""
MATCH (source:Entity {name: $source_name, entity_type: $source_type})
MATCH (target:Entity {name: $target_name, entity_type: $target_type})
CALL apoc.merge.relationship(
source,
$rel_type,
{source_chunk: $chunk_id},
{context: $context, confidence: $confidence, created_at: timestamp()},
target,
{}
) YIELD rel
RETURN rel
""",
source_name=relationship.source_entity,
source_type=relationship.source_type.value,
target_name=relationship.target_entity,
target_type=relationship.target_type.value,
rel_type=relationship.relationship_type.value,
chunk_id=chunk_id,
context=relationship.context or "",
confidence=relationship.confidence
)
async def store_extraction_result(
self,
result: ExtractionResult,
disambiguator: 'EntityDisambiguator'
) -> None:
"""Process a full extraction result."""
# First pass: upsert all entities
canonical_map = {}
for entity in result.entities:
canonical = await disambiguator.resolve(entity)
canonical_map[entity.name] = canonical
await self.upsert_entity(canonical)
# Second pass: upsert relationships using canonical names
for rel in result.relationships:
canonical_source = canonical_map.get(rel.source_entity)
canonical_target = canonical_map.get(rel.target_entity)
if not canonical_source or not canonical_target:
continue # Skip if entities weren't extracted
# Remap to canonical names
canonical_rel = rel.model_copy(update={
"source_entity": canonical_source.name,
"target_entity": canonical_target.name
})
await self.upsert_relationship(canonical_rel, result.chunk_id)
async def close(self):
await self.driver.close()
Tip: The
MERGEpattern in Cypher is your best friend and a common performance trap. AlwaysMERGEon a small, indexed set of properties. If youMERGEon properties that aren't indexed, Neo4j does a full node scan. With a large graph this becomes catastrophically slow. Our schema initialization creates the necessary indexes — don't skip it.
This is the core intellectual contribution of KG-RAG. We're building a retrieval system that uses vector search to find entry points into the graph, then traverses edges to collect context that would be invisible to pure semantic search.
# retrieval.py
from dataclasses import dataclass
from qdrant_client import AsyncQdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from neo4j import AsyncGraphDatabase
import uuid
@dataclass
class GraphPath:
"""Represents a traversed path through the knowledge graph."""
path_elements: list[str] # Sequence of nodes and relationships
start_entity: str
end_entity: str
hop_count: int
relevance_score: float
source_chunks: list[str] # Chunk IDs where relationships were found
@dataclass
class HybridRetrievalResult:
vector_chunks: list[dict] # Standard vector search results
graph_paths: list[GraphPath] # Multi-hop graph traversal results
seed_entities: list[str] # Entities identified in the query
query: str
class HybridRetriever:
def __init__(self):
self.qdrant = AsyncQdrantClient(
host=config.qdrant_host,
port=config.qdrant_port
)
self.neo4j = AsyncGraphDatabase.driver(
config.neo4j_uri,
auth=(config.neo4j_user, config.neo4j_password)
)
self.openai = AsyncOpenAI(api_key=config.openai_api_key)
async def extract_query_entities(self, query: str) -> list[str]:
"""
Extract entity mentions from the query for use as graph traversal seeds.
This is a lightweight extraction — we only need names, not full schema.
"""
response = await self.openai.chat.completions.create(
model=config.entity_extraction_model,
messages=[{
"role": "user",
"content": f"""Extract all named entities from this query.
Return only a JSON array of entity name strings.
Be inclusive — if something might be an entity, include it.
Query: {query}
Return format: ["entity1", "entity2", ...]"""
}],
response_format={"type": "json_object"},
temperature=0.0
)
import json
result = json.loads(response.choices[0].message.content)
# Handle both {"entities": [...]} and direct array responses
if isinstance(result, list):
return result
return result.get("entities", []) or result.get("items", [])
async def vector_search(
self,
query: str,
top_k: int = 10
) -> list[dict]:
"""Standard vector similarity search."""
embedding_response = await self.openai.embeddings.create(
model=config.embedding_model,
input=query
)
query_vector = embedding_response.data[0].embedding
results = await self.qdrant.search(
collection_name=config.collection_name,
query_vector=query_vector,
limit=top_k,
with_payload=True
)
return [
{
"chunk_id": hit.payload.get("chunk_id"),
"text": hit.payload.get("text"),
"source": hit.payload.get("source"),
"score": hit.score
}
for hit in results
]
async def find_seed_nodes(self, entity_names: list[str]) -> list[str]:
"""
Match query entity mentions to actual graph nodes.
Uses full-text search for fuzzy matching.
"""
matched_nodes = []
async with self.neo4j.session() as session:
for name in entity_names:
result = await session.run(
"""
CALL db.index.fulltext.queryNodes(
'entity_name_index', $search_term
) YIELD node, score
WHERE score > 1.0
RETURN node.name AS name, node.entity_type AS type, score
ORDER BY score DESC
LIMIT 3
""",
search_term=name
)
records = await result.data()
matched_nodes.extend([r["name"] for r in records])
return list(set(matched_nodes))
async def multi_hop_traversal(
self,
seed_entities: list[str],
max_hops: int = 3,
max_paths: int = 20
) -> list[GraphPath]:
"""
Core multi-hop traversal.
Finds paths connecting seed entities through the knowledge graph.
"""
if not seed_entities:
return []
paths = []
async with self.neo4j.session() as session:
# Single-entity neighborhood exploration
for seed in seed_entities:
result = await session.run(
"""
MATCH (start:Entity {name: $seed})
CALL apoc.path.expandConfig(start, {
maxLevel: $max_hops,
relationshipFilter: 'AUTHORED|STUDIED|TREATS|CAUSES|
PARTICIPATED_IN|FUNDED_BY|CITES|
COLLABORATED_WITH|AFFILIATED_WITH',
uniqueness: 'NODE_PATH',
limit: $max_paths
}) YIELD path
WHERE length(path) > 0
RETURN
[node in nodes(path) | node.name] AS node_names,
[node in nodes(path) | node.entity_type] AS node_types,
[rel in relationships(path) | type(rel)] AS rel_types,
[rel in relationships(path) | rel.source_chunk] AS chunks,
length(path) AS hop_count
ORDER BY hop_count
LIMIT $max_paths
""",
seed=seed,
max_hops=max_hops,
max_paths=max_paths
)
records = await result.data()
for record in records:
# Interleave nodes and relationships for readable path
path_elements = []
nodes = record["node_names"]
rels = record["rel_types"]
types = record["node_types"]
for i, node in enumerate(nodes):
path_elements.append(f"({node}:{types[i]})")
if i < len(rels):
path_elements.append(f"-[:{rels[i]}]->")
chunk_ids = [c for c in record["chunks"] if c]
paths.append(GraphPath(
path_elements=path_elements,
start_entity=nodes[0] if nodes else "",
end_entity=nodes[-1] if nodes else "",
hop_count=record["hop_count"],
relevance_score=1.0 / (record["hop_count"] + 1), # Closer = more relevant
source_chunks=chunk_ids
))
# If multiple seed entities: find connecting paths between them
if len(seed_entities) >= 2:
for i in range(len(seed_entities)):
for j in range(i + 1, len(seed_entities)):
result = await session.run(
"""
MATCH (a:Entity {name: $entity_a}),
(b:Entity {name: $entity_b})
CALL apoc.algo.dijkstra(a, b,
'AUTHORED|STUDIED|TREATS|CAUSES|
PARTICIPATED_IN|COLLABORATED_WITH|CITES',
'confidence'
) YIELD path, weight
WHERE length(path) <= $max_hops
RETURN
[node in nodes(path) | node.name] AS node_names,
[node in nodes(path) | node.entity_type] AS node_types,
[rel in relationships(path) | type(rel)] AS rel_types,
[rel in relationships(path) | rel.source_chunk] AS chunks,
length(path) AS hop_count,
weight
LIMIT 5
""",
entity_a=seed_entities[i],
entity_b=seed_entities[j],
max_hops=max_hops
)
connecting_records = await result.data()
for record in connecting_records:
nodes = record["node_names"]
rels = record["rel_types"]
types = record["node_types"]
path_elements = []
for idx, node in enumerate(nodes):
path_elements.append(f"({node}:{types[idx]})")
if idx < len(rels):
path_elements.append(f"-[:{rels[idx]}]->")
paths.append(GraphPath(
path_elements=path_elements,
start_entity=nodes[0] if nodes else "",
end_entity=nodes[-1] if nodes else "",
hop_count=record["hop_count"],
relevance_score=2.0 / (record["hop_count"] + 1), # Connecting paths are more valuable
source_chunks=[c for c in record["chunks"] if c]
))
# Sort by relevance and deduplicate
paths.sort(key=lambda p: p.relevance_score, reverse=True)
return self._deduplicate_paths(paths)[:max_paths]
def _deduplicate_paths(self, paths: list[GraphPath]) -> list[GraphPath]:
"""Remove paths that are subsets of longer paths."""
seen = set()
unique_paths = []
for path in paths:
key = "->".join(path.path_elements)
if key not in seen:
seen.add(key)
unique_paths.append(path)
return unique_paths
async def retrieve(
self,
query: str,
vector_top_k: int = 10,
max_hops: int = 3
) -> HybridRetrievalResult:
"""
Main retrieval entry point. Runs vector search and graph traversal
concurrently then combines results.
"""
import asyncio
# Extract entities for graph seeding
query_entities = await self.extract_query_entities(query)
# Run vector search and graph operations concurrently
vector_task = asyncio.create_task(
self.vector_search(query, vector_top_k)
)
seed_task = asyncio.create_task(
self.find_seed_nodes(query_entities)
)
vector_results, seed_nodes = await asyncio.gather(vector_task, seed_task)
# Graph traversal happens after seed resolution
graph_paths = await self.multi_hop_traversal(seed_nodes, max_hops)
return HybridRetrievalResult(
vector_chunks=vector_results,
graph_paths=graph_paths,
seed_entities=seed_nodes,
query=query
)
The concurrent execution of vector search and entity extraction is important. These are independent operations and running them sequentially adds unnecessary latency. In production systems, this latency difference becomes significant at scale.
Raw retrieval results are useless without thoughtful assembly into a prompt. This is where many KG-RAG implementations fall apart — they dump every graph path into the context window and hope the LLM figures it out.
# context_assembly.py
import tiktoken
from retrieval import HybridRetrievalResult, GraphPath
class ContextAssembler:
def __init__(self, max_context_tokens: int = 6000):
self.max_context_tokens = max_context_tokens
self.encoder = tiktoken.encoding_for_model("gpt-4o")
def count_tokens(self, text: str) -> int:
return len(self.encoder.encode(text))
def format_graph_path(self, path: GraphPath) -> str:
"""Convert a graph path to readable natural language."""
return " ".join(path.path_elements)
def group_paths_by_relevance(
self,
paths: list[GraphPath]
) -> dict[int, list[GraphPath]]:
"""Group paths by hop count for structured presentation."""
groups = {}
for path in paths:
groups.setdefault(path.hop_count, []).append(path)
return dict(sorted(groups.items()))
def assemble_context(
self,
retrieval_result: HybridRetrievalResult
) -> str:
"""
Build a structured context string that clearly separates
vector evidence from graph relationship evidence.
"""
sections = []
token_budget = self.max_context_tokens
# Section 1: Identified query entities
if retrieval_result.seed_entities:
entity_section = (
"## Identified Entities in Query\n"
+ ", ".join(retrieval_result.seed_entities)
+ "\n"
)
sections.append(entity_section)
token_budget -= self.count_tokens(entity_section)
# Section 2: Knowledge graph relationship paths
if retrieval_result.graph_paths:
path_section_parts = ["## Knowledge Graph Relationships\n"]
path_section_parts.append(
"The following relationship chains were found in the knowledge graph:\n"
)
grouped = self.group_paths_by_relevance(retrieval_result.graph_paths)
for hop_count, paths in grouped.items():
hop_label = f"### {hop_count}-hop relationships:\n"
path_section_parts.append(hop_label)
for path in paths[:5]: # Max 5 paths per hop level
path_text = self.format_graph_path(path) + "\n"
path_tokens = self.count_tokens(path_text)
if token_budget - path_tokens > 1000: # Keep 1000 tokens for chunks
path_section_parts.append(f"- {path_text}")
token_budget -= path_tokens
sections.append("".join(path_section_parts))
# Section 3: Vector-retrieved text chunks (semantic evidence)
if retrieval_result.vector_chunks:
chunk_section_parts = ["## Retrieved Text Passages\n"]
for i, chunk in enumerate(retrieval_result.vector_chunks):
chunk_text = (
f"### Passage {i+1} (relevance: {chunk['score']:.3f})\n"
f"Source: {chunk.get('source', 'unknown')}\n"
f"{chunk['text']}\n\n"
)
chunk_tokens = self.count_tokens(chunk_text)
if token_budget - chunk_tokens > 200:
chunk_section_parts.append(chunk_text)
token_budget -= chunk_tokens
else:
break
sections.append("".join(chunk_section_parts))
return "\n".join(sections)
async def generate_answer(
query: str,
context: str,
model: str = "gpt-4o"
) -> str:
client = AsyncOpenAI(api_key=config.openai_api_key)
system_prompt = """You are an expert research assistant with access to both
structured knowledge graph relationships and retrieved text passages.
When answering:
1. PRIORITIZE multi-hop reasoning: use the Knowledge Graph Relationships to trace
connections that span multiple entities
2. USE text passages to provide supporting evidence and detail
3. EXPLICITLY cite which relationship chains support your reasoning
4. ACKNOWLEDGE when the available evidence doesn't fully support a definitive answer
5. DISTINGUISH between directly stated facts and inferred connections"""
response = await client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"""Context:
{context}
Question: {query}
Provide a thorough answer that traces the relevant relationship chains and
cites specific evidence from both the knowledge graph and text passages."""}
],
temperature=0.2
)
return response.choices[0].message.content
# pipeline.py
import asyncio
from pathlib import Path
from extraction import extract_entities_and_relationships
from disambiguation import EntityDisambiguator
from graph_store import KnowledgeGraphStore
from retrieval import HybridRetriever
from context_assembly import ContextAssembler, generate_answer
import hashlib
class KGRAGPipeline:
def __init__(self):
self.graph_store = KnowledgeGraphStore()
self.retriever = HybridRetriever()
self.assembler = ContextAssembler()
self.disambiguator = EntityDisambiguator()
async def ingest_document(
self,
text: str,
source_name: str,
chunk_size: int = 800,
chunk_overlap: int = 100
) -> None:
"""Full ingestion: chunk, extract, embed, and store."""
chunks = self._chunk_text(text, chunk_size, chunk_overlap)
# Process chunks with bounded concurrency (don't hammer the API)
semaphore = asyncio.Semaphore(5)
async def process_chunk(chunk_text: str, idx: int):
async with semaphore:
chunk_id = hashlib.md5(
f"{source_name}_{idx}".encode()
).hexdigest()
# Extract entities and relationships
extraction_result = await extract_entities_and_relationships(
chunk_text, chunk_id
)
# Store in knowledge graph
await self.graph_store.store_extraction_result(
extraction_result,
self.disambiguator
)
# Embed and store in vector store
await self._store_chunk_vector(chunk_text, chunk_id, source_name)
print(f"Processed chunk {idx}: {len(extraction_result.entities)} entities, "
f"{len(extraction_result.relationships)} relationships")
tasks = [
process_chunk(chunk, i)
for i, chunk in enumerate(chunks)
]
await asyncio.gather(*tasks)
def _chunk_text(
self,
text: str,
chunk_size: int,
overlap: int
) -> list[str]:
"""Simple sentence-aware chunking."""
import re
sentences = re.split(r'(?<=[.!?])\s+', text)
chunks = []
current_chunk = []
current_size = 0
for sentence in sentences:
sentence_size = len(sentence.split())
if current_size + sentence_size > chunk_size and current_chunk:
chunks.append(" ".join(current_chunk))
# Keep overlap sentences
overlap_sentences = []
overlap_size = 0
for s in reversed(current_chunk):
if overlap_size + len(s.split()) <= overlap:
overlap_sentences.insert(0, s)
overlap_size += len(s.split())
else:
break
current_chunk = overlap_sentences
current_size = overlap_size
current_chunk.append(sentence)
current_size += sentence_size
if current_chunk:
chunks.append(" ".join(current_chunk))
return chunks
async def _store_chunk_vector(
self,
text: str,
chunk_id: str,
source: str
) -> None:
from openai import AsyncOpenAI
from qdrant_client.models import PointStruct
import uuid
openai_client = AsyncOpenAI(api_key=config.openai_api_key)
response = await openai_client.embeddings.create(
model=config.embedding_model,
input=text
)
vector = response.data[0].embedding
await self.retriever.qdrant.upsert(
collection_name=config.collection_name,
points=[PointStruct(
id=str(uuid.uuid4()),
vector=vector,
payload={
"chunk_id": chunk_id,
"text": text,
"source": source
}
)]
)
async def query(self, question: str) -> str:
"""End-to-end query: retrieve, assemble context, generate answer."""
retrieval_result = await self.retriever.retrieve(question)
context = self.assembler.assemble_context(retrieval_result)
answer = await generate_answer(question, context)
return answer
async def close(self):
await self.graph_store.close()
await self.retriever.neo4j.close()
await self.retriever.qdrant.close()
# Example usage
async def main():
pipeline = KGRAGPipeline()
# Initialize schema
await pipeline.graph_store.initialize_schema()
# Create Qdrant collection
from qdrant_client.models import Distance, VectorParams
await pipeline.retriever.qdrant.create_collection(
collection_name=config.collection_name,
vectors_config=VectorParams(
size=config.embedding_dim,
distance=Distance.COSINE
)
)
# Ingest a document
sample_text = Path("research_paper.txt").read_text()
await pipeline.ingest_document(sample_text, "ACCORD_trial_2008")
# Query with multi-hop reasoning
answer = await pipeline.query(
"Which researchers from the ACCORD trial later collaborated on "
"SGLT2 inhibitor studies, and what outcomes did those studies examine?"
)
print(answer)
await pipeline.close()
if __name__ == "__main__":
asyncio.run(main())
Build a KG-RAG system over a domain you care about. Here's a structured exercise using publicly available data:
Dataset: Download 20-30 Wikipedia articles about interconnected topics. For example: major technology companies (Apple, Google, Microsoft, Amazon), their founders, products, and acquisitions.
Tasks:
Ingest the corpus using the pipeline above. After ingestion, connect to Neo4j Browser (http://localhost:7474) and run MATCH (n) RETURN n LIMIT 100 to visualize your graph. You should see a rich network of companies, people, and products connected by relationships.
Test single-entity queries first to validate your vector search: "What products does Apple make?" This should work well with pure vector search as a baseline.
Test multi-hop queries that require graph traversal:
Measure the difference: For each multi-hop query, compare the answer from your KG-RAG system against a standard RAG system (same vector store, no graph). Log which specific graph paths were traversed to answer each question.
Tune your disambiguation threshold: Find two entities in your corpus that should be the same node but aren't (or two that are different but got merged). Adjust the 0.92 threshold in EntityDisambiguator and observe the effect.
The temptation is to extract every possible entity type. Resist it. Every entity type you add multiplies the disambiguation problem and graph complexity. Start with 5-7 entity types that are central to your domain. You can always add more.
If your graph looks like a hairball with 50 node types, you've over-extracted. Good graphs have a small number of high-connectivity node types.
In Neo4j, (A)-[:AUTHORED]->(B) and (B)-[:AUTHORED]->(A) are different. If your extraction pipeline isn't consistent about direction, you'll create traversal queries that miss relationships or double-count them. Enforce direction in your schema documentation and validate it in your extraction prompt.
Running entity extraction on entire documents wastes tokens and produces poor results. LLMs have limited attention — they do better extraction on 400-800 word chunks than on 10,000-word documents. Always chunk first.
This will kill you at scale. If you run MERGE (e:Entity {name: $name, entity_type: $type}) without an index on (name, entity_type), Neo4j does a full node scan on every merge operation. With 100,000 entities, ingestion slows from seconds to hours.
-- Check if your indexes exist
SHOW INDEXES
The LLM will occasionally return malformed JSON, hallucinate relationship types not in your schema, or return an empty extraction for valid content. The tenacity retry decorator helps with transient failures, but you also need schema validation:
try:
result = ExtractionResult(**raw)
except ValidationError as e:
# Log the failure, continue without this chunk's graph data
logger.warning(f"Extraction validation failed for chunk {chunk_id}: {e}")
return ExtractionResult(entities=[], relationships=[], chunk_id=chunk_id)
Setting max_hops=5 or higher on a large graph will produce astronomically long traversals and likely timeout. In practice, 2-3 hops is sufficient for most multi-hop reasoning. 4+ hops typically produce noise — relationships that are technically connected but semantically irrelevant to the query.
If multi_hop_traversal returns empty paths even though you've ingested documents:
MATCH (n:Entity) RETURN count(n)MATCH ()-[r]->() RETURN type(r), count(r) ORDER BY count(r) DESCCALL db.index.fulltext.queryNodes('entity_name_index', 'YourEntityName') YIELD node RETURN nodeRETURN apoc.version()If you're seeing many spurious relationships in your graph, the most likely culprits are:
temperature=0.1 or lower.You've built a system that genuinely solves a problem that pure vector search cannot. The architecture — entity extraction, disambiguation, graph storage, hybrid retrieval, and context-aware generation — is production-ready and extensible.
Here's what you've actually built and why each part matters:
The extraction pipeline transforms unstructured text into structured knowledge, with LLM-powered extraction for domain-specific entities and embedding-based disambiguation to prevent graph fragmentation.
The knowledge graph in Neo4j stores not just entities but the relationships between them — enabling Cypher queries that traverse paths no embedding space can represent.
The hybrid retriever combines the best of both worlds: vector search for semantic relevance as an entry point, graph traversal for relational completeness. Neither alone is sufficient.
The context assembler ensures the LLM receives structured, well-organized evidence that distinguishes graph-derived relational facts from chunk-retrieved semantic evidence.
Where to go from here:
MERGE pattern already supports this architecturally.The knowledge graph approach scales well because relationships are explicit — you don't need to pray that the right chunks land in your context window. When a new paper citing the ACCORD trial gets ingested, those edges are immediately traversable. That's a fundamentally different capability than anything pure vector search can offer.
Learning Path: Building with LLMs