Master GraphRAG: Boost AI Responses with Neo4j & LangChain

In this article, we’ll explore how to leverage LLama 3.1, a locally-run model, to perform GraphRAG operations using just 50 lines of code. This powerful technique combines the strengths of knowledge graphs and vector databases to enhance retrieval-augmented generation, offering a more nuanced approach to information retrieval and processing.

Understanding GraphRAG: The Fusion of Knowledge Graphs and Vector Databases

GraphRAG is an innovative method that enhances retrieval-augmented generation by considering the relationships between entities and documents. At its core, GraphRAG revolves around two key concepts: nodes and relationships, building upon traditional vector databases for enhanced retrieval.

The Power of Integration

One of the most effective GraphRAG architectures integrates knowledge graphs with vector databases. This approach leverages the strengths of both systems to gather relevant information:

1. Knowledge Graph Construction: Captures relationships between vector chunks, including document hierarchies.
2. Contextual Enrichment: Provides structured entity information near the chunks retrieved from vector search.
3. Enhanced Prompting: The enriched prompt, complete with valuable additional context, is fed into the Large Language Model (LLM) for processing.
4. Response Generation: The LLM generates a response based on the enriched input.
5. User Delivery: The generated answer is returned to the user.

This architecture is particularly well-suited for applications such as customer support, semantic search, and personalized recommendations.

Nodes and Relationships: The Building Blocks of GraphRAG

• Nodes: Represent entities or concepts extracted from data chunks, such as people, organizations, events, or locations. Each node in the knowledge graph contains attributes and characteristics that provide additional context for the entity.
• Relationships: Define connections between nodes, which can include various types of associations:
• Hierarchical (e.g., parent-child)
• Temporal (e.g., before-after)
• Causal (cause-effect)

These relationships also have attributes describing their nature and strength. When dealing with numerous documents, this approach results in a comprehensive graph depicting the interconnections among all documents.

A Simple Example

Consider a dataset where nodes represent entities like Apple Inc. and Tim Cook. The relationship between these nodes could be described as “Tim Cook is the CEO of Apple Inc.” This simple example illustrates how GraphRAG can capture and represent complex real-world relationships.

The Power and Challenge of GraphRAG

While GraphRAG offers a powerful approach to information retrieval and processing, it comes with a significant computational cost. Extracting entities from each document and using an LLM to compute the relationship graph can be resource-intensive. This is where locally-run models like LLama 3.1 shine, offering the ability to implement this approach without relying on external API calls or cloud-based services.

Hands-On Tutorial: Implementing GraphRAG with LLama 3.1

In this tutorial, we’ll combine LangChain, LLama, Ollama, and Neo4j as our graph database to create an information graph about a large Italian family with multiple restaurants. This scenario provides ample opportunities to model complex relationships.

Setting Up the Environment

1. Pull the LLama 3.1 8b model using Ollama:

   ollama pull llama2:8b

2. Configure and Run Neo4j Database:
   Use the provided Docker Compose file to set up Neo4j. Navigate to the neo4j folder and run:

   docker compose up

3. Install Dependencies:
   Install the required packages using pip:

   %pip install --upgrade --quiet langchain langchain-community langchain-openai langchain-ollama langchain-experimental neo4j tiktoken yfiles_jupyter_graphs python-dotenv

Setting Up the Environment

Before diving into the implementation, we need to set up our development environment with the necessary dependencies.

Installing Required Packages

To get started, we’ll install the following packages:

%pip install --upgrade --quiet langchain langchain-community langchain-openai langchain-ollama langchain-experimental neo4j tiktoken yfiles_jupyter_graphs python-dotenv

This command installs LangChain and its various components, OpenAI’s LangChain, Ollama, LangChain Experimental (which contains the graph solution), Neo4j, and additional utilities for visualization and environment management.

Importing Essential Classes

After installation, we import the required classes:

from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.output_parsers import StrOutputParser
import os
from langchain_community.graphs import Neo4jGraph
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI
from langchain_community.chat_models import ChatOllama
from langchain_experimental.graph_transformers import LLMGraphTransformer
from neo4j import GraphDatabase
from yfiles_jupyter_graphs import GraphWidget
from langchain_community.vectorstores import Neo4jVector
from langchain_openai import OpenAIEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores.neo4j_vector import remove_lucene_chars
from dotenv import load_dotenv

