Ruff

Author: NathalieCharbel

examples/customize/build_graph/components/resolvers/spacy_entity_resolver_pre_filter.py

@@ -22,7 +22,7 @@ async def main(driver: neo4j.Driver) -> None:
         driver,
         # let's filter all entities that belong to a certain docId
         filter_query="WHERE (entity)-[:FROM_CHUNK]->(:Chunk)-[:FROM_DOCUMENT]->(doc:"
-                     "Document {id = 'docId'}",
+        "Document {id = 'docId'}",
         # optionally, change the properties used for resolution (default is "name")
         # resolve_properties=["name", "ssn"],
         # the similarity threshold (default is 0.8)

src/neo4j_graphrag/experimental/components/resolver.py

@@ -35,9 +35,9 @@ class EntityResolver(Component, abc.ABC):
     """
 
     def __init__(
-            self,
-            driver: neo4j.Driver,
-            filter_query: Optional[str] = None,
+        self,
+        driver: neo4j.Driver,
+        filter_query: Optional[str] = None,
     ) -> None:
         self.driver = driver_config.override_user_agent(driver)
         self.filter_query = filter_query
@@ -74,11 +74,11 @@ class SinglePropertyExactMatchResolver(EntityResolver):
     """
 
     def __init__(
-            self,
-            driver: neo4j.Driver,
-            filter_query: Optional[str] = None,
-            resolve_property: str = "name",
-            neo4j_database: Optional[str] = None,
+        self,
+        driver: neo4j.Driver,
+        filter_query: Optional[str] = None,
+        resolve_property: str = "name",
+        neo4j_database: Optional[str] = None,
     ) -> None:
         super().__init__(driver, filter_query)
         self.resolve_property = resolve_property
@@ -174,13 +174,13 @@ class SpaCySemanticMatchResolver(EntityResolver):
     """
 
     def __init__(
-            self,
-            driver: neo4j.Driver,
-            filter_query: Optional[str] = None,
-            resolve_properties: Optional[List[str]] = None,
-            similarity_threshold: float = 0.8,
-            spacy_model: str = "en_core_web_lg",
-            neo4j_database: Optional[str] = None,
+        self,
+        driver: neo4j.Driver,
+        filter_query: Optional[str] = None,
+        resolve_properties: Optional[List[str]] = None,
+        similarity_threshold: float = 0.8,
+        spacy_model: str = "en_core_web_lg",
+        neo4j_database: Optional[str] = None,
     ) -> None:
         super().__init__(driver, filter_query)
         self.resolve_properties = resolve_properties or ["name"]
@@ -230,8 +230,9 @@ async def run(self) -> ResolutionStats:
             node_embeddings = {}
             for ent in entities:
                 # concatenate all textual properties (if non-null) into a single string
-                texts = [str(ent[p]) for p in self.resolve_properties if
-                         p in ent and ent[p]]
+                texts = [
+                    str(ent[p]) for p in self.resolve_properties if p in ent and ent[p]
+                ]
                 combined_text = " ".join(texts).strip()
                 if combined_text:
                     node_embeddings[ent["id"]] = self.nlp(combined_text).vector
@@ -258,8 +259,9 @@ async def run(self) -> ResolutionStats:
                         "YIELD node RETURN id(node)"
                     )
                     result, _, _ = self.driver.execute_query(
-                        merge_query, {"ids": list(node_id_set)},
-                        database_=self.neo4j_database
+                        merge_query,
+                        {"ids": list(node_id_set)},
+                        database_=self.neo4j_database,
                     )
                     merged_count += len(result)
 
@@ -307,7 +309,7 @@ def _load_or_download_spacy_model(model_name: str):
         except OSError as e:
             # The exact error message can differ slightly depending on spaCy version,
             # so you may want to be broader or narrower with handling logic:
-            if 'doesn\'t seem to be a Python package or a valid path' in str(e):
+            if "doesn't seem to be a Python package or a valid path" in str(e):
                 print(f"Model '{model_name}' not found. Downloading...")
                 spacy.cli.download(model_name)
                 return spacy.load(model_name)

tests/unit/experimental/components/test_resolver.py

@@ -58,8 +58,8 @@ async def test_simple_resolver_custom_filter(driver: MagicMock) -> None:
             )
         ]
     )
-    
-    
+
+
 @pytest.mark.asyncio
 async def test_spacy_resolver_match_on_name_property(driver: MagicMock) -> None:
     driver.execute_query.side_effect = [
@@ -70,7 +70,7 @@ async def test_spacy_resolver_match_on_name_property(driver: MagicMock) -> None:
                         "lab": "Person",
                         "labelCluster": [
                             {"id": 1, "name": "Alice"},
-                            {"id": 2, "name": "Alice"}
+                            {"id": 2, "name": "Alice"},
                         ],
                     }
                 )
@@ -94,6 +94,7 @@ async def test_spacy_resolver_match_on_name_property(driver: MagicMock) -> None:
 
     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 = [
@@ -104,7 +105,7 @@ async def test_spacy_resolver_no_merge(driver: MagicMock) -> None:
                         "lab": "Person",
                         "labelCluster": [
                             {"id": 1, "name": "Alice"},
-                            {"id": 2, "name": "Bob"}
+                            {"id": 2, "name": "Bob"},
                         ],
                     }
                 )
@@ -122,9 +123,10 @@ async def test_spacy_resolver_no_merge(driver: MagicMock) -> None:
 
     assert driver.execute_query.call_count == 1
 
+
 @pytest.mark.asyncio
 async def test_spacy_resolver_match_on_multiple_text_properties(
-        driver: MagicMock
+    driver: MagicMock,
 ) -> None:
     driver.execute_query.side_effect = [
         (
@@ -133,16 +135,8 @@ async def test_spacy_resolver_match_on_multiple_text_properties(
                     {
                         "lab": "Person",
                         "labelCluster": [
-                            {
-                                "id": 10,
-                                "name": "John Smith",
-                                "ssn": "23-45-6789"
-                            },
-                            {
-                                "id": 11,
-                                "name": "Jonathan Smith",
-                                "ssn": "23-45-6789"
-                            }
+                            {"id": 10, "name": "John Smith", "ssn": "23-45-6789"},
+                            {"id": 11, "name": "Jonathan Smith", "ssn": "23-45-6789"},
                         ],
                     }
                 )
@@ -158,8 +152,7 @@ async def test_spacy_resolver_match_on_multiple_text_properties(
     ]
 
     resolver = SpaCySemanticMatchResolver(
-        driver=driver,
-        resolve_properties=["name", "ssn"]
+        driver=driver, resolve_properties=["name", "ssn"]
     )
     res = await resolver.run()
     assert isinstance(res, ResolutionStats)