neo4j_graphrag_quickstart_vector_retriever_openai_llm.py

1import os
2from neo4j import GraphDatabase
3from neo4j_graphrag.llm import OpenAILLM
4from neo4j_graphrag.retrievers import VectorRetriever
5from neo4j_graphrag.generation import GraphRAG
6
7# Connection details
8URI = "neo4j+s://<your-neo4j-instance-id>.databases.neo4j.io"
9AUTH = ("neo4j", "<your-password>")
10
11# Initialize the LLM
12# Ensure OPENAI_API_KEY is set in your environment variables
13llm = OpenAILLM(model_name="gpt-4o")
14
15# Initialize the Retriever
16# This assumes you have an existing Vector Index named 'vector' in Neo4j
17retriever = VectorRetriever(
18    driver=GraphDatabase.driver(URI, auth=AUTH),
19    index_name="vector",
20    embedder=None  # Can be an OpenAIEmbeddings object or similar
21)
22
23# Initialize the GraphRAG pipeline
24rag = GraphRAG(retriever=retriever, llm=llm)
25
26# Run a query
27query = "What are the main benefits of using Neo4j for GraphRAG?"
28response = rag.search(query_text=query)
29
30print(response.answer)