Skip to content

Commit

Permalink
Adds GenAI client object (#1)
Browse files Browse the repository at this point in the history
* Adds GenAIClient class and examples
  • Loading branch information
willtai authored Mar 4, 2024
1 parent 1c20c75 commit 7f27093
Show file tree
Hide file tree
Showing 18 changed files with 892 additions and 112 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1 +1,4 @@
dist/
**/__pycache__/*
*.py[cod]
.mypy_cache/
6 changes: 6 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0 # Use the ref you want to point at
hooks:
- id: trailing-whitespace
# - id: ...
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
# Neo4j GenAI package for Python

This repository contains the official Neo4j GenAI features for Python.
54 changes: 54 additions & 0 deletions examples/similarity_search_for_text.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
from typing import List
from neo4j import GraphDatabase
from neo4j.exceptions import DatabaseError
from neo4j_genai.client import GenAIClient

from random import random
from neo4j_genai.embeddings import Embeddings

URI = "neo4j://localhost:7687"
AUTH = ("neo4j", "password")

INDEX_NAME = "embedding-name"
DIMENSION = 1536

# Connect to Neo4j database
driver = GraphDatabase.driver(URI, auth=AUTH)


# Create Embeddings object
class CustomEmbeddings(Embeddings):
def embed_query(self, text: str) -> List[float]:
return [random() for _ in range(DIMENSION)]


embeddings = CustomEmbeddings()

# Initialize the client
client = GenAIClient(driver, embeddings)

# Creating the index
client.create_index(
INDEX_NAME,
label="Document",
property="propertyKey",
dimensions=DIMENSION,
similarity_fn="euclidean",
)

# Upsert the query
vector = [random() for _ in range(DIMENSION)]
insert_query = (
"MERGE (n:Document)"
"WITH n "
"CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
"RETURN n"
)
parameters = {
"vector": vector,
}
client.database_query(insert_query, params=parameters)

# Perform the similarity search for a text query
query_text = "hello world"
print(client.similarity_search(INDEX_NAME, query_text=query_text, top_k=5))
42 changes: 42 additions & 0 deletions examples/similarity_search_for_vector.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
from neo4j import GraphDatabase
from neo4j_genai.client import GenAIClient

from random import random

URI = "neo4j://localhost:7687"
AUTH = ("neo4j", "password")

INDEX_NAME = "embedding-name"
DIMENSION = 1536

# Connect to Neo4j database
driver = GraphDatabase.driver(URI, auth=AUTH)

# Initialize the client
client = GenAIClient(driver)

# Creating the index
client.create_index(
INDEX_NAME,
label="Document",
property="propertyKey",
dimensions=DIMENSION,
similarity_fn="euclidean",
)

# Upsert the vector
vector = [random() for _ in range(DIMENSION)]
insert_query = (
"MERGE (n:Document)"
"WITH n "
"CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
"RETURN n"
)
parameters = {
"vector": vector,
}
client.database_query(insert_query, params=parameters)

# Perform the similarity search for a vector query
query_vector = [random() for _ in range(DIMENSION)]
print(client.similarity_search(INDEX_NAME, query_vector=query_vector, top_k=5))
Binary file removed neo4j_genai/__pycache__/__init__.cpython-311.pyc
Binary file not shown.
Binary file removed neo4j_genai/__pycache__/tools.cpython-311.pyc
Binary file not shown.
35 changes: 0 additions & 35 deletions neo4j_genai/tests/test.py

This file was deleted.

68 changes: 0 additions & 68 deletions neo4j_genai/tests/test2.py

This file was deleted.

2 changes: 0 additions & 2 deletions neo4j_genai/tools.py

This file was deleted.

Loading

0 comments on commit 7f27093

Please sign in to comment.