load_dotenv()

These imports provide us with the necessary tools to work with LangChain, Neo4j, and various AI models.

Preparing the Dataset

We’ll use a sample dataset stored in a file named dummy_text.txt. This dataset contains information about an Italian family, including names, relationships, and other details that we’ll represent in our graph.

To process this data:

1. Load the text file using TextLoader.
2. Split the text into manageable chunks using RecursiveCharacterTextSplitter.

Creating the Graph

The LLM Graph Transformer is responsible for converting our text documents into a format that Neo4j can process.

llm = ChatOllama() if os.getenv("llm_type") != "openai" else ChatOpenAI()
graph_transformer = LLMGraphTransformer.from_llm(llm)
graph_documents = graph_transformer.convert_to_graph_documents(documents)

This process may take a few minutes, even for small datasets. The result is a graph document containing nodes (e.g., Family, Love, Tradition) and relationships between them.

Visualizing the Graph

After storing the graph documents in Neo4j, we can visualize the knowledge graph:

driver = GraphDatabase.driver(os.getenv("NEO4J_URI"), auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD")))
with driver.session() as session:
    result = session.run("MATCH (s)-[r:mentions]->(t) RETURN s, r, t LIMIT 100")
    # Visualization code here

The visualization reveals entities like Petro, who likes the kitchen and the sea, and is a parent to Sophia. This comprehensive knowledge graph represents the relationships and attributes extracted from our text data.

Implementing Vector Storage

In addition to the graph database, we’ll create a vector store using Neo4jVector:

vector_store = Neo4jVector.from_existing_graph(embedding=OpenAIEmbeddings())
retriever = vector_store.as_retriever()

This allows us to perform vector searches alongside our graph-based queries.

Preparing Entities for Graph Database Queries

To effectively query our graph database, we need to extract entities from user queries. We’ll create a custom Entities model and a prompt template for this purpose:

class Entities(BaseModel):
    entities: list[str] = Field(description="List of extracted entities")

entity_extraction_prompt = ChatPromptTemplate.from_messages([
    ("system", "Extract organizations, persons, and business entities from the text."),
    ("human", "{input}")
])

entity_chain = entity_extraction_prompt | llm | Entities.parse()

Creating a Hybrid Retriever

We’ll combine our graph-based retriever with the vector store retriever to create a hybrid approach:

def full_retriever(query: str):
    graph_docs = graph_retriever(query)
    vector_docs = retriever.invoke(query)
    return graph_docs + vector_docs

Building the Final Chain

Our final RAG chain combines the hybrid retriever with a language model to generate responses:

template = """Answer the question based on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

chain = (
    RunnablePassthrough.assign(
        context=full_retriever
    )
    | prompt
    | llm
    | StrOutputParser()
)

Testing the System

Let’s test our Graph RAG system with a sample question:

question = "Who is Nonna Lucia? Did she teach anyone about restaurants or cooking?"
result = chain.invoke(question)
print(result)

Output:

Nonna Lucia is the matriarch of the Caruso family and a culinary mentor. She taught her grandchildren the art of Sicilian cooking, including recipes for Caponata and fresh pasta.

This response demonstrates how our system effectively combines information from the graph database and vector store to provide a comprehensive and accurate answer.

Conclusion

By integrating Neo4j with LangChain and implementing a Graph RAG approach, we’ve created a powerful system capable of generating more contextually aware and accurate responses. This method leverages the strengths of both graph databases and vector embeddings, resulting in a robust AI assistant that can handle complex queries with improved precision.

As we’ve discussed in our previous article on advanced NLP techniques, the combination of structured knowledge (graphs) and unstructured data (text embeddings) can significantly enhance the capabilities of AI systems. This Graph RAG implementation is a prime example of how these techniques can be applied in practice.

For those interested in further exploration, consider diving into entity recognition techniques to improve the accuracy of entity extraction in your Graph RAG system.

By mastering these advanced techniques, developers and data scientists can create more sophisticated AI applications that provide deeper insights and more accurate responses across a wide range of domains.