diff --git a/examples/vector_cypher_retrieval.py b/examples/vector_cypher_retrieval.py new file mode 100644 index 000000000..3671b2444 --- /dev/null +++ b/examples/vector_cypher_retrieval.py @@ -0,0 +1,68 @@ +from neo4j import GraphDatabase +from neo4j_genai import VectorCypherRetriever + +import random +import string +from neo4j_genai.embedder import Embedder +from neo4j_genai.indexes import create_vector_index, drop_index + + +URI = "neo4j://localhost:7687" +AUTH = ("neo4j", "password") + +INDEX_NAME = "embedding-name" +DIMENSION = 1536 + +# Connect to Neo4j database +driver = GraphDatabase.driver(URI, auth=AUTH) + + +# Create Embedder object +class CustomEmbedder(Embedder): + def embed_query(self, text: str) -> list[float]: + return [random.random() for _ in range(DIMENSION)] + + +# Generate random strings +def random_str(n: int) -> str: + return "".join([random.choice(string.ascii_letters) for _ in range(n)]) + + +embedder = CustomEmbedder() + +# Creating the index +drop_index(driver, INDEX_NAME) +create_vector_index( + driver, + INDEX_NAME, + label="Document", + property="propertyKey", + dimensions=DIMENSION, + similarity_fn="euclidean", +) + +# Initialize the retriever +retrieval_query = "MATCH (node)-[:AUTHORED_BY]->(author:Author)" "RETURN author.name" +retriever = VectorCypherRetriever(driver, INDEX_NAME, retrieval_query, embedder) + +# Upsert the query +vector = [random.random() for _ in range(DIMENSION)] +insert_query = ( + "MERGE (doc:Document {id: $id})" + "WITH doc " + "CALL db.create.setNodeVectorProperty(doc, 'propertyKey', $vector)" + "WITH doc " + "MERGE (author:Author {name: $authorName})" + "MERGE (doc)-[:AUTHORED_BY]->(author)" + "RETURN doc, author" +) +parameters = { + "id": random.randint(0, 10000), + "vector": vector, + "authorName": random_str(10), +} +driver.execute_query(insert_query, parameters) + +# Perform the search +query_text = "Find me the closest text" +print(retriever.search(query_text=query_text, top_k=5))