267 lines
9.7 KiB
Python
267 lines
9.7 KiB
Python
|
# Copyright (c) "Neo4j"
|
||
|
|
# Neo4j Sweden AB [https://neo4j.com]
|
||
|
|
# #
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
# #
|
||
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
# #
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
import neo4j
|
||
|
|
import pytest
|
||
|
|
from neo4j_graphrag.indexes import (
|
||
|
|
retrieve_fulltext_index_info,
|
||
|
|
retrieve_vector_index_info,
|
||
|
|
upsert_vectors,
|
||
|
|
)
|
||
|
|
from neo4j_graphrag.types import EntityType
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_vector_index_info_happy_path(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_vector_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="vector-index-name",
|
||
|
|
label_or_type="Document",
|
||
|
|
embedding_property="vectorProperty",
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "vector-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "VECTOR"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["vectorProperty"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
options = index_info.get("options")
|
||
|
|
assert isinstance(options, dict)
|
||
|
|
index_config = options.get("indexConfig")
|
||
|
|
assert isinstance(index_config, dict)
|
||
|
|
embedding_dimension = index_config.get("vector.dimensions")
|
||
|
|
assert isinstance(embedding_dimension, int)
|
||
|
|
assert embedding_dimension == 1536
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_vector_index_info_no_index_name(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_vector_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="",
|
||
|
|
label_or_type="Document",
|
||
|
|
embedding_property="vectorProperty",
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "vector-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "VECTOR"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["vectorProperty"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
options = index_info.get("options")
|
||
|
|
assert isinstance(options, dict)
|
||
|
|
index_config = options.get("indexConfig")
|
||
|
|
assert isinstance(index_config, dict)
|
||
|
|
embedding_dimension = index_config.get("vector.dimensions")
|
||
|
|
assert isinstance(embedding_dimension, int)
|
||
|
|
assert embedding_dimension == 1536
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_vector_index_info_no_label_or_property(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_vector_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="vector-index-name",
|
||
|
|
label_or_type="",
|
||
|
|
embedding_property="",
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "vector-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "VECTOR"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["vectorProperty"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
options = index_info.get("options")
|
||
|
|
assert isinstance(options, dict)
|
||
|
|
index_config = options.get("indexConfig")
|
||
|
|
assert isinstance(index_config, dict)
|
||
|
|
embedding_dimension = index_config.get("vector.dimensions")
|
||
|
|
assert isinstance(embedding_dimension, int)
|
||
|
|
assert embedding_dimension == 1536
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_vector_index_info_wrong_info(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_vector_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="err",
|
||
|
|
label_or_type="err",
|
||
|
|
embedding_property="err",
|
||
|
|
)
|
||
|
|
assert index_info is None
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_fulltext_index_info_happy_path(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_fulltext_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="fulltext-index-name",
|
||
|
|
label_or_type="Document",
|
||
|
|
text_properties=["short_text_property"],
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "fulltext-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "FULLTEXT"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["short_text_property"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_fulltext_index_info_no_index_name(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_fulltext_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="",
|
||
|
|
label_or_type="Document",
|
||
|
|
text_properties=["short_text_property"],
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "fulltext-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "FULLTEXT"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["short_text_property"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_fulltext_index_info_no_label_or_properties(
|
||
|
|
driver: neo4j.Driver,
|
||
|
|
) -> None:
|
||
|
|
index_info = retrieve_fulltext_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="fulltext-index-name",
|
||
|
|
label_or_type="",
|
||
|
|
text_properties=[""],
|
||
|
|
)
|
||
|
|
assert index_info is not None
|
||
|
|
index_name = index_info.get("name")
|
||
|
|
assert index_name == "fulltext-index-name"
|
||
|
|
index_type = index_info.get("type")
|
||
|
|
assert index_type == "FULLTEXT"
|
||
|
|
labels_or_types = index_info.get("labelsOrTypes")
|
||
|
|
assert labels_or_types == ["Document"]
|
||
|
|
properties = index_info.get("properties")
|
||
|
|
assert properties == ["short_text_property"]
|
||
|
|
entity_type = index_info.get("entityType")
|
||
|
|
assert entity_type == "NODE"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.usefixtures("setup_neo4j_for_retrieval")
|
||
|
|
def test_retrieve_fulltext_index_info_wrong_info(driver: neo4j.Driver) -> None:
|
||
|
|
index_info = retrieve_fulltext_index_info(
|
||
|
|
driver=driver,
|
||
|
|
index_name="err",
|
||
|
|
label_or_type="err",
|
||
|
|
text_properties=[""],
|
||
|
|
)
|
||
|
|
assert index_info is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_upsert_vectors_on_nodes(driver: neo4j.Driver) -> None:
|
||
|
|
driver.execute_query("MATCH (n) DETACH DELETE n;")
|
||
|
|
result = driver.execute_query(
|
||
|
|
"CREATE (p:Character {name: 'Paul Atreides'}), "
|
||
|
|
"(v:Character {name: 'Vladimir Harkonnen'}) "
|
||
|
|
"RETURN [elementId(p), elementId(v)] AS ids"
|
||
|
|
)
|
||
|
|
ids = result.records[0]["ids"]
|
||
|
|
assert len(ids) == 2
|
||
|
|
upsert_vectors(
|
||
|
|
driver=driver,
|
||
|
|
ids=ids,
|
||
|
|
embedding_property="embedding",
|
||
|
|
embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
|
||
|
|
entity_type=EntityType.NODE,
|
||
|
|
)
|
||
|
|
result = driver.execute_query(
|
||
|
|
"MATCH (c:Character) RETURN "
|
||
|
|
"elementId(c) as id, c.name as name, c.embedding as embedding"
|
||
|
|
)
|
||
|
|
records = [r.data() for r in result.records]
|
||
|
|
records_sorted = sorted(records, key=lambda x: x["id"])
|
||
|
|
expected_records = [
|
||
|
|
{"id": ids[0], "name": "Paul Atreides", "embedding": [1.0, 2.0, 3.0]},
|
||
|
|
{"id": ids[1], "name": "Vladimir Harkonnen", "embedding": [4.0, 5.0, 6.0]},
|
||
|
|
]
|
||
|
|
expected_records_sorted = sorted(expected_records, key=lambda x: x["id"])
|
||
|
|
assert records_sorted == expected_records_sorted
|
||
|
|
|
||
|
|
|
||
|
|
def test_upsert_vectors_on_relationships(driver: neo4j.Driver) -> None:
|
||
|
|
driver.execute_query("MATCH (n) DETACH DELETE n;")
|
||
|
|
result = driver.execute_query(
|
||
|
|
"CREATE (:Character {name: 'Paul Atreides'})-[a:IS_MEMBER_OF]->(:House {name: 'Atreides'}), "
|
||
|
|
"(:Character {name: 'Vladimir Harkonnen'})-[h:IS_MEMBER_OF]->(:House {name: 'Harkonnen'}) "
|
||
|
|
"RETURN [elementId(a), elementId(h)] AS ids"
|
||
|
|
)
|
||
|
|
ids = result.records[0]["ids"]
|
||
|
|
assert len(ids) == 2
|
||
|
|
upsert_vectors(
|
||
|
|
driver=driver,
|
||
|
|
ids=ids,
|
||
|
|
embedding_property="embedding",
|
||
|
|
embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
|
||
|
|
entity_type=EntityType.RELATIONSHIP,
|
||
|
|
)
|
||
|
|
result = driver.execute_query(
|
||
|
|
"MATCH (c:Character)-[r:IS_MEMBER_OF]->(h:House) RETURN "
|
||
|
|
"type(r) AS type, elementId(r) AS id, r.embedding AS embedding, "
|
||
|
|
"c.name as character, h.name AS house"
|
||
|
|
)
|
||
|
|
records = [r.data() for r in result.records]
|
||
|
|
records_sorted = sorted(records, key=lambda x: x["id"])
|
||
|
|
expected_records = [
|
||
|
|
{
|
||
|
|
"type": "IS_MEMBER_OF",
|
||
|
|
"id": ids[0],
|
||
|
|
"embedding": [1.0, 2.0, 3.0],
|
||
|
|
"character": "Paul Atreides",
|
||
|
|
"house": "Atreides",
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "IS_MEMBER_OF",
|
||
|
|
"id": ids[1],
|
||
|
|
"embedding": [4.0, 5.0, 6.0],
|
||
|
|
"character": "Vladimir Harkonnen",
|
||
|
|
"house": "Harkonnen",
|
||
|
|
},
|
||
|
|
]
|
||
|
|
expected_records_sorted = sorted(expected_records, key=lambda x: x["id"])
|
||
|
|
assert records_sorted == expected_records_sorted
|