Simple Neo4j GraphRAG Sample App

app.py

#!/usr/bin/env python3

"""
Simple Neo4j GraphRAG implementation using OpenAI embeddings and LLM

Installation instructions:
python -m venv venv
source venv/bin/activate
pip install neo4j 'neo4j-graphrag[openai]'
"""

import os
from neo4j import GraphDatabase
from neo4j_graphrag.embeddings import OpenAIEmbeddings
from neo4j_graphrag.generation import GraphRAG
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.retrievers import VectorRetriever

# Configuration
NEO4J_URI = "neo4j+s://demo.neo4jlabs.com"
NEO4J_USERNAME = "goodreads"
NEO4J_PASSWORD = "goodreads"
NEO4J_DATABASE = "goodreads"
VECTOR_INDEX_NAME = "book-descriptions"
LLM_MODEL = "gpt-4o"
EMBEDDING_MODEL = "text-embedding-3-small"

# Set your OpenAI API key here or as environment variable
OPENAI_API_KEY = "your-openai-api-key-here"  # Replace with your key

def create_graphrag_pipeline():
    """Initialize the GraphRAG components"""
    # Set OpenAI API key
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
    
    # Connect to Neo4j
    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
    
    # Create embeddings and LLM
    embedder = OpenAIEmbeddings(model=EMBEDDING_MODEL)
    llm = OpenAILLM(model_name=LLM_MODEL, model_params={"temperature": 0.0})
    
    # Create retriever
    retriever = VectorRetriever(
        driver=driver,
        index_name=VECTOR_INDEX_NAME,
        embedder=embedder,
        neo4j_database=NEO4J_DATABASE
    )
    
    # Create GraphRAG pipeline
    rag = GraphRAG(retriever=retriever, llm=llm)
    
    return rag, driver

def ask_question(question, top_k=5):
    """Ask a question and get an AI answer from the book database"""
    rag, driver = create_graphrag_pipeline()
    
    try:
        print(f"🔍 Question: {question}")
        
        # Search and get answer
        response = rag.search(
            query_text=question,
            retriever_config={"top_k": top_k}
        )
        
        print(f"🤖 Answer: {response.answer}")
        return response.answer
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None
    finally:
        driver.close()

if __name__ == "__main__":
    # Make sure to set your OpenAI API key above!
    
    # Ask questions about books
    questions = [
        "Which books are about motherhood and friendship?",
        "What are the best science fiction novels?",
        "Tell me about books that deal with mental health"
    ]
    
    for question in questions:
        print("=" * 60)
        ask_question(question, top_k=3)
        print()