56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
|
"""This example uses an example Movie database where movies' plots are embedded
|
||
|
|
using OpenAI embeddings. OPENAI_API_KEY needs to be set in the environment for
|
||
|
|
this example to run.
|
||
|
|
|
||
|
|
Also requires minimal Cypher knowledge to write the retrieval query.
|
||
|
|
|
||
|
|
It shows how to use a hybrid retriever to find context
|
||
|
|
similar to a query **text** using vector+text similarity
|
||
|
|
and graph traversal.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import neo4j
|
||
|
|
from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
|
||
|
|
from neo4j_graphrag.retrievers import HybridCypherRetriever
|
||
|
|
|
||
|
|
# Define database credentials
|
||
|
|
URI = "neo4j+s://demo.neo4jlabs.com"
|
||
|
|
AUTH = ("recommendations", "recommendations")
|
||
|
|
DATABASE = "recommendations"
|
||
|
|
INDEX_NAME = "moviePlotsEmbedding"
|
||
|
|
FULLTEXT_INDEX_NAME = "movieFulltext"
|
||
|
|
|
||
|
|
# for each Movie node matched by the vector search, retrieve more context:
|
||
|
|
# the name of all actors starring in that movie
|
||
|
|
RETRIEVAL_QUERY = """
|
||
|
|
RETURN node.title as movieTitle,
|
||
|
|
node.plot as moviePlot,
|
||
|
|
collect { MATCH (actor:Actor)-[:ACTED_IN]->(node) RETURN actor.name } AS actors,
|
||
|
|
score as similarityScore
|
||
|
|
"""
|
||
|
|
|
||
|
|
with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
|
||
|
|
# Initialize the retriever
|
||
|
|
retriever = HybridCypherRetriever(
|
||
|
|
driver=driver,
|
||
|
|
vector_index_name=INDEX_NAME,
|
||
|
|
fulltext_index_name=FULLTEXT_INDEX_NAME,
|
||
|
|
# note: embedder is optional if you only use query_vector
|
||
|
|
embedder=OpenAIEmbeddings(),
|
||
|
|
retrieval_query=RETRIEVAL_QUERY,
|
||
|
|
# optionally, configure how to format the results
|
||
|
|
# (see corresponding example in 'customize' directory)
|
||
|
|
# result_formatter=None,
|
||
|
|
# optionally, set neo4j database
|
||
|
|
neo4j_database=DATABASE,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Perform the similarity search for a text query
|
||
|
|
# (retrieve the top 5 most similar nodes)
|
||
|
|
query_text = "Who were the actors in Avatar?"
|
||
|
|
print(retriever.search(query_text=query_text, top_k=5))
|
||
|
|
|
||
|
|
# note: it is also possible to query from a query_vector directly:
|
||
|
|
# query_vector: list[float] = [...]
|
||
|
|
# retriever.search(query_vector=query_vector, top_k=5)
|