Beyond Vector Search: Building a Medical GraphRAG with Neo4j and LlamaIndex

We’ve all been there: you build a sleek RAG (Retrieval-Augmented Generation) pipeline using a standard vector database, only for it to fail miserably when asked a complex question like, "How does Metformin interact with dietary fiber in patients with Type 2 Diabetes?"

Standard vector search excels at finding "similar" chunks of text, but it’s notoriously bad at capturing multi-hop relationships and structural logic. This is where GraphRAG comes in. By combining the power of Neo4j with LlamaIndex and Ollama, we can transform flat PDF documents into a rich, interconnected Knowledge Graph. This approach significantly reduces LLM hallucinations in medical contexts by providing the model with explicit relationships rather than just semantic similarity.

If you're looking to master advanced Knowledge Graph architectures or production-ready RAG optimization, you’ll find deep-dive explorations and more production-ready examples over at the WellAlly Tech Blog, which served as a major inspiration for this architectural pattern.

🏗️ The Architecture: From Text to Triples

Traditional RAG treats your data like a library of disconnected sticky notes. GraphRAG treats it like a brain. Here is the data flow for our Chronic Disease Knowledge Base:

graph TD
    A[PDF Medical Literature] --> B{LlamaIndex Property Graph Index}
    B --> C[Ollama / LLM Entity Extraction]
    C --> D[Nodes: Disease, Symptom, Medication]
    C --> E[Edges: TREATS, CONTRAINDICATES, CAUSES]
    D --> F[(Neo4j Graph Database)]
    E --> F
    G[User Query] --> H[Graph + Vector Hybrid Retriever]
    H --> F
    F --> I[Context-Rich Prompt]
    I --> J[Ollama Response]

🛠️ Tech Stack

• Orchestration: LlamaIndex
• Graph Database: Neo4j (Local or Aura)
• Local LLM: Ollama (using Llama 3 or Mistral)
• Data Source: Chronic Disease PDFs (PubMed/Medical Journals)

💻 Step-by-Step Implementation

1. Setting up the Neo4j Environment

First, ensure you have a Neo4j instance running. You can use Docker or a free Neo4j Aura instance.

docker run -d \
    --name neo4j \
    -p 7474:7474 -p 7687:7687 \
    -e NEO4J_AUTH=neo4j/password \
    neo4j:latest

2. Defining the Schema and Extractor

In medical literature, we care about specific entities. We'll use LlamaIndex's SchemaLLMPathExtractor to ensure our graph follows a logical medical structure.

from llama_index.core.indices.property_graph import SchemaLLMPathExtractor
from llama_index.llms.ollama import Ollama

# Initialize our local powerhouse
llm = Ollama(model="llama3", request_timeout=360.0)

# Define our medical schema
entities = ["Disease", "Medication", "Symptom", "DietaryFactor"]
relations = ["TREATS", "CAUSES", "CONTRAINDICATES", "EXACERBATES"]
validation_schema = [
    ("Disease", "TREATS", "Medication"),
    ("Medication", "CONTRAINDICATES", "Disease"),
    ("DietaryFactor", "EXACERBATES", "Symptom"),
]

kg_extractor = SchemaLLMPathExtractor(
    llm=llm,
    possible_entities=entities,
    possible_relations=relations,
    kg_validation_schema=validation_schema,
    strict=True,
)

3. Ingesting PDFs into the Graph

Now, we load our chronic disease PDFs and let the LLM do the heavy lifting of extracting entities and relationships directly into Neo4j.

from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.graph_stores.neo4j import Neo4jGraphStore
from llama_index.core import PropertyGraphIndex

# Load medical documents
documents = SimpleDirectoryReader("./medical_data/").load_data()

# Setup Neo4j Storage
graph_store = Neo4jGraphStore(
    username="neo4j",
    password="password",
    url="bolt://localhost:7687",
)

storage_context = StorageContext.from_defaults(graph_store=graph_store)

# Create the Property Graph Index
index = PropertyGraphIndex.from_documents(
    documents,
    kg_extractors=[kg_extractor],
    storage_context=storage_context,
    llm=llm,
)

4. The Hybrid Query: Vector + Graph

The magic happens when we query. LlamaIndex will now look for the semantic vector matches and traverse the graph edges to find related concepts that aren't mentioned in the same text chunk.

query_engine = index.as_query_engine(
    sub_retrievers=["vector", "synonym", "path"],
    llm=llm
)

response = query_engine.query(
    "What are the risks of high sodium intake for a patient with Hypertension and Stage 2 CKD?"
)
print(f"🥑 Medical Insights: {response}")

🏥 Why this matters for Medical AI

In a standard RAG setup, the model might find a chunk about "Sodium" and another about "Hypertension." However, it might miss the subtle link where "High Sodium" triggers "Water Retention," which specifically worsens "Stage 2 CKD" (Chronic Kidney Disease).

By using Neo4j, we explicitly store the edge (Sodium)-[EXACERBATES]->(Hypertension). When the retriever hits the "Sodium" node, it automatically pulls in the connected "Hypertension" and "CKD" context, even if they were 50 pages apart in the PDF!

Pro Tip: For more advanced patterns on handling asynchronous data ingestion and multi-agent medical workflows, check out the specialized guides at wellally.tech/blog. They cover the "Day 2" operations of keeping your knowledge graph updated as new research emerges.

🎯 Conclusion

Moving from flat embeddings to a Knowledge Graph is a game-changer for high-stakes domains like healthcare. By combining Neo4j's structural integrity with LlamaIndex's ease of use, we've built a system that doesn't just "guess" based on word proximity—it "understands" based on defined relationships.

What's next?

1. Try adding a Semantic Cache to speed up repeated queries.
2. Implement Entity Resolution to ensure "High Blood Pressure" and "Hypertension" map to the same node.
3. Let me know in the comments: Are you using GraphRAG in production yet? 🥑💻