GraphRAG using the neo4j library

Introduction.

In this post, I created a knowledge graph using neo4j and built a RAG system using the knowledge graph as external information.

In this post, I have tried to build a RAG system easily using the neo4j library.

Information Sources

Getting Started With the Neo4j GraphRAG Python Package I tried the code in this article and modified some parts of the code, which I will comment in the text.
Enriching Vector Search With Graph Traversal Using the GraphRAG Python Package I tried the code in this article. It is almost the same as above.
Hybrid Retrieval for GraphRAG Applications Using the Neo4j GraphRAG Package for Python I have not tried the code in this article yet. I will try it in the future.

Execution Result

Connect to the prepared neo4j database

# Establish a connection to your Neo4j database

from neo4j import GraphDatabase
# Demo database credentials
URI = "neo4j+s://demo.neo4jlabs.com"
AUTH = ("recommendations", "recommendations")
# Connect to Neo4j database
driver = GraphDatabase.driver(URI, auth=AUTH)

Set OpenAI API key to environment variable

# Set OpenAI API key

import os

os.environ['OPENAI_API_KEY'] = 'sk-*****'

Prepare GraphRAG retriever

# Prepare GraphRAG's retriever

from neo4j_graphrag.retrievers import VectorRetriever
from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings

embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
retriever = VectorRetriever(
    driver,
    index_name="moviePlotsEmbedding",
    embedder=embedder,
    return_properties=["title", "plot"],
)

Changed package name from “neo4j-graphrag” to “neo4j_graphrag”.

Build a pipeline with GraphRAG

# GraphRAG

from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.generation import GraphRAG

# LLM
# Note: the OPENAI_API_KEY must be in the env vars
llm = OpenAILLM(model_name="gpt-4o-mini", model_params={"temperature": 0})

# Initialize the RAG pipline
rag = GraphRAG(retriever=retriever, llm=llm)

# Query the graph
query_text = "What movies are sad rommances?"
response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
print(response.answer)

Execution Result

Based on the context provided, the following movies can be considered sad romances:

1. **Autumn in New York** - This film is a romantic drama about an aging playboy who falls for a sweet, but terminally ill, young woman, which adds a poignant and sad element to the romance.

2. **Bed of Roses** - While it has romantic elements, the story of a young career girl being swept off her feet by a shy florist may also carry emotional weight, depending on the character's journeys.

These films explore themes of love intertwined with sadness and emotional challenges.

Conclusion

Using neo4j’s graphrag package, I have found that it is very easy to build a RAG system that uses the knowledge graph as external information.

Next, I am going to try a RAG system that performs hybrid search with GraphRAG for information source 3.