From 3d2bf72b3d5a50801d8d2fea00f65eb5e3081c66 Mon Sep 17 00:00:00 2001 From: estelle Date: Mon, 28 Apr 2025 09:40:58 +0200 Subject: [PATCH 1/2] Strict mode: if node/relationship is in schema but no properties are defined, do not filter allowed properties --- .../components/entity_relation_extractor.py | 10 ++++++++-- .../components/test_entity_relation_extractor.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py index f21a22fed..9967ef901 100644 --- a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py +++ b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py @@ -404,7 +404,10 @@ def _enforce_nodes( if not schema_entity: continue allowed_props = schema_entity.get("properties", []) - filtered_props = self._enforce_properties(node.properties, allowed_props) + if allowed_props: + filtered_props = self._enforce_properties(node.properties, allowed_props) + else: + filtered_props = node.properties if filtered_props: valid_nodes.append( Neo4jNode( @@ -469,7 +472,10 @@ def _enforce_relationships( continue allowed_props = schema_relation.get("properties", []) - filtered_props = self._enforce_properties(rel.properties, allowed_props) + if allowed_props: + filtered_props = self._enforce_properties(rel.properties, allowed_props) + else: + filtered_props = rel.properties valid_rels.append( Neo4jRelationship( diff --git a/tests/unit/experimental/components/test_entity_relation_extractor.py b/tests/unit/experimental/components/test_entity_relation_extractor.py index 21fe8807d..cebc3d31e 100644 --- a/tests/unit/experimental/components/test_entity_relation_extractor.py +++ b/tests/unit/experimental/components/test_entity_relation_extractor.py @@ -374,7 +374,7 @@ async def test_extractor_schema_enforcement_valid_nodes_with_empty_props() -> No result: Neo4jGraph = await extractor.run(chunks, schema=schema) - assert len(result.nodes) == 0 + assert len(result.nodes) == 1 @pytest.mark.asyncio From f11e4b481b7ce670122bf34029922bb4a9546e72 Mon Sep 17 00:00:00 2001 From: estelle Date: Mon, 28 Apr 2025 09:45:39 +0200 Subject: [PATCH 2/2] Ruff --- .../experimental/components/entity_relation_extractor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py index 9967ef901..112d8b2c6 100644 --- a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py +++ b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py @@ -405,7 +405,9 @@ def _enforce_nodes( continue allowed_props = schema_entity.get("properties", []) if allowed_props: - filtered_props = self._enforce_properties(node.properties, allowed_props) + filtered_props = self._enforce_properties( + node.properties, allowed_props + ) else: filtered_props = node.properties if filtered_props: