참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/test_indexes.py

@@ -0,0 +1,429 @@
+#  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.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import neo4j.exceptions
+import pytest
+from neo4j_graphrag.exceptions import Neo4jIndexError, Neo4jInsertionError
+from neo4j_graphrag.indexes import (
+    create_fulltext_index,
+    create_vector_index,
+    drop_index_if_exists,
+    upsert_vector,
+    upsert_vector_on_relationship,
+    upsert_vectors,
+)
+
+
+def test_create_vector_index_happy_path(driver: MagicMock) -> None:
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name OPTIONS "
+        "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(driver, "my-index", "People", "name", 2048, "cosine")
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_fail_if_exists(driver: MagicMock) -> None:
+    create_query = (
+        "CREATE VECTOR INDEX $name  FOR (n:People) ON n.name OPTIONS "
+        "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver, "my-index", "People", "name", 2048, "cosine", fail_if_exists=True
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_ensure_escaping(driver: MagicMock) -> None:
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name OPTIONS "
+        "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver, "my-complicated-`-index", "People", "name", 2048, "cosine"
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {
+            "name": "my-complicated-`-index",
+            "dimensions": 2048,
+            "similarity_fn": "cosine",
+        },
+        database_=None,
+    )
+
+
+def test_create_vector_index_negative_dimension(driver: MagicMock) -> None:
+    with pytest.raises(Neo4jIndexError) as excinfo:
+        create_vector_index(driver, "my-index", "People", "name", -5, "cosine")
+    assert "Error for inputs to create_vector_index" in str(excinfo)
+
+
+def test_create_vector_index_validation_error_dimensions(driver: MagicMock) -> None:
+    with pytest.raises(Neo4jIndexError) as excinfo:
+        create_vector_index(driver, "my-index", "People", "name", "no-dim", "cosine")  # type: ignore
+    assert "Error for inputs to create_vector_index" in str(excinfo)
+
+
+def test_create_vector_index_raises_error_with_neo4j_client_error(
+    driver: MagicMock,
+) -> None:
+    driver.execute_query.side_effect = neo4j.exceptions.ClientError
+    with pytest.raises(Neo4jIndexError):
+        create_vector_index(driver, "my-index", "People", "name", 2048, "cosine")
+
+
+def test_create_vector_index_validation_error_similarity_fn(driver: MagicMock) -> None:
+    with pytest.raises(Neo4jIndexError) as excinfo:
+        create_vector_index(driver, "my-index", "People", "name", 1536, "algebra")  # type: ignore
+    assert "Error for inputs to create_vector_index" in str(excinfo)
+
+
+def test_drop_index_if_exists(driver: MagicMock) -> None:
+    drop_query = "DROP INDEX $name IF EXISTS"
+
+    drop_index_if_exists(driver, "my-index")
+
+    driver.execute_query.assert_called_once_with(
+        drop_query,
+        {"name": "my-index"},
+        database_=None,
+    )
+
+
+def test_drop_index_if_exists_raises_error_with_neo4j_client_error(
+    driver: MagicMock,
+) -> None:
+    drop_query = "DROP INDEX $name IF EXISTS"
+
+    drop_index_if_exists(driver, "my-index")
+
+    driver.execute_query.assert_called_once_with(
+        drop_query,
+        {"name": "my-index"},
+        database_=None,
+    )
+
+
+def test_create_fulltext_index_happy_path(driver: MagicMock) -> None:
+    label = "node-label"
+    text_node_properties = ["property-1", "property-2"]
+    create_query = (
+        "CREATE FULLTEXT INDEX $name IF NOT EXISTS "
+        f"FOR (n:`{label}`) ON EACH "
+        f"[{', '.join(['n.`' + property + '`' for property in text_node_properties])}]"
+    )
+
+    create_fulltext_index(driver, "my-index", label, text_node_properties)
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index"},
+        database_=None,
+    )
+
+
+def test_create_fulltext_index_fail_if_exists(driver: MagicMock) -> None:
+    label = "node-label"
+    text_node_properties = ["property-1", "property-2"]
+    create_query = (
+        "CREATE FULLTEXT INDEX $name  "
+        f"FOR (n:`{label}`) ON EACH "
+        f"[{', '.join(['n.`' + property + '`' for property in text_node_properties])}]"
+    )
+
+    create_fulltext_index(
+        driver, "my-index", label, text_node_properties, fail_if_exists=True
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index"},
+        database_=None,
+    )
+
+
+def test_create_fulltext_index_raises_error_with_neo4j_client_error(
+    driver: MagicMock,
+) -> None:
+    label = "node-label"
+    text_node_properties = ["property-1", "property-2"]
+    driver.execute_query.side_effect = neo4j.exceptions.ClientError
+
+    with pytest.raises(Neo4jIndexError):
+        create_fulltext_index(driver, "my-index", label, text_node_properties)
+
+
+def test_create_fulltext_index_empty_node_properties(driver: MagicMock) -> None:
+    label = "node-label"
+    node_properties: list[str] = []
+
+    with pytest.raises(Neo4jIndexError) as excinfo:
+        create_fulltext_index(driver, "my-index", label, node_properties)
+
+    assert "Error for inputs to create_fulltext_index" in str(excinfo)
+
+
+def test_create_fulltext_index_ensure_escaping(driver: MagicMock) -> None:
+    label = "node-label"
+    text_node_properties = ["property-1", "property-2"]
+    create_query = (
+        "CREATE FULLTEXT INDEX $name IF NOT EXISTS "
+        f"FOR (n:`{label}`) ON EACH "
+        f"[{', '.join(['n.`' + property + '`' for property in text_node_properties])}]"
+    )
+
+    create_fulltext_index(driver, "my-complicated-`-index", label, text_node_properties)
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-complicated-`-index"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_with_filterable_properties(driver: MagicMock) -> None:
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name"
+        " WITH [n.`age`, n.`city`]"
+        " OPTIONS { indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver,
+        "my-index",
+        "People",
+        "name",
+        2048,
+        "cosine",
+        filterable_properties=["age", "city"],
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_with_single_filterable_property(
+    driver: MagicMock,
+) -> None:
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name"
+        " WITH [n.`year`]"
+        " OPTIONS { indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver,
+        "my-index",
+        "People",
+        "name",
+        2048,
+        "cosine",
+        filterable_properties=["year"],
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_with_empty_filterable_properties(
+    driver: MagicMock,
+) -> None:
+    """Empty list should produce same query as no filterable_properties."""
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name OPTIONS "
+        "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver,
+        "my-index",
+        "People",
+        "name",
+        2048,
+        "cosine",
+        filterable_properties=[],
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_create_vector_index_with_none_filterable_properties(
+    driver: MagicMock,
+) -> None:
+    """None (default) should produce same query as no filterable_properties."""
+    create_query = (
+        "CREATE VECTOR INDEX $name IF NOT EXISTS FOR (n:People) ON n.name OPTIONS "
+        "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
+    )
+
+    create_vector_index(
+        driver,
+        "my-index",
+        "People",
+        "name",
+        2048,
+        "cosine",
+        filterable_properties=None,
+    )
+
+    driver.execute_query.assert_called_once_with(
+        create_query,
+        {"name": "my-index", "dimensions": 2048, "similarity_fn": "cosine"},
+        database_=None,
+    )
+
+
+def test_upsert_vector_happy_path(driver: MagicMock) -> None:
+    id = 1
+    embedding_property = "embedding"
+    vector = [1.0, 2.0, 3.0]
+
+    upsert_vector(driver, id, embedding_property, vector)
+
+    upsert_query = (
+        "MATCH (n) "
+        "WHERE elementId(n) = $node_element_id "
+        "WITH n "
+        "CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) "
+        "RETURN n"
+    )
+
+    driver.execute_query.assert_called_once_with(
+        upsert_query,
+        {
+            "node_element_id": id,
+            "embedding_property": embedding_property,
+            "vector": vector,
+        },
+        database_=None,
+    )
+
+
+def test_upsert_vector_on_relationship_happy_path(driver: MagicMock) -> None:
+    id = 1
+    embedding_property = "embedding"
+    vector = [1.0, 2.0, 3.0]
+
+    upsert_vector_on_relationship(driver, id, embedding_property, vector)
+
+    upsert_query = (
+        "MATCH ()-[r]->() "
+        "WHERE elementId(r) = $rel_element_id "
+        "WITH r "
+        "CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) "
+        "RETURN r"
+    )
+
+    driver.execute_query.assert_called_once_with(
+        upsert_query,
+        {
+            "rel_element_id": id,
+            "embedding_property": embedding_property,
+            "vector": vector,
+        },
+        database_=None,
+    )
+
+
+def test_upsert_vector_on_relationship_raises_neo4j_insertion_error(
+    driver: MagicMock,
+) -> None:
+    id = 1
+    embedding_property = "embedding"
+    vector = [1.0, 2.0, 3.0]
+    driver.execute_query.side_effect = neo4j.exceptions.ClientError
+
+    with pytest.raises(Neo4jInsertionError) as excinfo:
+        upsert_vector_on_relationship(driver, id, embedding_property, vector)
+    assert "Upserting vector to Neo4j failed" in str(excinfo)
+
+
+def test_upsert_vector_raises_neo4j_insertion_error(
+    driver: MagicMock,
+) -> None:
+    id = 1
+    embedding_property = "embedding"
+    vector = [1.0, 2.0, 3.0]
+    driver.execute_query.side_effect = neo4j.exceptions.ClientError
+
+    with pytest.raises(Neo4jInsertionError) as excinfo:
+        upsert_vector(driver, id, embedding_property, vector)
+
+    assert "Upserting vector to Neo4j failed" in str(excinfo)
+
+
+def test_upsert_vectors_wrong_entity_type(driver: MagicMock) -> None:
+    with pytest.raises(ValueError) as exc_info:
+        upsert_vectors(
+            driver=driver,
+            ids=["1"],
+            embedding_property="embedding",
+            embeddings=[[1.0, 2.0, 3.0]],
+            neo4j_database="neo4j",
+            entity_type="WRONG_ENTITY_TYPE",  # type: ignore[arg-type]
+        )
+    assert str(exc_info.value) == "entity_type must be either 'NODE' or 'RELATIONSHIP'"
+
+
+def test_upsert_vectors_mismatched_ids_and_embeddings(driver: MagicMock) -> None:
+    with pytest.raises(ValueError) as exc_info:
+        upsert_vectors(
+            driver=driver,
+            ids=["1"],
+            embedding_property="embedding",
+            embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+            neo4j_database="neo4j",
+        )
+    assert str(exc_info.value) == "ids and embeddings must be the same length"
+
+
+def test_upsert_vectors_inconsistent_embedding_sizes(driver: MagicMock) -> None:
+    with pytest.raises(ValueError) as exc_info:
+        upsert_vectors(
+            driver=driver,
+            ids=["1", "2"],
+            embedding_property="embedding",
+            embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0]],
+            neo4j_database="neo4j",
+        )
+    assert str(exc_info.value) == "All embeddings must be of the same size"