참고소스 수정본

참고/neo4j-graphrag-python-main/tests/e2e/conftest.py

@@ -0,0 +1,225 @@
+#  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
+
+import os
+import random
+import string
+import uuid
+from typing import Any, Generator
+from unittest.mock import MagicMock
+
+import pytest
+from neo4j import Driver, GraphDatabase
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.indexes import (
+    create_fulltext_index,
+    create_vector_index,
+    drop_index_if_exists,
+)
+from neo4j_graphrag.llm import LLMInterface
+from neo4j_graphrag.retrievers import VectorRetriever
+
+from ..e2e.utils import EMBEDDING_BIOLOGY
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture(scope="module")
+def driver() -> Generator[Any, Any, Any]:
+    uri = "neo4j://localhost:7687"
+    auth = ("neo4j", "password")
+    driver = GraphDatabase.driver(uri, auth=auth)
+    yield driver
+    driver.close()
+
+
+@pytest.fixture(scope="function")
+def llm() -> MagicMock:
+    mock = MagicMock(spec=LLMInterface)
+    mock.supports_structured_output = False
+    return mock
+
+
+@pytest.fixture
+def embedder() -> Embedder:
+    embedder = MagicMock(spec=Embedder)
+    return embedder
+
+
+class RandomEmbedder(Embedder):
+    def embed_query(self, text: str) -> list[float]:
+        return [random.random() for _ in range(1536)]
+
+
+class BiologyEmbedder(Embedder):
+    def embed_query(self, text: str) -> list[float]:
+        if text == "biology":
+            return EMBEDDING_BIOLOGY
+        raise ValueError(f"Unknown embedding text: {text}")
+
+
+@pytest.fixture(scope="module")
+def random_embedder() -> RandomEmbedder:
+    return RandomEmbedder()
+
+
+@pytest.fixture(scope="module")
+def biology_embedder() -> BiologyEmbedder:
+    return BiologyEmbedder()
+
+
+@pytest.fixture(scope="function")
+def retriever_mock() -> MagicMock:
+    return MagicMock(spec=VectorRetriever)
+
+
+@pytest.fixture
+def harry_potter_text() -> str:
+    with open(os.path.join(BASE_DIR, "data/documents/harry_potter.txt"), "r") as f:
+        text = f.read()
+    return text
+
+
+@pytest.fixture
+def harry_potter_text_part1() -> str:
+    with open(
+        os.path.join(BASE_DIR, "data/documents/harry_potter_part1.txt"), "r"
+    ) as f:
+        text = f.read()
+    return text
+
+
+@pytest.fixture
+def harry_potter_text_part2() -> str:
+    with open(
+        os.path.join(BASE_DIR, "data/documents/harry_potter_part2.txt"), "r"
+    ) as f:
+        text = f.read()
+    return text
+
+
+@pytest.fixture(scope="module")
+def setup_neo4j_for_retrieval(driver: Driver) -> None:
+    vector_index_name = "vector-index-name"
+    fulltext_index_name = "fulltext-index-name"
+
+    # Delete data and drop indexes to prevent data leakage
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    drop_index_if_exists(driver, vector_index_name)
+    drop_index_if_exists(driver, fulltext_index_name)
+
+    # Create a vector index
+    create_vector_index(
+        driver,
+        vector_index_name,
+        label="Document",
+        embedding_property="vectorProperty",
+        dimensions=1536,
+        similarity_fn="euclidean",
+    )
+
+    # Create a fulltext index
+    create_fulltext_index(
+        driver,
+        fulltext_index_name,
+        label="Document",
+        node_properties=["short_text_property"],
+    )
+
+    # Insert 10 vectors and authors
+    vector = [random.random() for _ in range(1536)]
+
+    def random_str(n: int) -> str:
+        return "".join([random.choice(string.ascii_letters) for _ in range(n)])
+
+    for i in range(10):
+        insert_query = (
+            "MERGE (doc:Document {id: $id})"
+            "ON CREATE SET  doc.int_property = $i, "
+            "               doc.short_text_property = toString($i)"
+            "WITH doc "
+            "CALL db.create.setNodeVectorProperty(doc, 'vectorProperty', $vector)"
+            "WITH doc "
+            "MERGE (author:Author {name: $authorName})"
+            "MERGE (doc)-[:AUTHORED_BY]->(author)"
+            "RETURN doc, author"
+        )
+
+        parameters = {
+            "id": str(uuid.uuid4()),
+            "i": i,
+            "vector": vector,
+            "authorName": random_str(1536),
+        }
+        driver.execute_query(insert_query, parameters)
+
+
+@pytest.fixture(scope="module")
+def setup_neo4j_for_schema_query(driver: Driver) -> None:
+    # Delete all nodes in the graph
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    # Create two nodes and a relationship
+    driver.execute_query(
+        """
+        CREATE (la:LabelA {property_a: 'a'})
+        CREATE (lb:LabelB)
+        CREATE (lc:LabelC)
+        MERGE (la)-[:REL_TYPE]-> (lb)
+        MERGE (la)-[:REL_TYPE {rel_prop: 'abc'}]-> (lc)
+        """
+    )
+
+
+@pytest.fixture(scope="module")
+def setup_neo4j_for_schema_query_with_excluded_labels(driver: Driver) -> None:
+    # Delete all nodes in the graph
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    # Create two labels and a relationship to be excluded
+    driver.execute_query(
+        "CREATE (:_Bloom_Scene_{property_a: 'a'})-[:_Bloom_HAS_SCENE_{property_b: 'b'}]->(:_Bloom_Perspective_)"
+    )
+
+
+@pytest.fixture(scope="module")
+def setup_neo4j_for_kg_construction(driver: Driver) -> None:
+    # Delete all nodes and indexes in the graph
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    vector_index_name = "vector-index-name"
+    fulltext_index_name = "fulltext-index-name"
+    drop_index_if_exists(driver, vector_index_name)
+    drop_index_if_exists(driver, fulltext_index_name)
+
+    # Create a vector index with the dimensions used by the Hugging Face all-MiniLM-L6-v2 model
+    create_vector_index(
+        driver,
+        vector_index_name,
+        label="Document",
+        embedding_property="vectorProperty",
+        dimensions=3,
+        similarity_fn="euclidean",
+    )
+
+
+@pytest.fixture(scope="module")
+def setup_neo4j_for_kg_construction_with_chunks(
+    driver: Driver,
+    setup_neo4j_for_kg_construction: Any,
+) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query('CREATE (:Chunk {id: "0", index: 0, text: "some text"})')
+    driver.execute_query(
+        'CREATE (:Chunk {id: "1", index: 1, text: "some longer text"})'
+    )