참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_resolver.py

@@ -0,0 +1,303 @@
+#  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 types import SimpleNamespace
+from unittest.mock import MagicMock, call, patch
+
+import neo4j
+import numpy as np
+import pytest
+from neo4j_graphrag.experimental.components.resolver import (
+    FuzzyMatchResolver,
+    SinglePropertyExactMatchResolver,
+    SpaCySemanticMatchResolver,
+)
+from neo4j_graphrag.experimental.components.types import ResolutionStats
+
+
+class FakeNLPModel:
+    """
+    Stand-in for a spaCy NLP model for unit tests.
+    It returns an object with a `.vector` attribute so the resolver can compute cosine similarity.
+    """
+
+    def __call__(self, text: str) -> SimpleNamespace:
+        if "23-45-6789" in text or text == "Alice":
+            return SimpleNamespace(vector=np.array([1.0, 0.0], dtype=np.float64))
+        if text == "Bob":
+            return SimpleNamespace(vector=np.array([0.0, 1.0], dtype=np.float64))
+        return SimpleNamespace(vector=np.array([0.0, 0.0], dtype=np.float64))
+
+
+@pytest.mark.asyncio
+async def test_simple_resolver(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        ([neo4j.Record({"c": 2})], None, None),
+        ([neo4j.Record({"c": 1})], None, None),
+    ]
+    resolver = SinglePropertyExactMatchResolver(driver=driver)
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+    assert driver.execute_query.call_count == 2
+    driver.execute_query.assert_has_calls(
+        [call("MATCH (entity:__Entity__)  RETURN count(entity) as c", database_=None)]
+    )
+
+
+@pytest.mark.asyncio
+async def test_simple_resolver_custom_filter(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        ([neo4j.Record({"c": 2})], None, None),
+        ([neo4j.Record({"c": 1})], None, None),
+    ]
+    resolver = SinglePropertyExactMatchResolver(
+        driver=driver, filter_query="WHERE not entity:Resolved"
+    )
+    await resolver.run()
+    driver.execute_query.assert_has_calls(
+        [
+            call(
+                "MATCH (entity:__Entity__) WHERE not entity:Resolved RETURN count(entity) as c",
+                database_=None,
+            )
+        ]
+    )
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_match_on_name_property(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Alice"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 1})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_no_merge(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    res = await resolver.run()
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 0
+
+    assert driver.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_match_on_multiple_text_properties(
+    driver: MagicMock,
+) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 10, "name": "John Smith", "ssn": "23-45-6789"},
+                            {"id": 11, "name": "Jonathan Smith", "ssn": "23-45-6789"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 10})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(
+        driver=driver, resolve_properties=["name", "ssn"], nlp=FakeNLPModel()
+    )
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_no_merge(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        )
+    ]
+
+    resolver = FuzzyMatchResolver(driver=driver)
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 0
+
+    assert driver.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_multiple_properties(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 10, "name": "John Smith", "ssn": "123-45-6789"},
+                            {"id": 11, "name": "Jon Smith", "ssn": "123-45-6789"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 10})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = FuzzyMatchResolver(driver=driver, resolve_properties=["name", "ssn"])
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_normalization(driver: MagicMock) -> None:
+    # instantiate with a dummy driver
+    resolver = FuzzyMatchResolver(driver=driver)
+
+    sim = resolver.compute_similarity("  ALICE  ", "alice!")
+    assert sim == 1
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_caching(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Alice"},
+                            {"id": 3, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 1})],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 3})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    # patch spaCy NLP call to track how often embeddings are computed
+    with patch.object(resolver, "nlp", wraps=resolver.nlp) as mock_nlp:
+        await resolver.run()
+
+    # "Alice" should be embedded only once, despite being used twice.
+    # "Bob" should be embedded once.
+    assert mock_nlp.call_count == 2, (
+        f"Expected spaCy to embed each unique text once. Got {mock_nlp.call_count} "
+        f"calls."
+    )
+
+    # "Alice" and "Bob" are expected to be the only two distinct texts passed to spaCy.
+    called_texts = {call.args[0] for call in mock_nlp.call_args_list}
+    assert called_texts == {"Alice", "Bob"}