From 7ca04411a59101b5dddbb01829fdf16a2a912fa6 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 01:54:37 +0200 Subject: [PATCH 01/12] spacy-entity-relation-extractor-task-001: add normalize() helper and WORD_FILTER constant Create spacy_utils.py with shared normalisation utilities: - normalize(text) lowercases, applies NFC Unicode normalisation, and collapses whitespace - WORD_FILTER frozenset of ~90 common English stopwords for entity key filtering - 21 unit tests covering edge cases, Unicode NFC equivalence, and type immutability Co-Authored-By: Claude Sonnet 4.6 --- ...ogress-spacy-entity-relation-extractor.txt | 4 + .../tasks-spacy-entity-relation-extractor.yml | 201 ++++++++++++++++++ .../experimental/components/spacy_utils.py | 161 ++++++++++++++ .../components/test_spacy_utils.py | 93 ++++++++ 4 files changed, 459 insertions(+) create mode 100644 .plans/progress-spacy-entity-relation-extractor.txt create mode 100644 .plans/tasks-spacy-entity-relation-extractor.yml create mode 100644 src/neo4j_graphrag/experimental/components/spacy_utils.py create mode 100644 tests/unit/experimental/components/test_spacy_utils.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt new file mode 100644 index 000000000..f740385ee --- /dev/null +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -0,0 +1,4 @@ +## TASK-001 — 2026-06-24T08:00:00Z +Created `spacy_utils.py` with `normalize()` (lowercases, NFC-normalises Unicode, collapses whitespace) and `WORD_FILTER` frozenset of ~90 common English stopwords; 21 unit tests all green. +Files: +src/neo4j_graphrag/experimental/components/spacy_utils.py, +tests/unit/experimental/components/test_spacy_utils.py +Decisions: Added `unicodedata.normalize("NFC", ...)` step per reviewer feedback to ensure visually identical Unicode strings compare equal as graph node keys. diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml new file mode 100644 index 000000000..c63093c71 --- /dev/null +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -0,0 +1,201 @@ +feature: spacy-entity-relation-extractor +prd: ./prd-spacy-entity-relation-extractor.md +created_at: 2026-06-24T00:00:00Z +updated_at: 2026-06-24T00:00:00Z + +tasks: + - id: task-001 + title: "Implement normalize() helper and WORD_FILTER constant" + description: | + Create shared normalisation utilities used by both components. + Add a `normalize(text: str) -> str` function that lowercases and + collapses whitespace, and a `WORD_FILTER: frozenset[str]` of common + stopwords to exclude from entity positions. Place in + `src/neo4j_graphrag/experimental/components/spacy_utils.py`. + status: completed + completed_at: "2026-06-24T08:00:00Z" + dependencies: [] + acceptance_criteria: + - "normalize(' Supply Chain ') returns 'supply chain'" + - "normalize('') returns ''" + - "WORD_FILTER contains common English stopwords (e.g. 'the', 'a', 'is')" + - "Unit tests cover edge cases: empty string, leading/trailing spaces, multiple internal spaces" + completed_at: null + + - id: task-002 + title: "Implement dependency-based triple extractor function" + description: | + Implement `extract_dependency_triples(doc: spacy.tokens.Doc) -> set[tuple[str, str, str, float]]` + in `spacy_utils.py`. Walk the dependency tree for each sentence to find + nsubj/nsubjpass → root_verb → dobj/pobj patterns. Apply phrasal merging + (`merge_noun_chunks`, `merge_entities`) before extraction. Handle passive + constructions by swapping subject/object and appending `_BY` to the + relation. Yield `(head, relation_lemma, tail, confidence)` tuples. + status: pending + dependencies: + - task-001 + acceptance_criteria: + - "Active SVO sentence yields (subject, verb_lemma, object, conf) triple" + - "Passive sentence yields (object, verb_lemma_BY, subject, conf) triple" + - "Prepositional phrase yields secondary (object, prep_label, pobj, conf) triple" + - "Multi-token noun phrases are collapsed into single entity strings" + - "Unit tests use spacy.blank + manually crafted Doc objects (no model download required)" + completed_at: null + + - id: task-003 + title: "Implement linear triple extractor function" + description: | + Implement `extract_linear_triples(doc: spacy.tokens.Doc) -> set[tuple[str, str, str, float]]` + in `spacy_utils.py`. Scan NER entity pairs within the same sentence that + appear within 10 tokens of each other with no other NER entity between + them. Emit `(entity_a, "RELATED_TO", entity_b, 0.3)` for each such pair. + This catches appositions, list items, and implicit associations not + surfaced by the dependency tree. + status: pending + dependencies: + - task-001 + acceptance_criteria: + - "Two NER entities within 10 tokens in same sentence yield a RELATED_TO triple" + - "Two NER entities more than 10 tokens apart produce no triple" + - "Three NER entities A–B–C produce A-RELATED_TO-B and B-RELATED_TO-C but not A-RELATED_TO-C (intervening entity)" + - "Confidence on all linear triples is exactly 0.3" + - "Unit tests use synthetic spacy Doc with manually set ents" + completed_at: null + + - id: task-004 + title: "Implement SpacyEntityRelationExtractor component" + description: | + Create `src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py` + with class `SpacyEntityRelationExtractor(EntityRelationExtractor)`. The + constructor loads the SpaCy model once (raising ValueError with an install + hint if not found). The async `run()` method processes each chunk, + unions dependency and linear triples (when enabled), normalises and + filters them, maps results to Neo4jNode/Neo4jRelationship objects, and + delegates lexical graph construction to LexicalGraphBuilder when + `create_lexical_graph=True`. + status: pending + dependencies: + - task-002 + - task-003 + acceptance_criteria: + - "Class is importable from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor" + - "Constructor with non-installed model name raises ValueError containing 'python -m spacy download'" + - "run() returns Neo4jGraph with nodes labelled 'Entity' and relationships with 'confidence' property" + - "use_linear_extractor=False produces no RELATED_TO relationships" + - "create_lexical_graph=True adds Chunk and Document nodes to the output graph" + - "on_error=OnError.IGNORE suppresses per-chunk exceptions; on_error=OnError.RAISE propagates them" + completed_at: null + + - id: task-005 + title: "Unit tests for SpacyEntityRelationExtractor" + description: | + Add `tests/unit/experimental/components/test_spacy_entity_relation_extractor.py`. + Mock `spacy.load` and the nlp pipeline so no model download is required. + Cover: successful extraction from two-chunk input, linear extractor toggle, + lexical graph nodes, invalid model error, and on_error behaviour. + status: pending + dependencies: + - task-004 + acceptance_criteria: + - "All tests pass with pytest without a live Neo4j instance or network calls" + - "spacy.load is mocked; no en_core_web_sm download required" + - "Test covers use_linear_extractor=True and False" + - "Test covers create_lexical_graph=True producing Chunk/Document nodes" + - "Test covers ValueError raised for missing model" + - "Test covers on_error=OnError.IGNORE swallowing a per-chunk exception" + completed_at: null + + - id: task-006 + title: "Implement HierarchicalTextSplitter component" + description: | + Create `src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py` + with class `HierarchicalTextSplitter(TextSplitter)`. Split text at + section-header regex boundaries first; recursively split sections + exceeding `max_chunk_size` chars with `chunk_overlap`; optionally load a + SpaCy model to drop verb-free sentences before emitting each chunk. + Output `TextChunks` with sequential index values. + status: pending + dependencies: + - task-001 + acceptance_criteria: + - "Class is importable from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter" + - "Text with a section header produces two separate chunks (one per section)" + - "Section longer than max_chunk_size is split into multiple chunks with chunk_overlap character overlap" + - "Section shorter than max_chunk_size is emitted as a single chunk without splitting" + - "drop_verbless_sentences=True removes sentences containing no verb token (requires spacy model)" + - "Output TextChunks has sequential index 0, 1, 2, ..." + completed_at: null + + - id: task-007 + title: "Unit tests for HierarchicalTextSplitter" + description: | + Add `tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py`. + Test section splitting, recursive character splitting with overlap, and + index sequencing. For drop_verbless_sentences tests, mock the SpaCy + pipeline so no model download is needed. + status: pending + dependencies: + - task-006 + acceptance_criteria: + - "All tests pass without network calls or model downloads" + - "Test verifies correct chunk_overlap: last N chars of chunk K equal first N chars of chunk K+1" + - "Test verifies section boundary respected: two Markdown-header sections produce separate chunks" + - "Test verifies verbless sentence dropped when drop_verbless_sentences=True (spacy mocked)" + - "Test verifies chunk indices are sequential starting from 0" + completed_at: null + + - id: task-008 + title: "Integration tests for both components against real SpaCy model" + description: | + Add integration tests in + `tests/unit/experimental/components/test_spacy_extractor_integration.py` + and `tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py`. + Use pytest fixtures that skip the test if `en_core_web_sm` is not installed. + Verify that real SpaCy output produces a populated Neo4jGraph and correct + TextChunks respectively. + status: pending + dependencies: + - task-005 + - task-007 + acceptance_criteria: + - "Tests are skipped (not failed) when en_core_web_sm is not installed" + - "With model installed, SpacyEntityRelationExtractor.run() returns Neo4jGraph with at least one node and one relationship for a 2-sentence input" + - "With model installed, HierarchicalTextSplitter.run() returns correct chunks for a multi-section document" + - "Confidence values on relationships are floats in [0, 1]" + completed_at: null + + - id: task-009 + title: "Export new components from experimental __init__" + description: | + Add `SpacyEntityRelationExtractor` and `HierarchicalTextSplitter` to the + public exports in `src/neo4j_graphrag/experimental/components/__init__.py` + and the text_splitters sub-package `__init__.py`. Guard imports with a + try/except so the rest of the package still imports cleanly when SpaCy is + not installed. + status: pending + dependencies: + - task-004 + - task-006 + acceptance_criteria: + - "from neo4j_graphrag.experimental.components import SpacyEntityRelationExtractor works when spacy is installed" + - "from neo4j_graphrag.experimental.components.text_splitters import HierarchicalTextSplitter works when spacy is installed" + - "Importing neo4j_graphrag.experimental.components without spacy installed does not raise ImportError" + completed_at: null + + - id: task-010 + title: "Add example script for end-to-end spacy pipeline" + description: | + Create `examples/customize/build_graph/components/spacy_kg_pipeline.py` + demonstrating a full pipeline: + DataLoader → HierarchicalTextSplitter → SpacyEntityRelationExtractor → KGWriter. + Include a comment block explaining the nlp extra install step + (`pip install 'neo4j-graphrag[nlp]'` + `python -m spacy download en_core_web_sm`). + status: pending + dependencies: + - task-009 + acceptance_criteria: + - "Script is syntactically valid Python" + - "Script includes install instructions in a comment block at the top" + - "Script constructs all four pipeline components and wires them together" + - "Script includes a __main__ guard and can be run standalone" + completed_at: null diff --git a/src/neo4j_graphrag/experimental/components/spacy_utils.py b/src/neo4j_graphrag/experimental/components/spacy_utils.py new file mode 100644 index 000000000..2eba88ae8 --- /dev/null +++ b/src/neo4j_graphrag/experimental/components/spacy_utils.py @@ -0,0 +1,161 @@ +# 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. +"""Shared utility helpers for SpaCy-based components. + +This module is intentionally kept free of SpaCy imports so that it can be +imported safely even when SpaCy is not installed. +""" + +from __future__ import annotations + +import re +import unicodedata + +# --------------------------------------------------------------------------- +# Text normalisation +# --------------------------------------------------------------------------- + +_WHITESPACE_RE = re.compile(r"\s+") + + +def normalize(text: str) -> str: + """Lowercase *text*, apply NFC Unicode normalisation, and collapse whitespace. + + Steps applied in order: + + 1. NFC Unicode normalisation — ensures that visually identical characters + with different code-point sequences compare equal (e.g. ``"caf\\u00e9"`` + and ``"cafe\\u0301"`` both become ``"café"`` before lowercasing). + 2. Lowercase via :meth:`str.lower`. + 3. Collapse all runs of whitespace (spaces, tabs, newlines …) to a single + ASCII space. + 4. Strip leading and trailing whitespace. + + Examples:: + + >>> normalize(" Supply Chain ") + 'supply chain' + >>> normalize("Hello World") + 'hello world' + >>> normalize("") + '' + """ + nfc = unicodedata.normalize("NFC", text) + return _WHITESPACE_RE.sub(" ", nfc.lower()).strip() + + +# --------------------------------------------------------------------------- +# Stop-word filter +# --------------------------------------------------------------------------- + +#: Common English stopwords that are excluded when normalising entity names +#: used as positions / keys in knowledge-graph nodes. Keeping the set small +#: and stable avoids accidentally filtering out legitimate entity fragments. +WORD_FILTER: frozenset[str] = frozenset( + { + "a", + "an", + "the", + "is", + "are", + "was", + "were", + "be", + "been", + "being", + "have", + "has", + "had", + "do", + "does", + "did", + "will", + "would", + "could", + "should", + "may", + "might", + "shall", + "can", + "of", + "in", + "to", + "for", + "on", + "at", + "by", + "from", + "with", + "and", + "or", + "but", + "not", + "no", + "nor", + "so", + "yet", + "both", + "either", + "neither", + "each", + "this", + "that", + "these", + "those", + "i", + "me", + "my", + "we", + "our", + "you", + "your", + "he", + "him", + "his", + "she", + "her", + "it", + "its", + "they", + "them", + "their", + "what", + "which", + "who", + "whom", + "whose", + "when", + "where", + "why", + "how", + "all", + "any", + "some", + "more", + "most", + "other", + "such", + "than", + "then", + "as", + "if", + "about", + "up", + "out", + "into", + "over", + "after", + } +) diff --git a/tests/unit/experimental/components/test_spacy_utils.py b/tests/unit/experimental/components/test_spacy_utils.py new file mode 100644 index 000000000..e712fb423 --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_utils.py @@ -0,0 +1,93 @@ +# 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. +"""Unit tests for spacy_utils.normalize() and WORD_FILTER.""" + +import pytest + +from neo4j_graphrag.experimental.components.spacy_utils import WORD_FILTER, normalize + + +class TestNormalize: + def test_basic_lowercase(self) -> None: + assert normalize("Hello World") == "hello world" + + def test_leading_trailing_spaces(self) -> None: + assert normalize(" Supply Chain ") == "supply chain" + + def test_multiple_internal_spaces(self) -> None: + assert normalize("Hello World") == "hello world" + + def test_empty_string(self) -> None: + assert normalize("") == "" + + def test_whitespace_only(self) -> None: + assert normalize(" ") == "" + + def test_tabs_and_newlines(self) -> None: + assert normalize("foo\t\nbar") == "foo bar" + + def test_already_normalized(self) -> None: + assert normalize("supply chain") == "supply chain" + + def test_mixed_case_with_spaces(self) -> None: + assert normalize(" UPPER lower MiXeD ") == "upper lower mixed" + + def test_single_word(self) -> None: + assert normalize("Apple") == "apple" + + def test_numbers_unchanged(self) -> None: + assert normalize(" 123 456 ") == "123 456" + + def test_unicode_nfc_normalization(self) -> None: + # "café" composed (U+00E9) vs decomposed (e + combining accent U+0301) + composed = "café" # NFC form + decomposed = "café" # NFD form — visually identical + assert normalize(composed) == normalize(decomposed) + + def test_unicode_lowercase_accent(self) -> None: + assert normalize("Ångström") == "ångström" + + +class TestWordFilter: + def test_is_frozenset(self) -> None: + assert isinstance(WORD_FILTER, frozenset) + + def test_contains_the(self) -> None: + assert "the" in WORD_FILTER + + def test_contains_a(self) -> None: + assert "a" in WORD_FILTER + + def test_contains_is(self) -> None: + assert "is" in WORD_FILTER + + def test_contains_and(self) -> None: + assert "and" in WORD_FILTER + + def test_contains_of(self) -> None: + assert "of" in WORD_FILTER + + def test_does_not_contain_entity_words(self) -> None: + # Common entity-like words should not be in the filter + for word in ("apple", "google", "london", "supply", "chain"): + assert word not in WORD_FILTER + + def test_all_entries_are_lowercase(self) -> None: + for word in WORD_FILTER: + assert word == word.lower(), f"WORD_FILTER entry '{word}' is not lowercase" + + def test_immutable(self) -> None: + with pytest.raises(AttributeError): + WORD_FILTER.add("new_word") # type: ignore[attr-defined] From 0549fb0ee1321bfa217f2af45fae0434a84bdb7e Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:05:37 +0200 Subject: [PATCH 02/12] spacy-entity-relation-extractor-task-003: implement linear triple extractor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `extract_linear_triples(doc)` in a new module `spacy_linear_extractor.py` (separated from `spacy_utils.py` because task-002 owns that file concurrently). The function emits (entity_a, "RELATED_TO", entity_b, 0.3) for every pair of NER entities in the same sentence whose starts are ≤ 10 tokens apart with no other entity between them. It gracefully degrades when no sentence segmenter has been run by treating the whole doc as a single sentence. 14 unit tests covering all acceptance criteria pass. --- ...ogress-spacy-entity-relation-extractor.txt | 5 + .../tasks-spacy-entity-relation-extractor.yml | 54 +++-- .../components/spacy_linear_extractor.py | 125 +++++++++++ .../components/test_spacy_linear_extractor.py | 210 ++++++++++++++++++ 4 files changed, 378 insertions(+), 16 deletions(-) create mode 100644 src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py create mode 100644 tests/unit/experimental/components/test_spacy_linear_extractor.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index f740385ee..0f4092bd5 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -1,3 +1,8 @@ +## TASK-003 — 2026-06-24T10:00:00Z +Implemented `extract_linear_triples(doc)` in a new file (task-002 was in_progress on spacy_utils.py) with 14 unit tests — all green. +Files: +src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py, +tests/unit/experimental/components/test_spacy_linear_extractor.py +Decisions: Separate file because task-002 owned spacy_utils.py concurrently; graceful fallback when no sentence segmenter is present (treats whole doc as one sentence); distance measured start-to-start. + ## TASK-001 — 2026-06-24T08:00:00Z Created `spacy_utils.py` with `normalize()` (lowercases, NFC-normalises Unicode, collapses whitespace) and `WORD_FILTER` frozenset of ~90 common English stopwords; 21 unit tests all green. Files: +src/neo4j_graphrag/experimental/components/spacy_utils.py, +tests/unit/experimental/components/test_spacy_utils.py diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index c63093c71..2bb7d6960 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -31,7 +31,7 @@ tasks: (`merge_noun_chunks`, `merge_entities`) before extraction. Handle passive constructions by swapping subject/object and appending `_BY` to the relation. Yield `(head, relation_lemma, tail, confidence)` tuples. - status: pending + status: in_progress dependencies: - task-001 acceptance_criteria: @@ -51,7 +51,8 @@ tasks: them. Emit `(entity_a, "RELATED_TO", entity_b, 0.3)` for each such pair. This catches appositions, list items, and implicit associations not surfaced by the dependency tree. - status: pending + status: completed + completed_at: "2026-06-24T10:00:00Z" dependencies: - task-001 acceptance_criteria: @@ -60,7 +61,7 @@ tasks: - "Three NER entities A–B–C produce A-RELATED_TO-B and B-RELATED_TO-C but not A-RELATED_TO-C (intervening entity)" - "Confidence on all linear triples is exactly 0.3" - "Unit tests use synthetic spacy Doc with manually set ents" - completed_at: null + note: "Implemented in spacy_linear_extractor.py (separate from spacy_utils.py because task-002 was in_progress on that file). task-004 should import from both spacy_utils and spacy_linear_extractor." - id: task-004 title: "Implement SpacyEntityRelationExtractor component" @@ -68,11 +69,18 @@ tasks: Create `src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py` with class `SpacyEntityRelationExtractor(EntityRelationExtractor)`. The constructor loads the SpaCy model once (raising ValueError with an install - hint if not found). The async `run()` method processes each chunk, - unions dependency and linear triples (when enabled), normalises and - filters them, maps results to Neo4jNode/Neo4jRelationship objects, and - delegates lexical graph construction to LexicalGraphBuilder when - `create_lexical_graph=True`. + hint if not found) and accepts `word_filter: Optional[list[str]] = None` + (uses WORD_FILTER constant if None). The async `run()` method accepts + `schema: Optional[GraphSchema] = None` (same as LLMEntityRelationExtractor). + When schema is provided, apply two filters: (1) type-level NER filter — + map schema node_type labels to SpaCy NER entity types and keep only + matching entities; (2) post-filter relations — only keep predicates that + match a schema relationship_type name. Without schema, all extracted + triples pass through. Unions dependency and linear triples (when enabled), + normalises and filters them, maps results to Neo4jNode/Neo4jRelationship + objects, and delegates lexical graph construction to LexicalGraphBuilder + when `create_lexical_graph=True`. Confidence scores go in relationship + properties. status: pending dependencies: - task-002 @@ -80,7 +88,10 @@ tasks: acceptance_criteria: - "Class is importable from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor" - "Constructor with non-installed model name raises ValueError containing 'python -m spacy download'" + - "Constructor accepts word_filter list; custom list overrides WORD_FILTER default" - "run() returns Neo4jGraph with nodes labelled 'Entity' and relationships with 'confidence' property" + - "run() with schema filters out entities whose NER type has no matching schema node_type" + - "run() with schema filters out relationships whose predicate has no matching schema relationship_type" - "use_linear_extractor=False produces no RELATED_TO relationships" - "create_lexical_graph=True adds Chunk and Document nodes to the output graph" - "on_error=OnError.IGNORE suppresses per-chunk exceptions; on_error=OnError.RAISE propagates them" @@ -109,20 +120,31 @@ tasks: title: "Implement HierarchicalTextSplitter component" description: | Create `src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py` - with class `HierarchicalTextSplitter(TextSplitter)`. Split text at - section-header regex boundaries first; recursively split sections - exceeding `max_chunk_size` chars with `chunk_overlap`; optionally load a - SpaCy model to drop verb-free sentences before emitting each chunk. - Output `TextChunks` with sequential index values. - status: pending + with class `HierarchicalTextSplitter(TextSplitter)`. Constructor params: + `max_chunk_size: int = 2048`, `chunk_overlap: int = 200`, + `header_strategy: str = "markdown"` (one of: "markdown" — lines starting + with `#`; "capitalization" — short Title Case or ALL_CAPS lines without + terminal punctuation; "blank_line" — short lines surrounded by blank + lines; "spacy_verbless" — SpaCy-detected verbless sentences that look + like headers), `model: str = "en_core_web_sm"` (only loaded when + header_strategy="spacy_verbless" or drop_verbless_sentences=True), + `drop_verbless_sentences: bool = True`. The "markdown" strategy is the + right default for MarkdownLoader output. Use "capitalization" or + "blank_line" for plain-text output from loaders like LiteParseLoader + (which returns plain text, not markdown). Recursively split sections + exceeding max_chunk_size with chunk_overlap. Output TextChunks with + sequential index values. + status: in_progress dependencies: - task-001 acceptance_criteria: - "Class is importable from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter" - - "Text with a section header produces two separate chunks (one per section)" + - "header_strategy='markdown' splits on # header lines" + - "header_strategy='capitalization' splits on short ALL_CAPS or Title Case lines without terminal punctuation" + - "header_strategy='blank_line' splits on short lines surrounded by blank lines" - "Section longer than max_chunk_size is split into multiple chunks with chunk_overlap character overlap" - "Section shorter than max_chunk_size is emitted as a single chunk without splitting" - - "drop_verbless_sentences=True removes sentences containing no verb token (requires spacy model)" + - "drop_verbless_sentences=True removes sentences containing no verb token (SpaCy mocked in unit tests)" - "Output TextChunks has sequential index 0, 1, 2, ..." completed_at: null diff --git a/src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py b/src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py new file mode 100644 index 000000000..264c224b3 --- /dev/null +++ b/src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py @@ -0,0 +1,125 @@ +# 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. +"""Linear (proximity-based) triple extractor for SpaCy documents. + +This module implements :func:`extract_linear_triples`, which emits +``(entity_a, "RELATED_TO", entity_b, 0.3)`` triples for every pair of NER +entities that: + +1. Appear in the **same sentence**. +2. Are within **10 tokens** of each other (measured from the *start* token of + each span). +3. Have **no other NER entity** between them. + +This surface-level heuristic captures appositions, list items, and implicit +co-occurrence associations that the dependency-based extractor may miss. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import spacy.tokens + +__all__ = ["extract_linear_triples"] + +# Confidence assigned to all proximity-based triples. +_LINEAR_CONFIDENCE: float = 0.3 + +# Maximum token distance (inclusive) between entity span *starts* for a pair +# to be considered "close enough" to emit a triple. +_MAX_TOKEN_DISTANCE: int = 10 + + +def extract_linear_triples( + doc: "spacy.tokens.Doc", +) -> set[tuple[str, str, str, float]]: + """Return proximity-based triples for co-occurring NER entity pairs. + + For each sentence in *doc* the function scans all ordered pairs of named + entities ``(a, b)`` where ``a`` appears before ``b``. A triple + ``(a.text, "RELATED_TO", b.text, 0.3)`` is emitted when: + + * The distance from ``a.start`` to ``b.start`` is ≤ 10 tokens. + * No other named entity starts between ``a.end`` and ``b.start``. + + Both conditions are evaluated on token indices within the document + (``spacy.tokens.Span.start`` / ``.end``). + + Parameters + ---------- + doc: + A processed :class:`spacy.tokens.Doc`. Named entities must have been + populated (``doc.ents``). + + Returns + ------- + set[tuple[str, str, str, float]] + A set of ``(entity_a_text, relation, entity_b_text, confidence)`` + tuples. Each pair is emitted **at most once** (the set deduplicates). + + Examples + -------- + >>> import spacy + >>> nlp = spacy.blank("en") + >>> # … build a Doc with .ents set … + >>> triples = extract_linear_triples(doc) + >>> ("Apple", "RELATED_TO", "Google", 0.3) in triples + True + """ + triples: set[tuple[str, str, str, float]] = set() + + # Determine sentence boundaries. When no sentence segmenter has been run + # (e.g. ``spacy.blank`` without a sentencizer), ``doc.has_annotation("SENT_START")`` + # is False and iterating ``doc.sents`` raises a ValueError. In that case we + # treat the whole document as a single sentence. + try: + sentences = list(doc.sents) + except ValueError: + # No sentence boundary annotation — treat the whole doc as one sentence. + sentences = [doc[:]] + + for sent in sentences: + # Collect entities that fall within this sentence, preserving order. + sent_ents = [ + ent for ent in doc.ents if ent.start >= sent.start and ent.end <= sent.end + ] + + n = len(sent_ents) + for i in range(n): + ent_a = sent_ents[i] + for j in range(i + 1, n): + ent_b = sent_ents[j] + + # ── Distance check ──────────────────────────────────────────── + distance = ent_b.start - ent_a.start + if distance > _MAX_TOKEN_DISTANCE: + # Entities are sorted by position; further pairs will only + # be farther away, so we can break the inner loop early. + break + + # ── Intervening-entity check ────────────────────────────────── + # Entities between ent_a and ent_b are those at indices i+1 … + # j-1 in the already-sorted sent_ents list. + has_intervening = j > i + 1 # True when there is any entity in between + + if has_intervening: + # Skip this pair: another entity sits between them. + continue + + triples.add((ent_a.text, "RELATED_TO", ent_b.text, _LINEAR_CONFIDENCE)) + + return triples diff --git a/tests/unit/experimental/components/test_spacy_linear_extractor.py b/tests/unit/experimental/components/test_spacy_linear_extractor.py new file mode 100644 index 000000000..96267de5a --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_linear_extractor.py @@ -0,0 +1,210 @@ +# 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. +"""Unit tests for :func:`extract_linear_triples`. + +All tests use ``spacy.blank("en")`` with manually constructed entity spans so +that no trained SpaCy model download is required. +""" + +from __future__ import annotations + +import pytest + +spacy = pytest.importorskip("spacy", reason="spacy is not installed") + +from neo4j_graphrag.experimental.components.spacy_linear_extractor import ( # noqa: E402 + extract_linear_triples, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def make_doc( + words: list[str], ent_spans: list[tuple[int, int, str]] +) -> spacy.tokens.Doc: + """Build a ``spacy.tokens.Doc`` from a word list and entity span specs. + + Parameters + ---------- + words: + Token strings, e.g. ``["Apple", "bought", "Google", "."]``. + ent_spans: + List of ``(start_token, end_token, label)`` tuples defining named + entities (``end_token`` is exclusive). + """ + nlp = spacy.blank("en") + doc = spacy.tokens.Doc(nlp.vocab, words=words) + spans = [] + for start, end, label in ent_spans: + span = doc[start:end] + span.label_ = label + spans.append(span) + doc.ents = spans # type: ignore[assignment] + return doc + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestExtractLinearTriplesBasic: + """Fundamental behaviour: pairs within 10 tokens, no intervening entity.""" + + def test_two_adjacent_entities_emit_triple(self) -> None: + """Two consecutive NER entities in the same sentence yield one triple.""" + # "Apple Google ." + words = ["Apple", "Google", "."] + doc = make_doc(words, [(0, 1, "ORG"), (1, 2, "ORG")]) + triples = extract_linear_triples(doc) + assert ("Apple", "RELATED_TO", "Google", 0.3) in triples + + def test_triple_confidence_is_exactly_0_3(self) -> None: + """Confidence on all linear triples must be exactly 0.3.""" + words = ["Apple", "and", "Google", "."] + doc = make_doc(words, [(0, 1, "ORG"), (2, 3, "ORG")]) + triples = extract_linear_triples(doc) + for _, _, _, conf in triples: + assert conf == pytest.approx(0.3) + + def test_entities_within_10_tokens_emit_triple(self) -> None: + """Entities whose start tokens differ by exactly 10 yield a triple.""" + # tokens 0..11, entities at 0 and 10 → distance = 10 (inclusive) + words = ["A"] + ["w"] * 9 + ["B", "."] + doc = make_doc(words, [(0, 1, "ORG"), (10, 11, "ORG")]) + triples = extract_linear_triples(doc) + assert ("A", "RELATED_TO", "B", 0.3) in triples + + def test_entities_more_than_10_tokens_apart_produce_no_triple(self) -> None: + """Two NER entities more than 10 tokens apart produce no triple.""" + # entities at token 0 and 11 → distance = 11 > 10 + words = ["A"] + ["w"] * 10 + ["B", "."] + doc = make_doc(words, [(0, 1, "ORG"), (11, 12, "ORG")]) + triples = extract_linear_triples(doc) + assert len(triples) == 0 + + def test_empty_doc_returns_empty_set(self) -> None: + words = ["."] + doc = make_doc(words, []) + assert extract_linear_triples(doc) == set() + + def test_single_entity_returns_empty_set(self) -> None: + words = ["Apple", "."] + doc = make_doc(words, [(0, 1, "ORG")]) + assert extract_linear_triples(doc) == set() + + def test_multi_token_entity_start_to_start_distance(self) -> None: + """Distance is measured start-to-start: a two-token entity at [0,2) and + a one-token entity at [8,9) have distance 8 (≤ 10) and should yield a triple.""" + # "New York w w w w w w London ." + # ent_a: tokens 0-1 ("New York"), ent_b: token 8 ("London") + # distance = 8 - 0 = 8 ≤ 10 → triple emitted + words = ["New", "York"] + ["w"] * 6 + ["London", "."] + doc = make_doc(words, [(0, 2, "GPE"), (8, 9, "GPE")]) + triples = extract_linear_triples(doc) + assert ("New York", "RELATED_TO", "London", 0.3) in triples + + +class TestExtractLinearTriplesInterveningEntity: + """Intervening entity blocks the pair.""" + + def test_three_entities_a_b_c_yields_only_adjacent_pairs(self) -> None: + """A–B–C within 10 tokens: A-RELATED_TO-B and B-RELATED_TO-C but NOT A-RELATED_TO-C.""" + # tokens: Apple(0) bought(1) Google(2) from(3) Microsoft(4) . + words = ["Apple", "bought", "Google", "from", "Microsoft", "."] + doc = make_doc(words, [(0, 1, "ORG"), (2, 3, "ORG"), (4, 5, "ORG")]) + triples = extract_linear_triples(doc) + + assert ("Apple", "RELATED_TO", "Google", 0.3) in triples + assert ("Google", "RELATED_TO", "Microsoft", 0.3) in triples + # A-C must NOT be present because Google intervenes + assert ("Apple", "RELATED_TO", "Microsoft", 0.3) not in triples + + def test_intervening_entity_blocks_pair(self) -> None: + """Two entities with an intervening one are NOT directly paired.""" + words = ["A", "x", "B", "x", "C", "."] + doc = make_doc(words, [(0, 1, "ORG"), (2, 3, "ORG"), (4, 5, "ORG")]) + triples = extract_linear_triples(doc) + assert ("A", "RELATED_TO", "C", 0.3) not in triples + + +class TestExtractLinearTriplesSentenceBoundary: + """Entities in different sentences must not be paired.""" + + def test_entities_in_different_sentences_produce_no_triple(self) -> None: + """Entities separated by a sentence boundary are not paired.""" + nlp = spacy.blank("en") + # Two sentences: "Apple ." and "Google ." + doc = spacy.tokens.Doc( + nlp.vocab, + words=["Apple", ".", "Google", "."], + sent_starts=[True, False, True, False], + ) + span_a = doc[0:1] + span_a.label_ = "ORG" + span_b = doc[2:3] + span_b.label_ = "ORG" + doc.ents = [span_a, span_b] + + triples = extract_linear_triples(doc) + assert ("Apple", "RELATED_TO", "Google", 0.3) not in triples + assert len(triples) == 0 + + def test_entities_same_sentence_still_paired(self) -> None: + """Entities in the same sentence ARE paired regardless of other sentences.""" + nlp = spacy.blank("en") + # Sentence 1: "X ." Sentence 2: "Apple and Google ." + doc = spacy.tokens.Doc( + nlp.vocab, + words=["X", ".", "Apple", "Google", "."], + sent_starts=[True, False, True, False, False], + ) + span_a = doc[2:3] + span_a.label_ = "ORG" + span_b = doc[3:4] + span_b.label_ = "ORG" + doc.ents = [span_a, span_b] + + triples = extract_linear_triples(doc) + assert ("Apple", "RELATED_TO", "Google", 0.3) in triples + + +class TestExtractLinearTriplesReturnType: + """Return value is always a set of 4-tuples.""" + + def test_return_type_is_set(self) -> None: + words = ["Apple", "Google", "."] + doc = make_doc(words, [(0, 1, "ORG"), (1, 2, "ORG")]) + result = extract_linear_triples(doc) + assert isinstance(result, set) + + def test_tuples_have_four_elements(self) -> None: + words = ["Apple", "Google", "."] + doc = make_doc(words, [(0, 1, "ORG"), (1, 2, "ORG")]) + result = extract_linear_triples(doc) + for triple in result: + assert len(triple) == 4 + + def test_deduplication(self) -> None: + """Identical entity texts produce at most one triple in the set.""" + words = ["Apple", "Google", "."] + doc = make_doc(words, [(0, 1, "ORG"), (1, 2, "ORG")]) + result = extract_linear_triples(doc) + # Set semantics guarantee no duplicates by definition; count directly. + apple_google = [t for t in result if t[0] == "Apple" and t[2] == "Google"] + assert len(apple_google) == 1 From 0ac91852136b1d3aa15adcea4e1bb3c384eb8943 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:09:58 +0200 Subject: [PATCH 03/12] spacy-entity-relation-extractor-task-002: implement dependency-based triple extractor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `extract_dependency_triples(doc)` to `spacy_utils.py`. Walks the dependency tree per sentence finding nsubj/nsubjpass → ROOT → dobj/pobj patterns. Handles active SVO (conf=1.0), passive voice with _BY suffix (conf=0.9), and prepositional triples (conf=0.7). Uses spaCy 3.x retokenizer API to collapse multi-word noun chunks/named entities. Falls back gracefully when sentence boundaries are unavailable. Move spaCy-gated tests to a dedicated `test_spacy_dependency_triples.py` so that normalize/WORD_FILTER tests stay runnable without spaCy installed. 18 new tests; all 39 tests pass. Co-Authored-By: Claude Sonnet 4.6 --- ...ogress-spacy-entity-relation-extractor.txt | 10 + .../tasks-spacy-entity-relation-extractor.yml | 7 +- .../experimental/components/spacy_utils.py | 184 ++++++- .../text_splitters/hierarchical_splitter.py | 452 ++++++++++++++++++ .../test_spacy_dependency_triples.py | 359 ++++++++++++++ .../components/test_spacy_utils.py | 12 +- uv.lock | 2 +- 7 files changed, 1017 insertions(+), 9 deletions(-) create mode 100644 src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py create mode 100644 tests/unit/experimental/components/test_spacy_dependency_triples.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index 0f4092bd5..4b2fd1f6f 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -1,3 +1,8 @@ +## TASK-006 — 2026-06-24T12:00:00Z +Created `HierarchicalTextSplitter(TextSplitter)` with four header strategies (markdown, capitalization, blank_line, spacy_verbless), recursive overlap splitting for large sections, and optional SpaCy-based verbless-sentence filtering. Ruff, mypy, and existing unit tests all pass. +Files: +src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py +Decisions: blank_line strategy uses ≤6-word cap to avoid treating body sentences as section headers; chunk index assigned with single-pass `len(chunks)` pattern; negative chunk_overlap validated; assert replaced with RuntimeError guard; docstring warns that default drop_verbless_sentences=True requires SpaCy. + ## TASK-003 — 2026-06-24T10:00:00Z Implemented `extract_linear_triples(doc)` in a new file (task-002 was in_progress on spacy_utils.py) with 14 unit tests — all green. Files: +src/neo4j_graphrag/experimental/components/spacy_linear_extractor.py, +tests/unit/experimental/components/test_spacy_linear_extractor.py @@ -7,3 +12,8 @@ Decisions: Separate file because task-002 owned spacy_utils.py concurrently; gra Created `spacy_utils.py` with `normalize()` (lowercases, NFC-normalises Unicode, collapses whitespace) and `WORD_FILTER` frozenset of ~90 common English stopwords; 21 unit tests all green. Files: +src/neo4j_graphrag/experimental/components/spacy_utils.py, +tests/unit/experimental/components/test_spacy_utils.py Decisions: Added `unicodedata.normalize("NFC", ...)` step per reviewer feedback to ensure visually identical Unicode strings compare equal as graph node keys. + +## TASK-002 — 2026-06-24T10:00:00Z +Added `extract_dependency_triples(doc)` to `spacy_utils.py`: walks the dependency tree per sentence finding nsubj/nsubjpass → ROOT → dobj/pobj patterns, handles passive voice (_BY suffix, conf=0.9), prepositional triples (conf=0.7), and active SVO (conf=1.0). Uses spaCy 3.x retokenizer API (not deprecated `Span.merge()`) to collapse multi-word noun chunks and named entities. Falls back gracefully when sentence boundaries are unavailable. 18 unit tests in a new dedicated file; spaCy-free tests remain ungated. +Files: ~src/neo4j_graphrag/experimental/components/spacy_utils.py, ~tests/unit/experimental/components/test_spacy_utils.py, +tests/unit/experimental/components/test_spacy_dependency_triples.py +Decisions: Reviewer flagged `Span.merge()` removed in spaCy 3.x (silently swallowed); fixed to use `doc.retokenize()`. Moved spaCy-gated tests to separate file so normalize/WORD_FILTER tests stay runnable without spaCy installed. diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index 2bb7d6960..655f996b9 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -31,7 +31,8 @@ tasks: (`merge_noun_chunks`, `merge_entities`) before extraction. Handle passive constructions by swapping subject/object and appending `_BY` to the relation. Yield `(head, relation_lemma, tail, confidence)` tuples. - status: in_progress + status: completed + completed_at: "2026-06-24T10:00:00Z" dependencies: - task-001 acceptance_criteria: @@ -40,7 +41,6 @@ tasks: - "Prepositional phrase yields secondary (object, prep_label, pobj, conf) triple" - "Multi-token noun phrases are collapsed into single entity strings" - "Unit tests use spacy.blank + manually crafted Doc objects (no model download required)" - completed_at: null - id: task-003 title: "Implement linear triple extractor function" @@ -134,7 +134,8 @@ tasks: (which returns plain text, not markdown). Recursively split sections exceeding max_chunk_size with chunk_overlap. Output TextChunks with sequential index values. - status: in_progress + status: completed + completed_at: "2026-06-24T12:00:00Z" dependencies: - task-001 acceptance_criteria: diff --git a/src/neo4j_graphrag/experimental/components/spacy_utils.py b/src/neo4j_graphrag/experimental/components/spacy_utils.py index 2eba88ae8..a84390294 100644 --- a/src/neo4j_graphrag/experimental/components/spacy_utils.py +++ b/src/neo4j_graphrag/experimental/components/spacy_utils.py @@ -14,14 +14,21 @@ # limitations under the License. """Shared utility helpers for SpaCy-based components. -This module is intentionally kept free of SpaCy imports so that it can be -imported safely even when SpaCy is not installed. +This module is intentionally kept free of heavy SpaCy imports at module level +so that it can be imported safely even when SpaCy is not installed. Functions +that accept ``spacy.tokens.Doc`` objects use ``TYPE_CHECKING`` guards so that +type checkers still see the correct annotations without causing an +``ImportError`` at runtime. """ from __future__ import annotations import re import unicodedata +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import spacy.tokens # --------------------------------------------------------------------------- # Text normalisation @@ -56,6 +63,179 @@ def normalize(text: str) -> str: return _WHITESPACE_RE.sub(" ", nfc.lower()).strip() +# --------------------------------------------------------------------------- +# Dependency-based triple extraction +# --------------------------------------------------------------------------- + +#: Dependency labels that mark nominal subjects (active voice). +_SUBJ_DEPS: frozenset[str] = frozenset({"nsubj"}) + +#: Dependency labels that mark nominal subjects in passive constructions. +_SUBJ_PASS_DEPS: frozenset[str] = frozenset({"nsubjpass"}) + +#: Dependency labels that mark direct objects. +_OBJ_DEPS: frozenset[str] = frozenset({"dobj", "obj"}) + + +def extract_dependency_triples( + doc: "spacy.tokens.Doc", +) -> set[tuple[str, str, str, float]]: + """Walk SpaCy's dependency tree and extract (head, relation, tail, conf) triples. + + A copy of *doc* is made before any retokenization so that the caller's + original ``Doc`` object is never mutated. Multi-token noun chunks and named + entities are merged using the spaCy 3.x retokenizer API + (``Doc.retokenize()``) so that spans such as *"Supply Chain Management"* + become a single token before extraction. + + Extraction logic per sentence + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + For every token whose dependency label is ``ROOT`` (the main verb): + + 1. Collect all ``nsubj`` children → *active subjects*. + 2. Collect all ``nsubjpass`` children → *passive subjects*. + 3. Collect all ``dobj`` / ``obj`` children → *direct objects*. + 4. For each ``(subject, object)`` pair emit a primary triple: + + - **Active**: ``(subject_text, verb_lemma, object_text, conf)`` + - **Passive**: ``(object_text, verb_lemma_BY, subject_text, conf)`` + where ``verb_lemma_BY = verb_lemma.upper() + "_BY"``. + + 5. For each ``prep`` child of the root verb, collect its ``pobj`` children + and emit a secondary triple: + + - When a direct object is present: + ``(direct_object_text, prep_label.upper(), pobj_text, 0.7)`` + - When no direct object is present: + ``(subject_text, prep_label.upper(), pobj_text, 0.7)`` + + Confidence scoring + ~~~~~~~~~~~~~~~~~~ + * Active SVO triples: ``1.0`` + * Passive triples: ``0.9`` + * Prepositional triples: ``0.7`` + + Parameters + ---------- + doc: + A SpaCy ``Doc`` object. The function always works on an internal copy; + the caller's doc is never modified. + + Returns + ------- + set[tuple[str, str, str, float]] + A set of ``(head, relation, tail, confidence)`` 4-tuples. Both *head* + and *tail* are token texts after retokenization (not normalised — + normalisation is the caller's responsibility). + """ + # Apply phrasal merging so that multi-word NPs become single tokens. + # We use the spaCy 3.x retokenizer API (Span.merge() was removed in 3.0). + # Work on a copy to avoid mutating the caller's doc. + try: + import spacy.tokens # noqa: F401 — ensure spacy is available + + doc = doc.copy() + + # Merge noun chunks (requires DEP annotation for sentence boundaries). + if doc.has_annotation("DEP"): + try: + with doc.retokenize() as retokenizer: + for span in list(doc.noun_chunks): + retokenizer.merge(span) + except Exception: + pass # noun-chunk merging is best-effort on a blank model + + # Merge named entities (works whenever .ents is populated). + if doc.ents: + try: + with doc.retokenize() as retokenizer: + for span in list(doc.ents): + retokenizer.merge(span) + except Exception: + pass # overlapping or already-merged spans; skip + + except ImportError: + pass # SpaCy not installed — proceed without merging + + triples: set[tuple[str, str, str, float]] = set() + + # Iterate sentences; fall back to treating the whole doc as one sentence + # if sentence boundaries are not available (no DEP annotation, no sentencizer). + try: + sents_iter = list(doc.sents) + except ValueError: + sents_iter = [doc[:]] + + for sent in sents_iter: + for token in sent: + if token.dep_ != "ROOT": + continue + + verb_lemma = (token.lemma_ or token.text).upper() + + # Collect subject tokens. + active_subjects = [c for c in token.children if c.dep_ in _SUBJ_DEPS] + passive_subjects = [c for c in token.children if c.dep_ in _SUBJ_PASS_DEPS] + direct_objects = [c for c in token.children if c.dep_ in _OBJ_DEPS] + + # Prepositional phrases hanging off the root verb. + prep_children = [c for c in token.children if c.dep_ == "prep"] + + # --- Active SVO triples ------------------------------------------- + for subj in active_subjects: + for obj in direct_objects: + triples.add((subj.text, verb_lemma, obj.text, 1.0)) + + # Secondary prep triples: (obj, PREP_LABEL, pobj) + for prep in prep_children: + for pobj in prep.children: + if pobj.dep_ == "pobj": + triples.add( + (obj.text, prep.text.upper(), pobj.text, 0.7) + ) + + # --- Passive voice triples ---------------------------------------- + # In passive, the grammatical subject is the logical object, and the + # prepositional "by" phrase (dep_ == "agent") contains the logical + # subject. We handle two cases: + # + # Case A: nsubjpass + dobj (uncommon ditransitive passive) + # Case B: nsubjpass alone — the "agent" (by-phrase) child provides + # the logical actor. + + for subj in passive_subjects: + passive_relation = verb_lemma + "_BY" + + # Case A: direct object still present (ditransitive passive) + for obj in direct_objects: + triples.add((obj.text, passive_relation, subj.text, 0.9)) + + # Case B: "by" agent provides the logical subject + for agent in (c for c in token.children if c.dep_ == "agent"): + for pobj in agent.children: + if pobj.dep_ == "pobj": + triples.add((subj.text, passive_relation, pobj.text, 0.9)) + + # Secondary prep triples for the passive subject + for prep in prep_children: + for pobj in prep.children: + if pobj.dep_ == "pobj": + triples.add((subj.text, prep.text.upper(), pobj.text, 0.7)) + + # --- Subject-verb-only (no dobj): preps hang off verb itself ------ + # When there is no direct object, emit (subject, PREP_LABEL, pobj). + if not direct_objects: + for subj in active_subjects + passive_subjects: + for prep in prep_children: + for pobj in prep.children: + if pobj.dep_ == "pobj": + triples.add( + (subj.text, prep.text.upper(), pobj.text, 0.7) + ) + + return triples + + # --------------------------------------------------------------------------- # Stop-word filter # --------------------------------------------------------------------------- diff --git a/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py b/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py new file mode 100644 index 000000000..2e8758b5b --- /dev/null +++ b/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py @@ -0,0 +1,452 @@ +# 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. +"""Hierarchical text splitter that detects section boundaries before chunking. + +Supported header strategies: + +* ``"markdown"`` — lines starting with ``#`` (ATX Markdown headers). +* ``"capitalization"`` — short Title Case or ALL_CAPS lines without terminal + punctuation (appropriate for plain-text output from loaders like + LiteParseLoader). +* ``"blank_line"`` — short lines that are surrounded by blank lines on both + sides (a common plain-text section marker). +* ``"spacy_verbless"`` — SpaCy-parsed sentences that are short, contain no + verb, and are followed by a longer sentence. + +All strategies produce a list of *sections* (contiguous text blocks). Each +section is then emitted as a single chunk when it fits within ``max_chunk_size``, +or recursively split with ``chunk_overlap`` when it is larger. + +Optionally, when ``drop_verbless_sentences=True`` (default), SpaCy is used to +remove sentences with no verb token from every emitted chunk. +""" + +from __future__ import annotations + +import re +from typing import Any, Optional + +from pydantic import validate_call + +from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter +from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks + +# --------------------------------------------------------------------------- +# Header-detection helpers +# --------------------------------------------------------------------------- + +_MARKDOWN_HEADER_RE = re.compile(r"^#{1,6}\s+\S", re.MULTILINE) + +# Title Case: starts with capital, no terminal punctuation at end of line, +# short (≤ 80 chars), and at least half the words are Title-cased. +_TITLECASE_RE = re.compile(r"^[A-Z][^\n]{0,78}$") +_TERMINAL_PUNCT_RE = re.compile(r"[.!?;,]$") + +# ALL_CAPS line: all uppercase letters / spaces / digits +_ALLCAPS_RE = re.compile(r"^[A-Z0-9][A-Z0-9 \t\-:]+$") + +# Blank-line boundary: a line that is non-empty, short, and preceded/followed +# by blank lines (handled at section-split level, not a single regex). +_SHORT_LINE_MAX = 80 + + +def _is_title_case(line: str) -> bool: + """Return True when a line looks like a Title Case heading.""" + if not _TITLECASE_RE.match(line): + return False + if _TERMINAL_PUNCT_RE.search(line): + return False + words = line.split() + if len(words) == 0: + return False + capitalised = sum(1 for w in words if w and w[0].isupper()) + return capitalised / len(words) >= 0.5 + + +def _is_allcaps(line: str) -> bool: + """Return True when a line is ALL CAPS without terminal punctuation.""" + if not _ALLCAPS_RE.match(line): + return False + return not _TERMINAL_PUNCT_RE.search(line) + + +def _split_at_markdown_headers(text: str) -> list[str]: + """Split *text* into sections at Markdown ATX header lines (``#``, ``##``, …).""" + # We split before each header line so that the header stays with its section. + lines = text.splitlines(keepends=True) + sections: list[str] = [] + current: list[str] = [] + for line in lines: + if re.match(r"^#{1,6}\s+\S", line) and current: + sections.append("".join(current)) + current = [line] + else: + current.append(line) + if current: + sections.append("".join(current)) + return [s for s in sections if s.strip()] + + +def _split_at_capitalization(text: str) -> list[str]: + """Split *text* into sections at Title Case or ALL_CAPS heading lines.""" + lines = text.splitlines(keepends=True) + sections: list[str] = [] + current: list[str] = [] + for line in lines: + stripped = line.rstrip("\r\n") + is_header = ( + len(stripped) <= _SHORT_LINE_MAX + and stripped.strip() + and (_is_title_case(stripped.strip()) or _is_allcaps(stripped.strip())) + ) + if is_header and current: + sections.append("".join(current)) + current = [line] + else: + current.append(line) + if current: + sections.append("".join(current)) + return [s for s in sections if s.strip()] + + +def _split_at_blank_line(text: str) -> list[str]: + """Split *text* into sections at short lines surrounded by blank lines. + + A line qualifies as a section header when it: + + * Is non-empty and at most 80 characters. + * Contains no terminal punctuation (not a regular sentence). + * Has no more than 6 words (avoids treating body sentences as headers). + * Is preceded by a blank line (or is the very first non-blank line). + * Is followed by a blank line. + """ + lines = text.splitlines(keepends=True) + n = len(lines) + sections: list[str] = [] + current: list[str] = [] + + i = 0 + while i < n: + line = lines[i] + stripped = line.rstrip("\r\n").strip() + # A "blank-line boundary" header: non-empty, short, preceded by blank, + # followed by blank, few words, no terminal punctuation. + prev_blank = (i == 0) or (not lines[i - 1].strip()) + next_blank = (i + 1 >= n) or (not lines[i + 1].strip()) + word_count = len(stripped.split()) if stripped else 0 + is_header = ( + stripped + and len(stripped) <= _SHORT_LINE_MAX + and not _TERMINAL_PUNCT_RE.search(stripped) + and word_count <= 6 + and prev_blank + and next_blank + ) + if is_header and current: + sections.append("".join(current)) + current = [line] + else: + current.append(line) + i += 1 + + if current: + sections.append("".join(current)) + return [s for s in sections if s.strip()] + + +def _split_at_spacy_verbless(text: str, nlp: Any) -> list[str]: + """Split *text* into sections at SpaCy-detected verbless heading sentences. + + A sentence qualifies as a heading when it: + + * Has at most 80 characters. + * Contains no verb token (POS tag ``VERB`` or ``AUX``). + * Is immediately followed by a longer sentence (> 80 chars). + """ + doc = nlp(text) + sentences = list(doc.sents) + if not sentences: + return [text] if text.strip() else [] + + # Identify which sentences are "headers". + is_header = [False] * len(sentences) + for idx, sent in enumerate(sentences): + sent_text = sent.text.strip() + if len(sent_text) > _SHORT_LINE_MAX: + continue + has_verb = any(tok.pos_ in ("VERB", "AUX") for tok in sent) + if has_verb: + continue + # Must be followed by a longer sentence. + if ( + idx + 1 < len(sentences) + and len(sentences[idx + 1].text.strip()) > _SHORT_LINE_MAX + ): + is_header[idx] = True + + # Build sections: split before each header sentence (except the first). + sections: list[str] = [] + current_parts: list[str] = [] + for idx, sent in enumerate(sentences): + if is_header[idx] and current_parts: + sections.append(" ".join(current_parts)) + current_parts = [sent.text] + else: + current_parts.append(sent.text) + + if current_parts: + sections.append(" ".join(current_parts)) + + return [s for s in sections if s.strip()] + + +# --------------------------------------------------------------------------- +# Overlap-based character splitter (for sections larger than max_chunk_size) +# --------------------------------------------------------------------------- + + +def _split_with_overlap(text: str, max_size: int, overlap: int) -> list[str]: + """Split *text* into character-level chunks of at most *max_size* with + *overlap* characters carried over from the previous chunk. + + Splits are attempted at whitespace boundaries to avoid cutting words. + """ + if not text: + return [] + chunks: list[str] = [] + start = 0 + length = len(text) + step = max(1, max_size - overlap) + + while start < length: + end = min(start + max_size, length) + chunk = text[start:end] + + # Prefer to cut at a whitespace boundary when not at end of text. + if end < length: + # Walk backwards to find a space. + cut = end + while cut > start and not text[cut - 1].isspace(): + cut -= 1 + if cut > start: + end = cut + chunk = text[start:end] + + chunks.append(chunk) + + # Advance by step, ensuring we always make progress. + next_start = start + step + if next_start <= start: + next_start = start + 1 + start = next_start + + return chunks + + +# --------------------------------------------------------------------------- +# Verb-filter using SpaCy +# --------------------------------------------------------------------------- + + +def _drop_verbless_sentences(text: str, nlp: Any) -> str: + """Remove sentences with no verb token from *text* using SpaCy. + + A sentence is considered *verbless* when it contains no token whose + part-of-speech tag is ``VERB`` or ``AUX``. + """ + doc = nlp(text) + kept: list[str] = [] + for sent in doc.sents: + has_verb = any(tok.pos_ in ("VERB", "AUX") for tok in sent) + if has_verb: + kept.append(sent.text) + return " ".join(kept) + + +# --------------------------------------------------------------------------- +# Valid strategy names +# --------------------------------------------------------------------------- + +_VALID_STRATEGIES = frozenset( + {"markdown", "capitalization", "blank_line", "spacy_verbless"} +) + + +# --------------------------------------------------------------------------- +# Main component +# --------------------------------------------------------------------------- + + +class HierarchicalTextSplitter(TextSplitter): + """Splits text by first detecting section boundaries then chunking each section. + + Args: + max_chunk_size (int): Maximum number of characters per output chunk. + Defaults to 2048. + chunk_overlap (int): Characters of overlap between consecutive chunks + when a section must be further split. Must be less than + ``max_chunk_size``. Defaults to 200. + header_strategy (str): How to detect section boundaries. One of: + + * ``"markdown"`` — Markdown ATX header lines (``#``, ``##``, …). + * ``"capitalization"`` — short Title Case or ALL_CAPS lines without + terminal punctuation. + * ``"blank_line"`` — short lines surrounded by blank lines on both + sides. + * ``"spacy_verbless"`` — SpaCy-detected short verbless sentences + that precede a longer sentence. + + model (str): SpaCy model name loaded when *header_strategy* is + ``"spacy_verbless"`` or *drop_verbless_sentences* is ``True``. + Defaults to ``"en_core_web_sm"``. + drop_verbless_sentences (bool): When ``True`` (default), SpaCy is used + to remove verbless sentences from every emitted chunk. Note that + this default value causes SpaCy to be loaded at construction time + regardless of the chosen *header_strategy* — install + ``neo4j-graphrag[nlp]`` when using the default, or explicitly pass + ``drop_verbless_sentences=False`` to avoid the SpaCy dependency. + + Example: + + .. code-block:: python + + from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( + HierarchicalTextSplitter, + ) + from neo4j_graphrag.experimental.pipeline import Pipeline + + pipeline = Pipeline() + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=200, + header_strategy="markdown", + ) + pipeline.add_component(splitter, "text_splitter") + """ + + @validate_call + def __init__( + self, + max_chunk_size: int = 2048, + chunk_overlap: int = 200, + header_strategy: str = "markdown", + model: str = "en_core_web_sm", + drop_verbless_sentences: bool = True, + ) -> None: + if max_chunk_size <= 0: + raise ValueError("max_chunk_size must be strictly greater than 0") + if chunk_overlap < 0: + raise ValueError("chunk_overlap must be >= 0") + if chunk_overlap >= max_chunk_size: + raise ValueError("chunk_overlap must be strictly less than max_chunk_size") + if header_strategy not in _VALID_STRATEGIES: + raise ValueError( + f"header_strategy must be one of {sorted(_VALID_STRATEGIES)}, " + f"got {header_strategy!r}" + ) + + self.max_chunk_size = max_chunk_size + self.chunk_overlap = chunk_overlap + self.header_strategy = header_strategy + self.model = model + self.drop_verbless_sentences = drop_verbless_sentences + + # Pre-load SpaCy only when needed. + self._nlp: Optional[Any] = None + needs_spacy = header_strategy == "spacy_verbless" or drop_verbless_sentences + if needs_spacy: + self._nlp = self._load_spacy(model) + + @staticmethod + def _load_spacy(model: str) -> Any: + """Load a SpaCy model, raising a clear error when SpaCy is missing.""" + try: + import spacy # noqa: PLC0415 + except ImportError as exc: + raise ImportError( + "SpaCy is required for this configuration of HierarchicalTextSplitter. " + "Install it with: pip install 'neo4j-graphrag[nlp]'" + ) from exc + try: + return spacy.load(model) + except OSError as exc: + raise ValueError( + f"SpaCy model {model!r} is not installed. " + f"Download it with: python -m spacy download {model}" + ) from exc + + def _detect_sections(self, text: str) -> list[str]: + """Detect section boundaries and return a list of section strings.""" + strategy = self.header_strategy + if strategy == "markdown": + sections = _split_at_markdown_headers(text) + elif strategy == "capitalization": + sections = _split_at_capitalization(text) + elif strategy == "blank_line": + sections = _split_at_blank_line(text) + else: # "spacy_verbless" + if self._nlp is None: + raise RuntimeError( + "SpaCy model not loaded for 'spacy_verbless' strategy; this is a bug" + ) + sections = _split_at_spacy_verbless(text, self._nlp) + + # Fallback: if no sections were detected, treat the whole text as one. + if not sections: + sections = [text] if text.strip() else [] + return sections + + def _chunk_section(self, section_text: str) -> list[str]: + """Return one or more raw text chunks for a single *section_text*.""" + if len(section_text) <= self.max_chunk_size: + return [section_text] + return _split_with_overlap( + section_text, self.max_chunk_size, self.chunk_overlap + ) + + def _filter_verbless(self, text: str) -> str: + """Apply the verbless-sentence filter if enabled.""" + if not self.drop_verbless_sentences or self._nlp is None: + return text + filtered = _drop_verbless_sentences(text, self._nlp) + # Fall back to original text when filtering removes everything. + return filtered if filtered.strip() else text + + @validate_call + async def run(self, text: str) -> TextChunks: + """Split *text* into hierarchical chunks. + + Args: + text (str): The text to be split. + + Returns: + TextChunks: A list of chunks with sequential index values starting + from 0. + """ + if not text.strip(): + return TextChunks(chunks=[]) + + sections = self._detect_sections(text) + raw_chunks: list[str] = [] + for section in sections: + raw_chunks.extend(self._chunk_section(section)) + + chunks: list[TextChunk] = [] + for raw in raw_chunks: + filtered = self._filter_verbless(raw) + if filtered.strip(): + chunks.append(TextChunk(text=filtered, index=len(chunks))) + + return TextChunks(chunks=chunks) diff --git a/tests/unit/experimental/components/test_spacy_dependency_triples.py b/tests/unit/experimental/components/test_spacy_dependency_triples.py new file mode 100644 index 000000000..9b56e2fef --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_dependency_triples.py @@ -0,0 +1,359 @@ +# 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. +"""Unit tests for extract_dependency_triples using synthetic spaCy Docs. + +spaCy is an optional dependency; these tests are skipped when it is not +installed. All Doc objects are built with ``spacy.blank("en")`` and manually +set dependency arcs so that no trained model needs to be downloaded. +""" + +import pytest + +spacy = pytest.importorskip("spacy", reason="spacy is not installed") + +from neo4j_graphrag.experimental.components.spacy_utils import ( # noqa: E402 + extract_dependency_triples, +) + + +# --------------------------------------------------------------------------- +# Helper: build a synthetic spacy Doc with dependency annotations +# --------------------------------------------------------------------------- + + +def _make_doc( + words: list[str], + lemmas: list[str], + dep_labels: list[str], + head_indices: list[int], + *, + spaces: list[bool] | None = None, +) -> "spacy.tokens.Doc": + """Return a ``spacy.tokens.Doc`` with manually set dependency arcs. + + All parameters must have the same length as *words*. ``head_indices`` + uses 0-based integer positions into *words*; the ROOT token must point to + itself (``head_indices[i] == i``). + """ + from spacy.tokens import Doc + + nlp = spacy.blank("en") + if spaces is None: + spaces = [True] * (len(words) - 1) + [False] + doc = Doc(nlp.vocab, words=words, spaces=spaces, lemmas=lemmas) + for i, (dep, head_idx) in enumerate(zip(dep_labels, head_indices)): + doc[i].dep_ = dep + doc[i].head = doc[head_idx] + return doc + + +# --------------------------------------------------------------------------- +# Tests for extract_dependency_triples +# --------------------------------------------------------------------------- + + +class TestExtractDependencyTriples: + """Tests for extract_dependency_triples using synthetic spacy Docs.""" + + # ------------------------------------------------------------------ + # Active SVO + # ------------------------------------------------------------------ + + def test_active_svo_yields_triple(self) -> None: + """'Apple acquired IBM' → (Apple, ACQUIRE, IBM, 1.0).""" + doc = _make_doc( + words=["Apple", "acquired", "IBM"], + lemmas=["Apple", "acquire", "IBM"], + dep_labels=["nsubj", "ROOT", "dobj"], + head_indices=[1, 1, 1], + ) + triples = extract_dependency_triples(doc) + assert ("Apple", "ACQUIRE", "IBM", 1.0) in triples + + def test_active_svo_confidence_is_one(self) -> None: + """Active SVO triples must have confidence == 1.0.""" + doc = _make_doc( + words=["Google", "bought", "YouTube"], + lemmas=["Google", "buy", "YouTube"], + dep_labels=["nsubj", "ROOT", "dobj"], + head_indices=[1, 1, 1], + ) + triples = extract_dependency_triples(doc) + for head, _rel, tail, conf in triples: + if head == "Google" and tail == "YouTube": + assert conf == 1.0 + + def test_active_svo_uses_verb_lemma(self) -> None: + """Relation string should be the uppercased lemma, not the surface form.""" + doc = _make_doc( + words=["Apple", "acquired", "IBM"], + lemmas=["Apple", "acquire", "IBM"], + dep_labels=["nsubj", "ROOT", "dobj"], + head_indices=[1, 1, 1], + ) + triples = extract_dependency_triples(doc) + relations = {rel for _, rel, _, _ in triples} + assert "ACQUIRE" in relations + assert "ACQUIRED" not in relations # surface form must NOT appear + + def test_no_object_no_svo_triple(self) -> None: + """A sentence with nsubj but no dobj should yield no SVO triple.""" + # "Apple runs" + doc = _make_doc( + words=["Apple", "runs"], + lemmas=["Apple", "run"], + dep_labels=["nsubj", "ROOT"], + head_indices=[1, 1], + ) + triples = extract_dependency_triples(doc) + # No SVO triple — only possibly a prep triple if there were one + svo_triples = [t for t in triples if t[1] not in ("TO", "IN", "ON", "AT")] + assert len(svo_triples) == 0 + + def test_no_subject_no_svo_triple(self) -> None: + """A sentence with dobj but no subject should yield no SVO triple.""" + # "(Ø) acquires IBM" — dobj without nsubj + doc = _make_doc( + words=["acquires", "IBM"], + lemmas=["acquire", "IBM"], + dep_labels=["ROOT", "dobj"], + head_indices=[0, 0], + ) + triples = extract_dependency_triples(doc) + assert len(triples) == 0 + + # ------------------------------------------------------------------ + # Passive voice + # ------------------------------------------------------------------ + + def test_passive_nsubjpass_agent_yields_triple(self) -> None: + """'IBM was acquired by Apple' → (IBM, ACQUIRE_BY, Apple, 0.9).""" + # IBM(nsubjpass) was(auxpass) acquired(ROOT) by(agent) Apple(pobj) + doc = _make_doc( + words=["IBM", "was", "acquired", "by", "Apple"], + lemmas=["IBM", "be", "acquire", "by", "Apple"], + dep_labels=["nsubjpass", "auxpass", "ROOT", "agent", "pobj"], + head_indices=[2, 2, 2, 2, 3], + ) + triples = extract_dependency_triples(doc) + assert ("IBM", "ACQUIRE_BY", "Apple", 0.9) in triples + + def test_passive_relation_has_by_suffix(self) -> None: + """Passive triples must carry the _BY suffix on the relation.""" + doc = _make_doc( + words=["IBM", "was", "acquired", "by", "Apple"], + lemmas=["IBM", "be", "acquire", "by", "Apple"], + dep_labels=["nsubjpass", "auxpass", "ROOT", "agent", "pobj"], + head_indices=[2, 2, 2, 2, 3], + ) + triples = extract_dependency_triples(doc) + relations = {rel for _, rel, _, _ in triples} + assert any(rel.endswith("_BY") for rel in relations) + + def test_passive_confidence_is_0_9(self) -> None: + """Passive triples must have confidence 0.9.""" + doc = _make_doc( + words=["IBM", "was", "acquired", "by", "Apple"], + lemmas=["IBM", "be", "acquire", "by", "Apple"], + dep_labels=["nsubjpass", "auxpass", "ROOT", "agent", "pobj"], + head_indices=[2, 2, 2, 2, 3], + ) + triples = extract_dependency_triples(doc) + passive = [t for t in triples if t[1].endswith("_BY")] + assert len(passive) > 0 + for _, _, _, conf in passive: + assert conf == 0.9 + + # ------------------------------------------------------------------ + # Prepositional triples + # ------------------------------------------------------------------ + + def test_prep_triple_without_dobj(self) -> None: + """'Apple moved to California' → (Apple, TO, California, 0.7).""" + # Apple(nsubj) moved(ROOT) to(prep) California(pobj) + doc = _make_doc( + words=["Apple", "moved", "to", "California"], + lemmas=["Apple", "move", "to", "California"], + dep_labels=["nsubj", "ROOT", "prep", "pobj"], + head_indices=[1, 1, 1, 2], + ) + triples = extract_dependency_triples(doc) + assert ("Apple", "TO", "California", 0.7) in triples + + def test_prep_triple_with_dobj(self) -> None: + """'Apple sold IBM to Google' → primary (Apple, SELL, IBM) + secondary (IBM, TO, Google).""" + # Apple(nsubj) sold(ROOT) IBM(dobj) to(prep) Google(pobj) + doc = _make_doc( + words=["Apple", "sold", "IBM", "to", "Google"], + lemmas=["Apple", "sell", "IBM", "to", "Google"], + dep_labels=["nsubj", "ROOT", "dobj", "prep", "pobj"], + head_indices=[1, 1, 1, 1, 3], + ) + triples = extract_dependency_triples(doc) + assert ("Apple", "SELL", "IBM", 1.0) in triples + assert ("IBM", "TO", "Google", 0.7) in triples + + def test_prep_triple_confidence_is_0_7(self) -> None: + """All prepositional triples must carry confidence 0.7.""" + doc = _make_doc( + words=["Apple", "moved", "to", "California"], + lemmas=["Apple", "move", "to", "California"], + dep_labels=["nsubj", "ROOT", "prep", "pobj"], + head_indices=[1, 1, 1, 2], + ) + triples = extract_dependency_triples(doc) + prep_triples = [t for t in triples if t[1] == "TO"] + assert len(prep_triples) > 0 + for _, _, _, conf in prep_triples: + assert conf == 0.7 + + # ------------------------------------------------------------------ + # Multi-token entity merging + # ------------------------------------------------------------------ + + def test_multiword_entity_collapsed_into_single_token(self) -> None: + """Named entities spanning multiple tokens are merged before extraction. + + A two-token entity set via ``doc.ents`` must appear as a single string + in the resulting triple head/tail. + """ + from spacy.tokens import Doc, Span + + nlp = spacy.blank("en") + # "Supply Chain acquired IBM" + words = ["Supply", "Chain", "acquired", "IBM"] + lemmas = ["Supply", "Chain", "acquire", "IBM"] + spaces = [True, True, True, False] + doc = Doc(nlp.vocab, words=words, spaces=spaces, lemmas=lemmas) + + # Set dependency arcs: "Supply Chain" fused as nsubj for "acquired" + # We treat token 0 ("Supply") as the head of the nsubj span. + doc[0].dep_ = "nsubj" + doc[0].head = doc[2] + doc[1].dep_ = "compound" + doc[1].head = doc[0] + doc[2].dep_ = "ROOT" + doc[2].head = doc[2] + doc[3].dep_ = "dobj" + doc[3].head = doc[2] + + # Mark "Supply Chain" (tokens 0-1) as a named entity so the retokenizer + # will merge it into a single token. + doc.ents = [Span(doc, 0, 2, label="ORG")] + + triples = extract_dependency_triples(doc) + + heads = {t[0] for t in triples} + # After merging, "Supply Chain" should be a single token string + assert any( + "Supply" in h for h in heads + ), f"expected merged entity in heads, got {heads}" + + # ------------------------------------------------------------------ + # Return type + # ------------------------------------------------------------------ + + def test_returns_set(self) -> None: + """Return value must be a set.""" + doc = _make_doc( + words=["Apple", "acquired", "IBM"], + lemmas=["Apple", "acquire", "IBM"], + dep_labels=["nsubj", "ROOT", "dobj"], + head_indices=[1, 1, 1], + ) + result = extract_dependency_triples(doc) + assert isinstance(result, set) + + def test_returns_4_tuples(self) -> None: + """Each element must be a 4-tuple (head, rel, tail, conf).""" + doc = _make_doc( + words=["Apple", "acquired", "IBM"], + lemmas=["Apple", "acquire", "IBM"], + dep_labels=["nsubj", "ROOT", "dobj"], + head_indices=[1, 1, 1], + ) + result = extract_dependency_triples(doc) + for item in result: + assert len(item) == 4 + head, rel, tail, conf = item + assert isinstance(head, str) + assert isinstance(rel, str) + assert isinstance(tail, str) + assert isinstance(conf, float) + + def test_empty_doc_returns_empty_set(self) -> None: + """An empty doc should yield an empty set.""" + from spacy.tokens import Doc + + nlp = spacy.blank("en") + doc = Doc(nlp.vocab, words=[], spaces=[]) + result = extract_dependency_triples(doc) + assert result == set() + + def test_no_root_token_returns_empty(self) -> None: + """If no token has dep ROOT, no triples are extracted.""" + doc = _make_doc( + words=["Apple", "IBM"], + lemmas=["Apple", "IBM"], + dep_labels=["nsubj", "dobj"], + head_indices=[1, 0], + ) + result = extract_dependency_triples(doc) + assert result == set() + + def test_deduplication_via_set(self) -> None: + """The same triple emitted for two different sentences is stored once.""" + # Two identical sentences in one doc — the set should deduplicate. + from spacy.tokens import Doc + + nlp = spacy.blank("en") + words = ["Apple", "acquired", "IBM", "Apple", "acquired", "IBM"] + lemmas = ["Apple", "acquire", "IBM", "Apple", "acquire", "IBM"] + spaces = [True, True, True, True, True, False] + doc = Doc(nlp.vocab, words=words, spaces=spaces, lemmas=lemmas) + # First sentence: tokens 0-2 + doc[0].dep_ = "nsubj" + doc[0].head = doc[1] + doc[1].dep_ = "ROOT" + doc[1].head = doc[1] + doc[2].dep_ = "dobj" + doc[2].head = doc[1] + # Second sentence: tokens 3-5 + doc[3].dep_ = "nsubj" + doc[3].head = doc[4] + doc[4].dep_ = "ROOT" + doc[4].head = doc[4] + doc[5].dep_ = "dobj" + doc[5].head = doc[4] + + result = extract_dependency_triples(doc) + # Should be exactly one triple (set deduplication) + assert result == {("Apple", "ACQUIRE", "IBM", 1.0)} + + def test_no_dep_annotations_falls_back_gracefully(self) -> None: + """When no DEP annotations are set, doc.sents raises ValueError. + + The function must fall back to treating the whole doc as one sentence + rather than raising. + """ + from spacy.tokens import Doc + + nlp = spacy.blank("en") + # Create doc with NO dependency annotations at all + doc = Doc(nlp.vocab, words=["Apple", "IBM"], spaces=[True, False]) + # No dep_ set — doc.sents will raise ValueError + result = extract_dependency_triples(doc) + # Should return empty set (no ROOT token found), not raise + assert isinstance(result, set) diff --git a/tests/unit/experimental/components/test_spacy_utils.py b/tests/unit/experimental/components/test_spacy_utils.py index e712fb423..4497012a9 100644 --- a/tests/unit/experimental/components/test_spacy_utils.py +++ b/tests/unit/experimental/components/test_spacy_utils.py @@ -12,7 +12,13 @@ # 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. -"""Unit tests for spacy_utils.normalize() and WORD_FILTER.""" +"""Unit tests for spacy_utils.normalize() and WORD_FILTER. + +These tests cover only the spaCy-free parts of the module and are safe to run +even when spaCy is not installed. Tests for ``extract_dependency_triples`` +live in ``test_spacy_dependency_triples.py`` and are gated on spaCy being +available. +""" import pytest @@ -52,8 +58,8 @@ def test_numbers_unchanged(self) -> None: def test_unicode_nfc_normalization(self) -> None: # "café" composed (U+00E9) vs decomposed (e + combining accent U+0301) - composed = "café" # NFC form - decomposed = "café" # NFD form — visually identical + composed = "café" # NFC form + decomposed = "café" # NFD form — visually identical assert normalize(composed) == normalize(decomposed) def test_unicode_lowercase_accent(self) -> None: diff --git a/uv.lock b/uv.lock index 6865845f0..e4e0a23f4 100644 --- a/uv.lock +++ b/uv.lock @@ -2974,7 +2974,7 @@ wheels = [ [[package]] name = "neo4j-graphrag" -version = "1.15.0" +version = "1.17.0" source = { editable = "." } dependencies = [ { name = "fsspec" }, From 5a4c2261b54ce64de09676b1cf27d614bacb90c5 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:18:55 +0200 Subject: [PATCH 04/12] spacy-entity-relation-extractor-task-007: unit tests for HierarchicalTextSplitter --- ...ogress-spacy-entity-relation-extractor.txt | 4 + .../tasks-spacy-entity-relation-extractor.yml | 6 +- .../test_hierarchical_splitter.py | 524 ++++++++++++++++++ 3 files changed, 531 insertions(+), 3 deletions(-) create mode 100644 tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index 4b2fd1f6f..8ce1f74ab 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -17,3 +17,7 @@ Decisions: Added `unicodedata.normalize("NFC", ...)` step per reviewer feedback Added `extract_dependency_triples(doc)` to `spacy_utils.py`: walks the dependency tree per sentence finding nsubj/nsubjpass → ROOT → dobj/pobj patterns, handles passive voice (_BY suffix, conf=0.9), prepositional triples (conf=0.7), and active SVO (conf=1.0). Uses spaCy 3.x retokenizer API (not deprecated `Span.merge()`) to collapse multi-word noun chunks and named entities. Falls back gracefully when sentence boundaries are unavailable. 18 unit tests in a new dedicated file; spaCy-free tests remain ungated. Files: ~src/neo4j_graphrag/experimental/components/spacy_utils.py, ~tests/unit/experimental/components/test_spacy_utils.py, +tests/unit/experimental/components/test_spacy_dependency_triples.py Decisions: Reviewer flagged `Span.merge()` removed in spaCy 3.x (silently swallowed); fixed to use `doc.retokenize()`. Moved spaCy-gated tests to separate file so normalize/WORD_FILTER tests stay runnable without spaCy installed. + +## TASK-007 — 2026-06-24T14:00:00Z +Added 20 unit tests for `HierarchicalTextSplitter` covering all 4 header strategies (markdown, capitalization, blank_line, spacy_verbless), recursive overlap splitting with exact last-N==first-N invariant verification, drop_verbless_sentences filter (SpaCy fully mocked via patch("spacy.load")), chunk index sequencing across sections, and edge cases (empty/whitespace input, invalid constructor args). No network calls or model downloads required. +Files: +tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index 655f996b9..a67618e6d 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -81,7 +81,7 @@ tasks: objects, and delegates lexical graph construction to LexicalGraphBuilder when `create_lexical_graph=True`. Confidence scores go in relationship properties. - status: pending + status: in_progress dependencies: - task-002 - task-003 @@ -156,7 +156,7 @@ tasks: Test section splitting, recursive character splitting with overlap, and index sequencing. For drop_verbless_sentences tests, mock the SpaCy pipeline so no model download is needed. - status: pending + status: completed dependencies: - task-006 acceptance_criteria: @@ -165,7 +165,7 @@ tasks: - "Test verifies section boundary respected: two Markdown-header sections produce separate chunks" - "Test verifies verbless sentence dropped when drop_verbless_sentences=True (spacy mocked)" - "Test verifies chunk indices are sequential starting from 0" - completed_at: null + completed_at: "2026-06-24T14:00:00Z" - id: task-008 title: "Integration tests for both components against real SpaCy model" diff --git a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py new file mode 100644 index 000000000..ba06e02d6 --- /dev/null +++ b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py @@ -0,0 +1,524 @@ +# 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. +"""Unit tests for HierarchicalTextSplitter. + +SpaCy is not downloaded in these tests. Where the splitter would normally +load a model (`drop_verbless_sentences=True` or `header_strategy="spacy_verbless"`), +the tests patch `spacy.load` with a lightweight fake nlp object so that no +network access or model installation is required. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +# Skip the entire module when spaCy is not installed at all. +spacy = pytest.importorskip("spacy") + +from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( # noqa: E402 + HierarchicalTextSplitter, +) + + +# --------------------------------------------------------------------------- +# Helpers for building fake SpaCy objects without a real model +# --------------------------------------------------------------------------- + + +def _make_fake_token(text: str, pos: str) -> MagicMock: + """Return a MagicMock that looks like a spaCy Token.""" + tok = MagicMock() + tok.text = text + tok.pos_ = pos + return tok + + +def _make_fake_sent(text: str, tokens: list[MagicMock]) -> MagicMock: + """Return a MagicMock that looks like a spaCy Span (sentence). + + Uses ``side_effect`` instead of ``return_value`` so that each call to + ``iter(sent)`` creates a *fresh* iterator — important if the sentence is + iterated more than once (e.g. across multiple ``run()`` calls or in future + multi-pass tests). + """ + sent = MagicMock() + sent.text = text + sent.__iter__ = MagicMock(side_effect=lambda: iter(tokens)) + return sent + + +def _make_fake_doc(sentences: list[MagicMock]) -> MagicMock: + """Return a MagicMock that looks like a spaCy Doc with .sents.""" + doc = MagicMock() + doc.sents = sentences + doc.__iter__ = MagicMock(return_value=iter([])) + return doc + + +def _make_nlp_returning_doc(sentences: list[MagicMock]) -> MagicMock: + """Return a callable MagicMock that acts as a spaCy nlp() pipeline.""" + nlp = MagicMock() + nlp.return_value = _make_fake_doc(sentences) + return nlp + + +# --------------------------------------------------------------------------- +# Tests: header_strategy="markdown" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_markdown_two_sections_produce_two_chunks() -> None: + """Two Markdown sections produce exactly two chunks.""" + text = "# Introduction\nThis is the intro section.\n# Conclusion\nThis is the conclusion." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 2 + assert "Introduction" in result.chunks[0].text + assert "Conclusion" in result.chunks[1].text + + +@pytest.mark.asyncio +async def test_markdown_single_section_produces_one_chunk() -> None: + """Text with no Markdown headers is treated as one section.""" + text = "No headers here. Just a single paragraph of text." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 1 + + +@pytest.mark.asyncio +async def test_markdown_three_sections_sequential_indices() -> None: + """Three Markdown sections produce chunks with sequential indices 0, 1, 2.""" + text = "# A\nSection A body.\n# B\nSection B body.\n# C\nSection C body." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 3 + for i, chunk in enumerate(result.chunks): + assert chunk.index == i + + +# --------------------------------------------------------------------------- +# Tests: header_strategy="capitalization" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_capitalization_allcaps_header_splits() -> None: + """ALL_CAPS lines without terminal punctuation act as section headers.""" + text = "INTRODUCTION\nThis paragraph describes the introduction.\nCONCLUSION\nThis paragraph wraps things up." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="capitalization", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 2 + assert "INTRODUCTION" in result.chunks[0].text + assert "CONCLUSION" in result.chunks[1].text + + +@pytest.mark.asyncio +async def test_capitalization_title_case_header_splits() -> None: + """Title Case lines without terminal punctuation act as section headers.""" + text = "First Section Title\nContent for the first section.\nSecond Section Title\nContent for the second section." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="capitalization", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 2 + + +@pytest.mark.asyncio +async def test_capitalization_indices_sequential() -> None: + """Chunk indices are sequential starting from 0.""" + text = "PART ONE\nBody of part one.\nPART TWO\nBody of part two.\nPART THREE\nBody of part three." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="capitalization", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + for i, chunk in enumerate(result.chunks): + assert chunk.index == i, f"chunk {i} has index {chunk.index}" + + +# --------------------------------------------------------------------------- +# Tests: header_strategy="blank_line" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blank_line_short_surrounded_header_splits() -> None: + """Short lines surrounded by blank lines act as section headers.""" + text = "\nOverview\n\nThis section covers the overview of the system.\n\nDetails\n\nThis section covers the fine-grained details." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="blank_line", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 2 + assert "Overview" in result.chunks[0].text + assert "Details" in result.chunks[1].text + + +@pytest.mark.asyncio +async def test_blank_line_indices_sequential() -> None: + """Chunk indices are sequential starting from 0 for blank_line strategy.""" + text = "\nPart A\n\nContent for part A goes here.\n\nPart B\n\nContent for part B goes here." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="blank_line", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + for i, chunk in enumerate(result.chunks): + assert chunk.index == i + + +# --------------------------------------------------------------------------- +# Tests: overlap splitting for large sections +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_large_section_split_with_overlap() -> None: + """A section larger than max_chunk_size is split into multiple chunks + with the last N characters of chunk K equal to the first N characters + of chunk K+1. + + Strategy: use no header strategy (treat whole text as one section) via a + single Markdown section; make the body very long and max_chunk_size small + relative to the body so that we get several chunks all in the body content, + and the overlap check applies to non-header chunks. + """ + # The section body is pure 'x' characters with no whitespace. + # With max_size=50, overlap=20, step=30: + # chunk 0: body[0:50] = 'x'*50 + # chunk 1: body[30:80] = 'x'*50, so overlap = body[30:50] = 'x'*20 + # We avoid the Markdown header prefix eating into the overlap window by + # using a plain-text body (no Markdown header at all, so the whole text + # is one section) and using drop_verbless_sentences=False. + body = "x" * 500 + overlap = 20 + max_size = 50 + splitter = HierarchicalTextSplitter( + max_chunk_size=max_size, + chunk_overlap=overlap, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(body) + chunks = result.chunks + assert len(chunks) > 1 + + # Verify overlap: last `overlap` chars of chunk[k] == first `overlap` chars of chunk[k+1]. + for k in range(len(chunks) - 1): + tail = chunks[k].text[-overlap:] + head = chunks[k + 1].text[:overlap] + assert tail == head, ( + f"Overlap mismatch between chunk {k} and {k + 1}: " + f"tail={tail!r}, head={head!r}" + ) + + +@pytest.mark.asyncio +async def test_small_section_emitted_as_single_chunk() -> None: + """A section shorter than max_chunk_size is emitted as a single chunk.""" + text = "# Tiny\nShort body." + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 1 + assert "Tiny" in result.chunks[0].text + assert "Short body" in result.chunks[0].text + + +@pytest.mark.asyncio +async def test_overlap_chunk_indices_sequential() -> None: + """Indices remain sequential when a section is split due to size.""" + body = "y" * 300 + text = f"# Big Section\n{body}" + splitter = HierarchicalTextSplitter( + max_chunk_size=60, + chunk_overlap=10, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) > 1 + for i, chunk in enumerate(result.chunks): + assert chunk.index == i + + +# --------------------------------------------------------------------------- +# Tests: drop_verbless_sentences=True (SpaCy mocked) +# --------------------------------------------------------------------------- + + +def _make_spacy_nlp_with_verbless() -> MagicMock: + """Return a fake nlp() that drops verbless sentences when called. + + Sentence 1 (verbless): "No verb here" — tokens have no VERB/AUX. + Sentence 2 (with verb): "The dog runs fast" — one VERB token. + """ + # Tokens for sentence 1 (verbless: no VERB/AUX tags) + sent1_tokens = [ + _make_fake_token("No", "DET"), + _make_fake_token("verb", "NOUN"), + _make_fake_token("here", "ADV"), + ] + sent1 = _make_fake_sent("No verb here", sent1_tokens) + + # Tokens for sentence 2 (contains a VERB) + sent2_tokens = [ + _make_fake_token("The", "DET"), + _make_fake_token("dog", "NOUN"), + _make_fake_token("runs", "VERB"), + _make_fake_token("fast", "ADV"), + ] + sent2 = _make_fake_sent("The dog runs fast", sent2_tokens) + + return _make_nlp_returning_doc([sent1, sent2]) + + +@pytest.mark.asyncio +async def test_drop_verbless_removes_verbless_sentence() -> None: + """When drop_verbless_sentences=True, sentences with no verb are removed.""" + fake_nlp = _make_spacy_nlp_with_verbless() + + with patch("spacy.load", return_value=fake_nlp): + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=True, + model="en_core_web_sm", + ) + + text = "# Section\nNo verb here. The dog runs fast." + result = await splitter.run(text) + + assert len(result.chunks) == 1 + chunk_text = result.chunks[0].text + # Verbless sentence should be gone; verbal sentence should remain. + assert "The dog runs fast" in chunk_text + assert "No verb here" not in chunk_text + + +@pytest.mark.asyncio +async def test_drop_verbless_keeps_verbal_sentences() -> None: + """When all sentences contain a verb, no text is removed.""" + sent_tokens = [ + _make_fake_token("She", "PRON"), + _make_fake_token("walks", "VERB"), + _make_fake_token("home", "NOUN"), + ] + sent = _make_fake_sent("She walks home", sent_tokens) + fake_nlp = _make_nlp_returning_doc([sent]) + + with patch("spacy.load", return_value=fake_nlp): + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=True, + model="en_core_web_sm", + ) + + text = "# Section\nShe walks home." + result = await splitter.run(text) + + assert len(result.chunks) == 1 + assert "She walks home" in result.chunks[0].text + + +# --------------------------------------------------------------------------- +# Tests: chunk index sequencing across multiple sections +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_indices_sequential_across_multiple_sections() -> None: + """Indices are global — they continue from the last chunk of the previous section.""" + text = ( + "# Alpha\nFirst section body.\n" + "# Beta\nSecond section body.\n" + "# Gamma\nThird section body." + ) + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + assert len(result.chunks) == 3 + for expected_index, chunk in enumerate(result.chunks): + assert chunk.index == expected_index + + +# --------------------------------------------------------------------------- +# Tests: header_strategy="spacy_verbless" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_spacy_verbless_strategy_splits_at_verbless_heading() -> None: + """header_strategy="spacy_verbless" uses SpaCy to detect verbless headings. + + A short sentence with no verb that precedes a longer sentence is treated as + a section heading and causes a split. This test mocks the nlp pipeline so + no model is downloaded. + """ + # Sentence 1 (verbless heading, short ≤ 80 chars): "Introduction" + sent1_tokens = [_make_fake_token("Introduction", "NOUN")] + sent1 = _make_fake_sent("Introduction", sent1_tokens) + + # Sentence 2 (long body sentence, > 80 chars, contains VERB): + long_body = "This section covers all the foundational concepts you need to understand before proceeding further." + sent2_tokens = [ + _make_fake_token("This", "DET"), + _make_fake_token("section", "NOUN"), + _make_fake_token("covers", "VERB"), + ] + sent2 = _make_fake_sent(long_body, sent2_tokens) + + # Sentence 3 (verbless heading): "Conclusion" + sent3_tokens = [_make_fake_token("Conclusion", "NOUN")] + sent3 = _make_fake_sent("Conclusion", sent3_tokens) + + # Sentence 4 (another long body): + long_body2 = "This final section wraps up all the topics discussed and provides closing remarks for the reader." + sent4_tokens = [ + _make_fake_token("This", "DET"), + _make_fake_token("section", "NOUN"), + _make_fake_token("wraps", "VERB"), + ] + sent4 = _make_fake_sent(long_body2, sent4_tokens) + + fake_nlp = _make_nlp_returning_doc([sent1, sent2, sent3, sent4]) + + with patch("spacy.load", return_value=fake_nlp): + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="spacy_verbless", + drop_verbless_sentences=False, + model="en_core_web_sm", + ) + + # Use arbitrary text — the nlp mock controls what sents are returned. + text = "Introduction. " + long_body + " Conclusion. " + long_body2 + result = await splitter.run(text) + + # Two verbless headings (sent1, sent3) each trigger a split → 2 sections. + assert len(result.chunks) == 2 + assert result.chunks[0].index == 0 + assert result.chunks[1].index == 1 + + +# --------------------------------------------------------------------------- +# Tests: edge cases +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_empty_text_returns_no_chunks() -> None: + """Empty input produces an empty chunk list.""" + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run("") + assert result.chunks == [] + + +@pytest.mark.asyncio +async def test_whitespace_only_returns_no_chunks() -> None: + """Whitespace-only input produces an empty chunk list.""" + splitter = HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + result = await splitter.run(" \n\t\n ") + assert result.chunks == [] + + +def test_invalid_header_strategy_raises() -> None: + """An unrecognised header_strategy raises ValueError at construction time.""" + with pytest.raises(ValueError, match="header_strategy must be one of"): + HierarchicalTextSplitter( + max_chunk_size=2048, + chunk_overlap=0, + header_strategy="unknown_strategy", + drop_verbless_sentences=False, + ) + + +def test_chunk_overlap_ge_max_chunk_size_raises() -> None: + """chunk_overlap >= max_chunk_size raises ValueError.""" + with pytest.raises(ValueError, match="chunk_overlap must be strictly less than max_chunk_size"): + HierarchicalTextSplitter( + max_chunk_size=100, + chunk_overlap=100, + header_strategy="markdown", + drop_verbless_sentences=False, + ) + + +def test_max_chunk_size_zero_raises() -> None: + """max_chunk_size=0 raises ValueError.""" + with pytest.raises(ValueError, match="max_chunk_size must be strictly greater than 0"): + HierarchicalTextSplitter( + max_chunk_size=0, + chunk_overlap=0, + header_strategy="markdown", + drop_verbless_sentences=False, + ) From 2f9ae33450d2905862d2db6ddb6f5c1a051e65fd Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:28:14 +0200 Subject: [PATCH 05/12] spacy-entity-relation-extractor-task-004: implement SpacyEntityRelationExtractor --- .../spacy_entity_relation_extractor.py | 406 ++++++++++++++++++ 1 file changed, 406 insertions(+) create mode 100644 src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py diff --git a/src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py new file mode 100644 index 000000000..df58d2743 --- /dev/null +++ b/src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py @@ -0,0 +1,406 @@ +# 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. +"""SpaCy-based entity and relation extractor. + +Uses SpaCy's NLP pipeline to extract entities and relationships from text +chunks without requiring an LLM. Extraction combines dependency-based +triple extraction and optional proximity-based linear triple extraction. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Optional + +from neo4j_graphrag.experimental.components.entity_relation_extractor import ( + EntityRelationExtractor, + OnError, +) +from neo4j_graphrag.experimental.components.lexical_graph import LexicalGraphBuilder +from neo4j_graphrag.experimental.components.schema import GraphSchema +from neo4j_graphrag.experimental.components.spacy_linear_extractor import ( + extract_linear_triples, +) +from neo4j_graphrag.experimental.components.spacy_utils import ( + WORD_FILTER, + extract_dependency_triples, + normalize, +) +from neo4j_graphrag.experimental.components.types import ( + DocumentInfo, + LexicalGraphConfig, + Neo4jGraph, + Neo4jNode, + Neo4jRelationship, + TextChunk, + TextChunks, +) + +if TYPE_CHECKING: + import spacy + import spacy.tokens + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# SpaCy NER label → schema node-type label mapping +# --------------------------------------------------------------------------- +# Maps SpaCy's built-in NER entity type labels to commonly used schema labels. +# Used when a GraphSchema is provided to filter out unrecognised entity types. +# The mapping is intentionally broad; a case-insensitive label-match on the +# schema side covers the rest (see _build_allowed_ner_labels). + +_SPACY_NER_TO_SCHEMA_LABEL: dict[str, str] = { + "ORG": "Organization", + "PERSON": "Person", + "GPE": "Location", + "LOC": "Location", + "FAC": "Location", + "NORP": "Group", + "PRODUCT": "Product", + "EVENT": "Event", + "WORK_OF_ART": "WorkOfArt", + "LAW": "Law", + "LANGUAGE": "Language", + "DATE": "Date", + "TIME": "Time", + "MONEY": "Money", + "QUANTITY": "Quantity", + "CARDINAL": "Cardinal", + "ORDINAL": "Ordinal", + "PERCENT": "Percent", +} + + +def _build_allowed_ner_labels(schema: GraphSchema) -> Optional[set[str]]: + """Return the set of SpaCy NER labels that match schema node_type labels. + + Returns ``None`` when the schema has no node_types (no filtering). Also + returns ``None`` — with a warning — when node_types are defined but none + can be mapped to a SpaCy NER label (e.g. purely custom types like + "Contract"). This "fail-open" behaviour avoids silently filtering out all + entities when the schema uses domain-specific labels. + + The matching strategy (in priority order): + + 1. Direct case-insensitive match of the schema label against the values in + :data:`_SPACY_NER_TO_SCHEMA_LABEL` (e.g. "Organization" → ``ORG``). + 2. Try the schema label uppercased directly as a SpaCy NER label (covers + schemas that already use ``ORG``, ``PERSON``, ``GPE``, etc.). + """ + if not schema.node_types: + return None + + allowed: set[str] = set() + schema_labels_lower = {nt.label.lower() for nt in schema.node_types} + + # Strategy 1: forward lookup in the canonical NER→schema-label map. + for spacy_ner, schema_default in _SPACY_NER_TO_SCHEMA_LABEL.items(): + if schema_default.lower() in schema_labels_lower: + allowed.add(spacy_ner) + + # Strategy 2: treat the schema label as a raw SpaCy NER label (e.g. "ORG"). + for label_lower in schema_labels_lower: + upper = label_lower.upper() + if upper in _SPACY_NER_TO_SCHEMA_LABEL: + allowed.add(upper) + + if not allowed: + logger.warning( + "GraphSchema has node_types but none could be mapped to a SpaCy NER label. " + "NER filtering will be skipped (all entity types are kept). " + "Consider using standard schema labels such as 'Organization', 'Person', or 'Location'." + ) + return None + + return allowed + + +def _build_allowed_relation_types(schema: GraphSchema) -> Optional[set[str]]: + """Return the set of normalised relation type strings allowed by the schema. + + Each relationship_type label is uppercased and has spaces replaced by + underscores to match the format produced by the triple extractors. + Returns ``None`` when the schema defines no relationship_types (no filtering). + """ + if not schema.relationship_types: + return None + return {rt.label.upper().replace(" ", "_") for rt in schema.relationship_types} + + +class SpacyEntityRelationExtractor(EntityRelationExtractor): + """Extract entities and relations from text using SpaCy NLP pipelines. + + This extractor does not require a large language model. It uses SpaCy's + dependency parser and/or named-entity recogniser to produce (head, relation, + tail) triples which are then converted to :class:`Neo4jNode` and + :class:`Neo4jRelationship` objects. + + Args: + model (str): SpaCy model name to load (e.g. ``"en_core_web_sm"``). + The model is loaded once at construction time. A + :class:`ValueError` is raised with install instructions if the + model is not available. + word_filter (Optional[list[str]]): Custom list of stopwords to exclude + when checking whether a normalised entity name is meaningful. When + ``None`` the built-in :data:`WORD_FILTER` constant is used. + use_linear_extractor (bool): When ``True`` (default), proximity-based + triples from :func:`extract_linear_triples` are unioned with the + dependency-based triples from :func:`extract_dependency_triples`. + create_lexical_graph (bool): Inherited from + :class:`EntityRelationExtractor`. When ``True`` (default), chunk + and document nodes are added to the returned graph. + on_error (OnError): Inherited from :class:`EntityRelationExtractor`. + Controls behaviour when a per-chunk exception occurs. + + Example:: + + from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, + ) + from neo4j_graphrag.experimental.components.types import TextChunks, TextChunk + + extractor = SpacyEntityRelationExtractor(model="en_core_web_sm") + chunks = TextChunks(chunks=[TextChunk(text="Apple bought Google.", index=0)]) + import asyncio + graph = asyncio.run(extractor.run(chunks=chunks)) + """ + + def __init__( + self, + model: str = "en_core_web_sm", + word_filter: Optional[list[str]] = None, + use_linear_extractor: bool = True, + create_lexical_graph: bool = True, + on_error: OnError = OnError.IGNORE, + **kwargs: Any, + ) -> None: + super().__init__( + on_error=on_error, + create_lexical_graph=create_lexical_graph, + **kwargs, + ) + self.use_linear_extractor = use_linear_extractor + self._word_filter: frozenset[str] = ( + frozenset(word_filter) if word_filter is not None else WORD_FILTER + ) + + try: + import spacy as _spacy + + self.nlp: spacy.language.Language = _spacy.load(model) + except (OSError, ImportError): + raise ValueError( + f"SpaCy model '{model}' is not installed or spacy is not available. " + f"Install with:\n" + f" pip install 'neo4j-graphrag[nlp]'\n" + f" python -m spacy download {model}" + ) from None + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _is_valid_entity(self, name: str) -> bool: + """Return True when *name* (after normalisation) passes the word filter.""" + norm = normalize(name) + return bool(norm) and norm not in self._word_filter + + def _triples_to_graph( + self, + triples: set[tuple[str, str, str, float]], + allowed_ner_labels: Optional[set[str]], + allowed_rel_types: Optional[set[str]], + ner_labels_by_text: dict[str, str], + ) -> Neo4jGraph: + """Convert a set of (head, relation, tail, conf) triples to a Neo4jGraph. + + Applies schema filters and word-filter cleaning. + + Args: + triples: Raw triples from the extractors. + allowed_ner_labels: SpaCy NER label whitelist (``None`` = no filter). + allowed_rel_types: Relation type whitelist (``None`` = no filter). + ner_labels_by_text: Mapping of entity text → SpaCy NER label for + entities found in the current document. Used for NER filtering. + + Returns: + A :class:`Neo4jGraph` with deduplicated nodes and relationships. + """ + nodes: dict[str, Neo4jNode] = {} + relationships: list[Neo4jRelationship] = [] + + for head_text, relation, tail_text, conf in triples: + # --- Word filter ------------------------------------------- + if not self._is_valid_entity(head_text): + continue + if not self._is_valid_entity(tail_text): + continue + + # --- Relation type filter ---------------------------------- + rel_type = relation.upper().replace(" ", "_") + if allowed_rel_types is not None and rel_type not in allowed_rel_types: + continue + + # --- NER type filter (only when schema is provided) -------- + if allowed_ner_labels is not None: + head_ner = ner_labels_by_text.get(head_text) + tail_ner = ner_labels_by_text.get(tail_text) + if head_ner is not None and head_ner not in allowed_ner_labels: + continue + if tail_ner is not None and tail_ner not in allowed_ner_labels: + continue + + # --- Build nodes ------------------------------------------- + head_id = normalize(head_text) + tail_id = normalize(tail_text) + + if head_id not in nodes: + nodes[head_id] = Neo4jNode( + id=head_id, + label="Entity", + properties={"name": head_text}, + ) + if tail_id not in nodes: + nodes[tail_id] = Neo4jNode( + id=tail_id, + label="Entity", + properties={"name": tail_text}, + ) + + # --- Build relationship ------------------------------------ + relationships.append( + Neo4jRelationship( + start_node_id=head_id, + end_node_id=tail_id, + type=rel_type, + properties={"confidence": conf}, + ) + ) + + return Neo4jGraph( + nodes=list(nodes.values()), + relationships=relationships, + ) + + def _extract_chunk( + self, + chunk: TextChunk, + allowed_ner_labels: Optional[set[str]], + allowed_rel_types: Optional[set[str]], + ) -> Neo4jGraph: + """Run NLP + extraction for a single chunk and return a raw graph.""" + doc: spacy.tokens.Doc = self.nlp(chunk.text) + + # Build NER label index for schema filtering. + # ent.text matches the original span text before any retokenization + # performed inside extract_dependency_triples (which works on a copy). + ner_labels_by_text: dict[str, str] = {ent.text: ent.label_ for ent in doc.ents} + + # Dependency-based triples. + triples: set[tuple[str, str, str, float]] = extract_dependency_triples(doc) + + # Union with linear (proximity) triples when requested. + if self.use_linear_extractor: + triples |= extract_linear_triples(doc) + + return self._triples_to_graph( + triples, + allowed_ner_labels, + allowed_rel_types, + ner_labels_by_text, + ) + + # ------------------------------------------------------------------ + # Component entry point + # ------------------------------------------------------------------ + + async def run( + self, + chunks: TextChunks, + document_info: Optional[DocumentInfo] = None, + lexical_graph_config: Optional[LexicalGraphConfig] = None, + schema: Optional[GraphSchema] = None, + **kwargs: Any, + ) -> Neo4jGraph: + """Extract entities and relations from all chunks. + + Args: + chunks: Text chunks to process. + document_info: Optional document metadata (used for lexical graph). + lexical_graph_config: Optional configuration for the lexical graph. + schema: Optional :class:`GraphSchema` to constrain extraction. + When provided, only entities matching schema node_type labels + (via SpaCy NER type mapping) and relations matching schema + relationship_type labels are kept. + + Returns: + A :class:`Neo4jGraph` containing the extracted entities and + relations, plus lexical-graph nodes/relationships when + ``create_lexical_graph=True``. + """ + # --- Schema-derived filters ------------------------------------ + allowed_ner_labels: Optional[set[str]] = None + allowed_rel_types: Optional[set[str]] = None + if schema is not None: + allowed_ner_labels = _build_allowed_ner_labels(schema) + allowed_rel_types = _build_allowed_relation_types(schema) + + # --- Lexical graph setup -------------------------------------- + lexical_graph_builder: Optional[LexicalGraphBuilder] = None + lexical_graph: Optional[Neo4jGraph] = None + + if self.create_lexical_graph: + config = lexical_graph_config or LexicalGraphConfig() + lexical_graph_builder = LexicalGraphBuilder(config=config) + lexical_graph_result = await lexical_graph_builder.run( + text_chunks=chunks, document_info=document_info + ) + lexical_graph = lexical_graph_result.graph + elif lexical_graph_config: + lexical_graph_builder = LexicalGraphBuilder(config=lexical_graph_config) + + # --- Per-chunk extraction ------------------------------------- + chunk_graphs: list[Neo4jGraph] = [] + for chunk in chunks.chunks: + try: + chunk_graph = self._extract_chunk( + chunk, allowed_ner_labels, allowed_rel_types + ) + self.update_ids(chunk_graph, chunk) + if lexical_graph_builder: + await lexical_graph_builder.process_chunk_extracted_entities( + chunk_graph, chunk + ) + chunk_graphs.append(chunk_graph) + except Exception as exc: + if self.on_error == OnError.RAISE: + raise + logger.warning( + f"Skipping chunk index={chunk.index} due to error: {exc}", + exc_info=True, + ) + + # --- Combine -------------------------------------------------- + graph = lexical_graph.model_copy(deep=True) if lexical_graph else Neo4jGraph() + for cg in chunk_graphs: + graph.nodes.extend(cg.nodes) + graph.relationships.extend(cg.relationships) + + logger.debug( + f"SpacyEntityRelationExtractor: extracted {len(graph.nodes)} nodes " + f"and {len(graph.relationships)} relationships from {len(chunks.chunks)} chunks." + ) + return graph From d876b447294dbf5e759e1a697dd28ecf704cfaa3 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:31:47 +0200 Subject: [PATCH 06/12] spacy-entity-relation-extractor-task-009: export new components from __init__ --- ...ogress-spacy-entity-relation-extractor.txt | 10 ++++++++ .../tasks-spacy-entity-relation-extractor.yml | 12 +++++----- .../experimental/components/__init__.py | 9 ++++++++ .../components/text_splitters/__init__.py | 23 +++++++++++++++++++ 4 files changed, 48 insertions(+), 6 deletions(-) diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index 8ce1f74ab..232b52f10 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -21,3 +21,13 @@ Decisions: Reviewer flagged `Span.merge()` removed in spaCy 3.x (silently swallo ## TASK-007 — 2026-06-24T14:00:00Z Added 20 unit tests for `HierarchicalTextSplitter` covering all 4 header strategies (markdown, capitalization, blank_line, spacy_verbless), recursive overlap splitting with exact last-N==first-N invariant verification, drop_verbless_sentences filter (SpaCy fully mocked via patch("spacy.load")), chunk index sequencing across sections, and edge cases (empty/whitespace input, invalid constructor args). No network calls or model downloads required. Files: +tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py + +## TASK-004 — 2026-06-24T15:00:00Z +Implemented `SpacyEntityRelationExtractor(EntityRelationExtractor)` in `src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py`. +Key design decisions: SpaCy model loaded once at construction time with user-friendly ValueError (covering both missing model and missing package); NER→schema label mapping via `_SPACY_NER_TO_SCHEMA_LABEL` forward-lookup + direct-uppercase Strategy (removed dead reverse map); fail-open NER filtering with `logger.warning` when schema types don't map to SpaCy NER labels; dependency + linear triples unioned and filtered by word filter, relation type, and NER type; lexical graph built via `LexicalGraphBuilder` following the same pattern as `LLMEntityRelationExtractor`; `on_error=IGNORE` logs and skips per-chunk exceptions. All types use lowercase builtins (`set[...]`, `tuple[...]`). ruff, mypy, pytest all clean. +Files: +src/neo4j_graphrag/experimental/components/spacy_entity_relation_extractor.py +Review feedback applied: removed `_SCHEMA_LABEL_TO_SPACY_NER` dead code, added warning log for unmapped schema types, broadened `except (OSError, ImportError)` in constructor, switched to lowercase type builtins. + +## TASK-009 — 2026-06-24T16:00:00Z +Exported SpacyEntityRelationExtractor and HierarchicalTextSplitter from __init__ files with ImportError guards. +Files: ~src/neo4j_graphrag/experimental/components/__init__.py, ~src/neo4j_graphrag/experimental/components/text_splitters/__init__.py diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index a67618e6d..f758b9a86 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -81,7 +81,8 @@ tasks: objects, and delegates lexical graph construction to LexicalGraphBuilder when `create_lexical_graph=True`. Confidence scores go in relationship properties. - status: in_progress + status: completed + completed_at: "2026-06-24T15:00:00Z" dependencies: - task-002 - task-003 @@ -95,7 +96,6 @@ tasks: - "use_linear_extractor=False produces no RELATED_TO relationships" - "create_lexical_graph=True adds Chunk and Document nodes to the output graph" - "on_error=OnError.IGNORE suppresses per-chunk exceptions; on_error=OnError.RAISE propagates them" - completed_at: null - id: task-005 title: "Unit tests for SpacyEntityRelationExtractor" @@ -104,7 +104,7 @@ tasks: Mock `spacy.load` and the nlp pipeline so no model download is required. Cover: successful extraction from two-chunk input, linear extractor toggle, lexical graph nodes, invalid model error, and on_error behaviour. - status: pending + status: completed dependencies: - task-004 acceptance_criteria: @@ -114,7 +114,7 @@ tasks: - "Test covers create_lexical_graph=True producing Chunk/Document nodes" - "Test covers ValueError raised for missing model" - "Test covers on_error=OnError.IGNORE swallowing a per-chunk exception" - completed_at: null + completed_at: "2026-06-24T16:00:00Z" - id: task-006 title: "Implement HierarchicalTextSplitter component" @@ -195,7 +195,8 @@ tasks: and the text_splitters sub-package `__init__.py`. Guard imports with a try/except so the rest of the package still imports cleanly when SpaCy is not installed. - status: pending + status: completed + completed_at: "2026-06-24T16:00:00Z" dependencies: - task-004 - task-006 @@ -203,7 +204,6 @@ tasks: - "from neo4j_graphrag.experimental.components import SpacyEntityRelationExtractor works when spacy is installed" - "from neo4j_graphrag.experimental.components.text_splitters import HierarchicalTextSplitter works when spacy is installed" - "Importing neo4j_graphrag.experimental.components without spacy installed does not raise ImportError" - completed_at: null - id: task-010 title: "Add example script for end-to-end spacy pipeline" diff --git a/src/neo4j_graphrag/experimental/components/__init__.py b/src/neo4j_graphrag/experimental/components/__init__.py index c0199c144..e712af224 100644 --- a/src/neo4j_graphrag/experimental/components/__init__.py +++ b/src/neo4j_graphrag/experimental/components/__init__.py @@ -12,3 +12,12 @@ # 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. + +try: + from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, + ) + + __all__ = ["SpacyEntityRelationExtractor"] +except ImportError: + pass diff --git a/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py b/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py index e69de29bb..d128c3410 100644 --- a/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py +++ b/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py @@ -0,0 +1,23 @@ +# 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. + +try: + from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( + HierarchicalTextSplitter, + ) + + __all__ = ["HierarchicalTextSplitter"] +except ImportError: + pass From d468c042ad8fc7056725b4b51061a1183ef314fb Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:31:57 +0200 Subject: [PATCH 07/12] spacy-entity-relation-extractor-task-005: unit tests for SpacyEntityRelationExtractor --- ...ogress-spacy-entity-relation-extractor.txt | 5 + .../test_spacy_entity_relation_extractor.py | 498 ++++++++++++++++++ 2 files changed, 503 insertions(+) create mode 100644 tests/unit/experimental/components/test_spacy_entity_relation_extractor.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index 232b52f10..08ac00263 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -1,3 +1,8 @@ +## TASK-005 — 2026-06-24T16:00:00Z +Unit tests for SpacyEntityRelationExtractor, spacy.load mocked. +13 tests covering: two-chunk extraction with Entity nodes + confidence property, use_linear_extractor=True/False toggle, create_lexical_graph=True adding Chunk/Document nodes, ValueError for missing/non-installed model, on_error=OnError.IGNORE suppression, custom word_filter exclusion, schema fail-open behaviour (unmapped NER types pass through), known NER type filtering. All 13 tests pass with no model download. +Files: +tests/unit/experimental/components/test_spacy_entity_relation_extractor.py + ## TASK-006 — 2026-06-24T12:00:00Z Created `HierarchicalTextSplitter(TextSplitter)` with four header strategies (markdown, capitalization, blank_line, spacy_verbless), recursive overlap splitting for large sections, and optional SpaCy-based verbless-sentence filtering. Ruff, mypy, and existing unit tests all pass. Files: +src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py diff --git a/tests/unit/experimental/components/test_spacy_entity_relation_extractor.py b/tests/unit/experimental/components/test_spacy_entity_relation_extractor.py new file mode 100644 index 000000000..61554ed26 --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_entity_relation_extractor.py @@ -0,0 +1,498 @@ +# 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. +"""Unit tests for SpacyEntityRelationExtractor. + +spaCy is an optional dependency; these tests are skipped when it is not +installed. ``spacy.load`` is mocked so no trained model download is required. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +pytest.importorskip("spacy", reason="spacy is not installed") + +from neo4j_graphrag.experimental.components.entity_relation_extractor import OnError # noqa: E402 +from neo4j_graphrag.experimental.components.schema import ( # noqa: E402 + GraphSchema, + NodeType, + RelationshipType, +) +from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( # noqa: E402 + SpacyEntityRelationExtractor, +) +from neo4j_graphrag.experimental.components.types import ( # noqa: E402 + DocumentInfo, + Neo4jGraph, + TextChunk, + TextChunks, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_mock_nlp( + dep_triples: set[tuple[str, str, str, float]] | None = None, + linear_triples: set[tuple[str, str, str, float]] | None = None, + ents: list[tuple[str, str]] | None = None, +) -> MagicMock: + """Return a mock ``spacy.language.Language`` callable. + + The mock's ``__call__`` method returns a Doc-like MagicMock whose + ``.ents`` attribute contains simple objects with ``.text`` and ``.label_`` + attributes. + + *dep_triples* and *linear_triples* are injected via patches on + ``extract_dependency_triples`` / ``extract_linear_triples`` rather than + via this mock, but the Doc object must still exist. + + Args: + dep_triples: Triples returned by extract_dependency_triples. + linear_triples: Triples returned by extract_linear_triples. + ents: List of (text, label_) for doc.ents. + """ + mock_nlp = MagicMock() + + # Build a simple entity list. + mock_ents = [] + for text, label in (ents or []): + ent = MagicMock() + ent.text = text + ent.label_ = label + mock_ents.append(ent) + + mock_doc = MagicMock() + mock_doc.ents = mock_ents + mock_nlp.return_value = mock_doc + + return mock_nlp + + +# --------------------------------------------------------------------------- +# 1. Successful extraction — two chunks, Entity nodes + relationships +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_successful_extraction_two_chunks() -> None: + """Two-chunk input produces Entity nodes and RELATED_TO relationships.""" + dep_triples: set[tuple[str, str, str, float]] = set() + linear_triples: set[tuple[str, str, str, float]] = { + ("Apple", "RELATED_TO", "Google", 0.3), + ("Microsoft", "RELATED_TO", "Azure", 0.3), + } + + mock_nlp = _make_mock_nlp( + ents=[("Apple", "ORG"), ("Google", "ORG")], + ) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + side_effect=[linear_triples, {("Microsoft", "RELATED_TO", "Azure", 0.3)}], + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + create_lexical_graph=False, + ) + chunks = TextChunks( + chunks=[ + TextChunk(text="Apple acquired Google.", index=0), + TextChunk(text="Microsoft launched Azure.", index=1), + ] + ) + result = await extractor.run(chunks=chunks) + + assert isinstance(result, Neo4jGraph) + # We expect at least 4 Entity nodes across both chunks (apple, google, microsoft, azure) + node_ids = {n.id for n in result.nodes} + assert len(node_ids) >= 2 # at minimum the chunk-0 entities + # All nodes should be labelled "Entity" + for node in result.nodes: + assert node.label == "Entity" + # All relationships should have a "confidence" property + for rel in result.relationships: + assert "confidence" in rel.properties + assert isinstance(rel.properties["confidence"], float) + + +# --------------------------------------------------------------------------- +# 2. use_linear_extractor=False — no RELATED_TO relationships +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_use_linear_extractor_false_no_related_to() -> None: + """With use_linear_extractor=False, linear triples are never produced.""" + dep_triples: set[tuple[str, str, str, float]] = set() + + mock_nlp = _make_mock_nlp(ents=[("Apple", "ORG"), ("Google", "ORG")]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ) as mock_dep, patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + ) as mock_linear: + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=False, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Apple bought Google.", index=0)]) + result = await extractor.run(chunks=chunks) + + # The linear extractor must NOT have been called. + mock_linear.assert_not_called() + # No nodes or relationships (dep_triples is empty). + assert result.nodes == [] + assert result.relationships == [] + + +# --------------------------------------------------------------------------- +# 3. use_linear_extractor=True — RELATED_TO relationships may be produced +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_use_linear_extractor_true_produces_related_to() -> None: + """With use_linear_extractor=True and NER entities, RELATED_TO rels appear.""" + dep_triples: set[tuple[str, str, str, float]] = set() + linear_triples: set[tuple[str, str, str, float]] = { + ("Apple", "RELATED_TO", "Google", 0.3), + } + + mock_nlp = _make_mock_nlp(ents=[("Apple", "ORG"), ("Google", "ORG")]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=linear_triples, + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Apple partnered with Google.", index=0)]) + result = await extractor.run(chunks=chunks) + + rel_types = {r.type for r in result.relationships} + assert "RELATED_TO" in rel_types + # Confidence should be 0.3 for linear triples. + for rel in result.relationships: + if rel.type == "RELATED_TO": + assert rel.properties["confidence"] == pytest.approx(0.3) + + +# --------------------------------------------------------------------------- +# 4. create_lexical_graph=True — Chunk/Document nodes present +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_lexical_graph_adds_chunk_nodes() -> None: + """create_lexical_graph=True adds Chunk (and Document) nodes to the graph.""" + dep_triples: set[tuple[str, str, str, float]] = set() + linear_triples: set[tuple[str, str, str, float]] = set() + + mock_nlp = _make_mock_nlp(ents=[]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=linear_triples, + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + create_lexical_graph=True, + ) + document_info = DocumentInfo(path="test_doc.txt") + chunks = TextChunks(chunks=[TextChunk(text="Some text.", index=0)]) + result = await extractor.run(chunks=chunks, document_info=document_info) + + labels = {n.label for n in result.nodes} + assert "Chunk" in labels + assert "Document" in labels + + +@pytest.mark.asyncio +async def test_create_lexical_graph_false_no_chunk_nodes() -> None: + """create_lexical_graph=False produces no Chunk or Document nodes.""" + dep_triples: set[tuple[str, str, str, float]] = set() + linear_triples: set[tuple[str, str, str, float]] = set() + + mock_nlp = _make_mock_nlp(ents=[]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=linear_triples, + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Some text.", index=0)]) + result = await extractor.run(chunks=chunks) + + labels = {n.label for n in result.nodes} + assert "Chunk" not in labels + assert "Document" not in labels + + +# --------------------------------------------------------------------------- +# 5. Constructor raises ValueError for non-installed model +# --------------------------------------------------------------------------- + + +def test_constructor_missing_model_raises_value_error() -> None: + """A missing SpaCy model raises ValueError with install instructions.""" + with patch("spacy.load", side_effect=OSError("Model not found")): + with pytest.raises(ValueError) as exc_info: + SpacyEntityRelationExtractor(model="en_core_web_sm") + + assert "python -m spacy download" in str(exc_info.value) + assert "en_core_web_sm" in str(exc_info.value) + + +def test_constructor_missing_spacy_package_raises_value_error() -> None: + """An ImportError (spacy not installed) is also wrapped in ValueError.""" + with patch("spacy.load", side_effect=ImportError("No module named 'spacy'")): + with pytest.raises(ValueError) as exc_info: + SpacyEntityRelationExtractor(model="en_core_web_sm") + + assert "python -m spacy download" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# 6. on_error=OnError.IGNORE suppresses per-chunk exceptions +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_on_error_ignore_suppresses_exception() -> None: + """on_error=OnError.IGNORE skips chunks that raise exceptions.""" + + def _raise(*args: object, **kwargs: object) -> None: + raise RuntimeError("NLP pipeline exploded") + + mock_nlp = MagicMock(side_effect=_raise) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=set(), + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + on_error=OnError.IGNORE, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="boom", index=0)]) + # Should not raise; the chunk is silently skipped. + result = await extractor.run(chunks=chunks) + + assert result.nodes == [] + assert result.relationships == [] + + +@pytest.mark.asyncio +async def test_on_error_raise_propagates_exception() -> None: + """on_error=OnError.RAISE re-raises per-chunk exceptions.""" + + def _raise(*args: object, **kwargs: object) -> None: + raise RuntimeError("NLP pipeline exploded") + + mock_nlp = MagicMock(side_effect=_raise) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=set(), + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + on_error=OnError.RAISE, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="boom", index=0)]) + with pytest.raises(RuntimeError, match="NLP pipeline exploded"): + await extractor.run(chunks=chunks) + + +# --------------------------------------------------------------------------- +# 7. Custom word_filter filters out entities in the list +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_custom_word_filter_excludes_entities() -> None: + """Entities whose normalised text appears in word_filter are excluded.""" + dep_triples: set[tuple[str, str, str, float]] = { + ("apple", "RELATED_TO", "stopword", 0.3), + ("apple", "RELATED_TO", "google", 0.3), + } + + mock_nlp = _make_mock_nlp(ents=[("apple", "ORG"), ("stopword", "ORG"), ("google", "ORG")]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=set(), + ): + # "stopword" is in our custom filter. + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + word_filter=["stopword"], + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="apple stopword google", index=0)]) + result = await extractor.run(chunks=chunks) + + # "stopword" should not appear as a node. + # Node IDs are prefixed with the chunk UUID: ":apple" — check the suffix. + node_id_suffixes = {n.id.split(":")[-1] for n in result.nodes} + assert "stopword" not in node_id_suffixes + # The (apple, google) triple should still be present. + assert any( + r.start_node_id.endswith("apple") and r.end_node_id.endswith("google") + for r in result.relationships + ) + + +# --------------------------------------------------------------------------- +# 8. Schema filtering: entities with unmapped NER types pass through (fail-open) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_schema_filtering_unmapped_ner_types_pass_through() -> None: + """Schema with only custom types that don't map to SpaCy NER → fail-open.""" + dep_triples: set[tuple[str, str, str, float]] = { + ("ContractA", "RELATED_TO", "ContractB", 0.3), + } + + # "Contract" is not in _SPACY_NER_TO_SCHEMA_LABEL, so fail-open applies. + mock_nlp = _make_mock_nlp(ents=[("ContractA", "MISC"), ("ContractB", "MISC")]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=set(), + ): + schema = GraphSchema( + node_types=(NodeType(label="Contract"),), + relationship_types=(RelationshipType(label="RELATED_TO"),), + ) + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="ContractA links ContractB.", index=0)]) + result = await extractor.run(chunks=chunks, schema=schema) + + # Fail-open: triples pass through because "Contract" has no SpaCy NER mapping. + # Node IDs are prefixed with the chunk UUID: ":contracta". + node_id_suffixes = {n.id.split(":")[-1] for n in result.nodes} + assert "contracta" in node_id_suffixes + assert "contractb" in node_id_suffixes + + +@pytest.mark.asyncio +async def test_schema_filtering_known_ner_type_filters_unmatched() -> None: + """Schema with Organization node_type: only ORG entities survive NER filter.""" + dep_triples: set[tuple[str, str, str, float]] = { + ("Apple", "RELATED_TO", "NewYork", 0.3), # NewYork is GPE, not ORG + } + + mock_nlp = _make_mock_nlp(ents=[("Apple", "ORG"), ("NewYork", "GPE")]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=set(), + ): + # Schema only allows Organization (→ ORG). GPE/Location excluded. + schema = GraphSchema( + node_types=(NodeType(label="Organization"),), + relationship_types=(RelationshipType(label="RELATED_TO"),), + ) + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Apple is in NewYork.", index=0)]) + result = await extractor.run(chunks=chunks, schema=schema) + + # The (Apple, NewYork) triple should be filtered out because NewYork is GPE. + node_ids = {n.id for n in result.nodes} + # "newyork" is a GPE; schema only allows ORG → triple is filtered. + assert "newyork" not in node_ids + + +# --------------------------------------------------------------------------- +# 9. Confidence property on relationships +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_relationships_have_confidence_property() -> None: + """All extracted relationships carry a numeric 'confidence' property.""" + dep_triples: set[tuple[str, str, str, float]] = { + ("Alpha", "WORKS_FOR", "Beta", 1.0), + } + + mock_nlp = _make_mock_nlp(ents=[]) + + with patch("spacy.load", return_value=mock_nlp), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_dependency_triples", + return_value=dep_triples, + ), patch( + "neo4j_graphrag.experimental.components.spacy_entity_relation_extractor.extract_linear_triples", + return_value=set(), + ): + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Alpha works for Beta.", index=0)]) + result = await extractor.run(chunks=chunks) + + assert len(result.relationships) == 1 + rel = result.relationships[0] + assert "confidence" in rel.properties + assert rel.properties["confidence"] == pytest.approx(1.0) From 3e567e1d344ec9006b262549c6bb17852ddcba73 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:34:01 +0200 Subject: [PATCH 08/12] spacy-entity-relation-extractor-task-010: add example spacy KG pipeline script --- ...ogress-spacy-entity-relation-extractor.txt | 4 + .../tasks-spacy-entity-relation-extractor.yml | 6 +- .../components/spacy_kg_pipeline.py | 111 ++++++++++++++++++ 3 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 examples/customize/build_graph/components/spacy_kg_pipeline.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index 08ac00263..d49893a1e 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -36,3 +36,7 @@ Review feedback applied: removed `_SCHEMA_LABEL_TO_SPACY_NER` dead code, added w ## TASK-009 — 2026-06-24T16:00:00Z Exported SpacyEntityRelationExtractor and HierarchicalTextSplitter from __init__ files with ImportError guards. Files: ~src/neo4j_graphrag/experimental/components/__init__.py, ~src/neo4j_graphrag/experimental/components/text_splitters/__init__.py + +## TASK-010 — 2026-06-24T17:00:00Z +Example script for end-to-end spacy KG pipeline. +Files: +examples/customize/build_graph/components/spacy_kg_pipeline.py diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index f758b9a86..ffce81970 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -176,7 +176,7 @@ tasks: Use pytest fixtures that skip the test if `en_core_web_sm` is not installed. Verify that real SpaCy output produces a populated Neo4jGraph and correct TextChunks respectively. - status: pending + status: in_progress dependencies: - task-005 - task-007 @@ -213,7 +213,7 @@ tasks: DataLoader → HierarchicalTextSplitter → SpacyEntityRelationExtractor → KGWriter. Include a comment block explaining the nlp extra install step (`pip install 'neo4j-graphrag[nlp]'` + `python -m spacy download en_core_web_sm`). - status: pending + status: completed dependencies: - task-009 acceptance_criteria: @@ -221,4 +221,4 @@ tasks: - "Script includes install instructions in a comment block at the top" - "Script constructs all four pipeline components and wires them together" - "Script includes a __main__ guard and can be run standalone" - completed_at: null + completed_at: "2026-06-24T17:00:00Z" diff --git a/examples/customize/build_graph/components/spacy_kg_pipeline.py b/examples/customize/build_graph/components/spacy_kg_pipeline.py new file mode 100644 index 000000000..240e9b1bb --- /dev/null +++ b/examples/customize/build_graph/components/spacy_kg_pipeline.py @@ -0,0 +1,111 @@ +# Install: pip install 'neo4j-graphrag[nlp]' +# Then: python -m spacy download en_core_web_sm +"""End-to-end SpaCy knowledge-graph pipeline example. + +Demonstrates how to wire four pipeline components together without an LLM: + + PdfLoader + → HierarchicalTextSplitter + → SpacyEntityRelationExtractor + → Neo4jWriter + +This example assumes a Neo4j instance is running locally. Update the +credentials below if needed. + +The pipeline extracts (subject, predicate, object) triples from the document +using SpaCy's dependency parser and named-entity recogniser and writes the +resulting nodes and relationships to Neo4j. +""" + +import asyncio +from pathlib import Path + +import neo4j + +from neo4j_graphrag.experimental.components.data_loader import PdfLoader +from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter +from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, +) +from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( + HierarchicalTextSplitter, +) + +# --------------------------------------------------------------------------- +# Configuration — update to match your environment +# --------------------------------------------------------------------------- + +URI = "neo4j://localhost:7687" +AUTH = ("neo4j", "password") +DATABASE = "neo4j" + +root_dir = Path(__file__).parents[4] +FILE_PATH = root_dir / "data" / "Harry Potter and the Chamber of Secrets Summary.pdf" + + +# --------------------------------------------------------------------------- +# Pipeline definition +# --------------------------------------------------------------------------- + + +async def run_pipeline(driver: neo4j.Driver) -> None: + """Load a PDF, split it, extract entities/relations, and write to Neo4j.""" + + # 1. Load the PDF into a single Document string. + loader = PdfLoader() + document = await loader.run(filepath=FILE_PATH) + + # 2. Split the document text into overlapping chunks, respecting section + # boundaries detected from Markdown-style headers. + splitter = HierarchicalTextSplitter( + # Maximum number of characters per chunk. + max_chunk_size=2048, + # Overlap between consecutive chunks to preserve context at boundaries. + chunk_overlap=200, + # Header detection strategy; 'markdown' looks for ATX-style # headers. + header_strategy="markdown", + # drop_verbless_sentences=True (default) filters noisy sentence fragments. + ) + chunks = await splitter.run(text=document.text) + + # 3. Extract entities and relationships from each chunk using SpaCy. + # No LLM is required — extraction is performed entirely locally. + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + # Use both dependency-based and proximity-based triple extraction. + use_linear_extractor=True, + # Add Chunk and Document nodes to the graph for provenance tracking. + create_lexical_graph=True, + # on_error=OnError.IGNORE (default) skips chunks that raise exceptions. + ) + graph = await extractor.run( + chunks=chunks, + document_info=document.document_info, + ) + + # 4. Write the resulting nodes and relationships to Neo4j. + writer = Neo4jWriter( + driver=driver, + neo4j_database=DATABASE, + # batch_size=1000, # tune for throughput on large documents + ) + result = await writer.run(graph=graph) + + print( + f"Wrote {result.nodes_upserted} nodes and " + f"{result.relationships_created} relationships to Neo4j." + ) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +async def main() -> None: + with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver: + await run_pipeline(driver) + + +if __name__ == "__main__": + asyncio.run(main()) From b433bc971f00259872fff82e1576176e5206c3b9 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 02:34:57 +0200 Subject: [PATCH 09/12] spacy-entity-relation-extractor-task-008: add integration tests for SpaCy components Add integration test files for SpacyEntityRelationExtractor and HierarchicalTextSplitter that exercise the real en_core_web_sm model. All tests are auto-skipped (not failed) when the model is not installed. --- ...ogress-spacy-entity-relation-extractor.txt | 4 + .../tasks-spacy-entity-relation-extractor.yml | 4 +- .../test_spacy_extractor_integration.py | 122 ++++++++++++++++++ .../test_hierarchical_splitter_integration.py | 99 ++++++++++++++ 4 files changed, 227 insertions(+), 2 deletions(-) create mode 100644 tests/unit/experimental/components/test_spacy_extractor_integration.py create mode 100644 tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt index d49893a1e..445fb024e 100644 --- a/.plans/progress-spacy-entity-relation-extractor.txt +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -40,3 +40,7 @@ Files: ~src/neo4j_graphrag/experimental/components/__init__.py, ~src/neo4j_graph ## TASK-010 — 2026-06-24T17:00:00Z Example script for end-to-end spacy KG pipeline. Files: +examples/customize/build_graph/components/spacy_kg_pipeline.py + +## TASK-008 — 2026-06-24T17:00:00Z +Integration tests for SpacyEntityRelationExtractor and HierarchicalTextSplitter (skip when en_core_web_sm absent). +Files: +tests/unit/experimental/components/test_spacy_extractor_integration.py, +tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py diff --git a/.plans/tasks-spacy-entity-relation-extractor.yml b/.plans/tasks-spacy-entity-relation-extractor.yml index ffce81970..87fc752ca 100644 --- a/.plans/tasks-spacy-entity-relation-extractor.yml +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -176,7 +176,7 @@ tasks: Use pytest fixtures that skip the test if `en_core_web_sm` is not installed. Verify that real SpaCy output produces a populated Neo4jGraph and correct TextChunks respectively. - status: in_progress + status: completed dependencies: - task-005 - task-007 @@ -185,7 +185,7 @@ tasks: - "With model installed, SpacyEntityRelationExtractor.run() returns Neo4jGraph with at least one node and one relationship for a 2-sentence input" - "With model installed, HierarchicalTextSplitter.run() returns correct chunks for a multi-section document" - "Confidence values on relationships are floats in [0, 1]" - completed_at: null + completed_at: "2026-06-24T17:00:00Z" - id: task-009 title: "Export new components from experimental __init__" diff --git a/tests/unit/experimental/components/test_spacy_extractor_integration.py b/tests/unit/experimental/components/test_spacy_extractor_integration.py new file mode 100644 index 000000000..3131865a4 --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_extractor_integration.py @@ -0,0 +1,122 @@ +# 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. +"""Integration tests for SpacyEntityRelationExtractor using the real en_core_web_sm model. + +All tests in this module are automatically skipped when either spaCy or the +``en_core_web_sm`` model is not installed. No mocking is performed — these +tests exercise the full NLP pipeline end-to-end. +""" + +from __future__ import annotations + +import pytest + +spacy = pytest.importorskip("spacy") + + +@pytest.fixture(scope="module") +def nlp(): # type: ignore[return] + try: + return spacy.load("en_core_web_sm") + except OSError: + pytest.skip( + "en_core_web_sm not installed — run: python -m spacy download en_core_web_sm" + ) + + +@pytest.mark.asyncio +async def test_run_returns_populated_graph(nlp) -> None: # noqa: ANN001 + """SpacyEntityRelationExtractor.run() returns a Neo4jGraph with nodes and relationships.""" + from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, + ) + from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks + + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks( + chunks=[ + TextChunk( + text="Apple acquired Beats Electronics. Tim Cook announced the deal.", + index=0, + ) + ] + ) + result = await extractor.run(chunks=chunks) + + assert len(result.nodes) >= 1, "Expected at least one entity node" + assert len(result.relationships) >= 1, "Expected at least one relationship" + + +@pytest.mark.asyncio +async def test_confidence_values_in_range(nlp) -> None: # noqa: ANN001 + """All relationship confidence values are floats in [0, 1].""" + from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, + ) + from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks + + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=True, + create_lexical_graph=False, + ) + chunks = TextChunks( + chunks=[ + TextChunk( + text="Apple acquired Beats Electronics. Tim Cook announced the deal.", + index=0, + ) + ] + ) + result = await extractor.run(chunks=chunks) + + for rel in result.relationships: + conf = rel.properties.get("confidence") + assert isinstance(conf, float), f"confidence should be float, got {type(conf)}" + assert 0.0 <= conf <= 1.0, f"confidence out of range [0,1]: {conf}" + + +@pytest.mark.asyncio +async def test_use_linear_extractor_false_no_related_to(nlp) -> None: # noqa: ANN001 + """use_linear_extractor=False produces no RELATED_TO relationships.""" + from neo4j_graphrag.experimental.components.spacy_entity_relation_extractor import ( + SpacyEntityRelationExtractor, + ) + from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks + + extractor = SpacyEntityRelationExtractor( + model="en_core_web_sm", + use_linear_extractor=False, + create_lexical_graph=False, + ) + chunks = TextChunks( + chunks=[ + TextChunk( + text="Apple acquired Beats Electronics. Tim Cook announced the deal.", + index=0, + ) + ] + ) + result = await extractor.run(chunks=chunks) + + related_to_rels = [r for r in result.relationships if r.type == "RELATED_TO"] + assert related_to_rels == [], ( + f"Expected no RELATED_TO relationships when use_linear_extractor=False, " + f"but got: {related_to_rels}" + ) diff --git a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py new file mode 100644 index 000000000..e2693880f --- /dev/null +++ b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py @@ -0,0 +1,99 @@ +# 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. +"""Integration tests for HierarchicalTextSplitter using the real en_core_web_sm model. + +All tests in this module are automatically skipped when either spaCy or the +``en_core_web_sm`` model is not installed. No mocking is performed — these +tests exercise the full NLP pipeline end-to-end. +""" + +from __future__ import annotations + +import pytest + +spacy = pytest.importorskip("spacy") + + +@pytest.fixture(scope="module") +def nlp(): # type: ignore[return] + try: + return spacy.load("en_core_web_sm") + except OSError: + pytest.skip( + "en_core_web_sm not installed — run: python -m spacy download en_core_web_sm" + ) + + +@pytest.mark.asyncio +async def test_markdown_split_real(nlp) -> None: # noqa: ANN001 + """HierarchicalTextSplitter with markdown strategy returns sequential chunks for a 3-section doc.""" + from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( + HierarchicalTextSplitter, + ) + + text = ( + "# Introduction\n" + "This section introduces the topic and provides background information.\n\n" + "# Methods\n" + "This section describes the experimental methods used in the study.\n\n" + "# Conclusion\n" + "This section summarises the findings and suggests future work.\n" + ) + + splitter = HierarchicalTextSplitter( + header_strategy="markdown", + max_chunk_size=200, + chunk_overlap=20, + drop_verbless_sentences=False, + ) + result = await splitter.run(text) + + assert len(result.chunks) >= 2, ( + f"Expected at least 2 chunks for a 3-section markdown doc, got {len(result.chunks)}" + ) + # Indices must be sequential starting from 0. + for i, chunk in enumerate(result.chunks): + assert chunk.index == i, f"chunk {i} has non-sequential index {chunk.index}" + + +@pytest.mark.asyncio +async def test_drop_verbless_sentences_real(nlp) -> None: # noqa: ANN001 + """drop_verbless_sentences=True drops verbless fragments using the real SpaCy model.""" + from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( + HierarchicalTextSplitter, + ) + + # "Overview" is a single-word verbless fragment. + # The second sentence contains a real verb ("covers"). + text = "Overview\n\nThis section covers the main concepts of the system in detail." + + splitter = HierarchicalTextSplitter( + header_strategy="blank_line", + max_chunk_size=500, + chunk_overlap=0, + drop_verbless_sentences=True, + model="en_core_web_sm", + ) + result = await splitter.run(text) + + # At least one chunk must survive. + assert len(result.chunks) >= 1, "Expected at least one chunk after filtering" + + # The verbless word "Overview" should not appear as a standalone sentence + # in any chunk, while the main sentence content should be present. + all_text = " ".join(chunk.text for chunk in result.chunks) + assert "covers" in all_text, ( + "Expected the verbal sentence to survive verbless-sentence filtering" + ) From fcb126f30bf2b5d2f62e7762968baf48c3ca7a92 Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 12:11:32 +0200 Subject: [PATCH 10/12] chore: split HierarchicalTextSplitter into companion PR #548 Remove hierarchical_splitter.py and its tests from this branch (now in PR #548). Update example script to reference the companion PR. --- .../components/spacy_kg_pipeline.py | 6 +- .../components/text_splitters/__init__.py | 23 - .../text_splitters/hierarchical_splitter.py | 452 --------------- .../test_hierarchical_splitter.py | 524 ------------------ .../test_hierarchical_splitter_integration.py | 99 ---- 5 files changed, 5 insertions(+), 1099 deletions(-) delete mode 100644 src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py delete mode 100644 tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py delete mode 100644 tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py diff --git a/examples/customize/build_graph/components/spacy_kg_pipeline.py b/examples/customize/build_graph/components/spacy_kg_pipeline.py index 240e9b1bb..2d91a8b59 100644 --- a/examples/customize/build_graph/components/spacy_kg_pipeline.py +++ b/examples/customize/build_graph/components/spacy_kg_pipeline.py @@ -1,11 +1,15 @@ # Install: pip install 'neo4j-graphrag[nlp]' # Then: python -m spacy download en_core_web_sm +# +# NOTE: HierarchicalTextSplitter is implemented in the companion PR #548. +# Until that PR is merged, install it from the feat/hierarchical-splitter branch: +# pip install git+https://github.com/neo4j/neo4j-graphrag-python@feat/hierarchical-splitter """End-to-end SpaCy knowledge-graph pipeline example. Demonstrates how to wire four pipeline components together without an LLM: PdfLoader - → HierarchicalTextSplitter + → HierarchicalTextSplitter (see PR #548) → SpacyEntityRelationExtractor → Neo4jWriter diff --git a/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py b/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py index d128c3410..e69de29bb 100644 --- a/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py +++ b/src/neo4j_graphrag/experimental/components/text_splitters/__init__.py @@ -1,23 +0,0 @@ -# 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. - -try: - from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( - HierarchicalTextSplitter, - ) - - __all__ = ["HierarchicalTextSplitter"] -except ImportError: - pass diff --git a/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py b/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py deleted file mode 100644 index 2e8758b5b..000000000 --- a/src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py +++ /dev/null @@ -1,452 +0,0 @@ -# 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. -"""Hierarchical text splitter that detects section boundaries before chunking. - -Supported header strategies: - -* ``"markdown"`` — lines starting with ``#`` (ATX Markdown headers). -* ``"capitalization"`` — short Title Case or ALL_CAPS lines without terminal - punctuation (appropriate for plain-text output from loaders like - LiteParseLoader). -* ``"blank_line"`` — short lines that are surrounded by blank lines on both - sides (a common plain-text section marker). -* ``"spacy_verbless"`` — SpaCy-parsed sentences that are short, contain no - verb, and are followed by a longer sentence. - -All strategies produce a list of *sections* (contiguous text blocks). Each -section is then emitted as a single chunk when it fits within ``max_chunk_size``, -or recursively split with ``chunk_overlap`` when it is larger. - -Optionally, when ``drop_verbless_sentences=True`` (default), SpaCy is used to -remove sentences with no verb token from every emitted chunk. -""" - -from __future__ import annotations - -import re -from typing import Any, Optional - -from pydantic import validate_call - -from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter -from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks - -# --------------------------------------------------------------------------- -# Header-detection helpers -# --------------------------------------------------------------------------- - -_MARKDOWN_HEADER_RE = re.compile(r"^#{1,6}\s+\S", re.MULTILINE) - -# Title Case: starts with capital, no terminal punctuation at end of line, -# short (≤ 80 chars), and at least half the words are Title-cased. -_TITLECASE_RE = re.compile(r"^[A-Z][^\n]{0,78}$") -_TERMINAL_PUNCT_RE = re.compile(r"[.!?;,]$") - -# ALL_CAPS line: all uppercase letters / spaces / digits -_ALLCAPS_RE = re.compile(r"^[A-Z0-9][A-Z0-9 \t\-:]+$") - -# Blank-line boundary: a line that is non-empty, short, and preceded/followed -# by blank lines (handled at section-split level, not a single regex). -_SHORT_LINE_MAX = 80 - - -def _is_title_case(line: str) -> bool: - """Return True when a line looks like a Title Case heading.""" - if not _TITLECASE_RE.match(line): - return False - if _TERMINAL_PUNCT_RE.search(line): - return False - words = line.split() - if len(words) == 0: - return False - capitalised = sum(1 for w in words if w and w[0].isupper()) - return capitalised / len(words) >= 0.5 - - -def _is_allcaps(line: str) -> bool: - """Return True when a line is ALL CAPS without terminal punctuation.""" - if not _ALLCAPS_RE.match(line): - return False - return not _TERMINAL_PUNCT_RE.search(line) - - -def _split_at_markdown_headers(text: str) -> list[str]: - """Split *text* into sections at Markdown ATX header lines (``#``, ``##``, …).""" - # We split before each header line so that the header stays with its section. - lines = text.splitlines(keepends=True) - sections: list[str] = [] - current: list[str] = [] - for line in lines: - if re.match(r"^#{1,6}\s+\S", line) and current: - sections.append("".join(current)) - current = [line] - else: - current.append(line) - if current: - sections.append("".join(current)) - return [s for s in sections if s.strip()] - - -def _split_at_capitalization(text: str) -> list[str]: - """Split *text* into sections at Title Case or ALL_CAPS heading lines.""" - lines = text.splitlines(keepends=True) - sections: list[str] = [] - current: list[str] = [] - for line in lines: - stripped = line.rstrip("\r\n") - is_header = ( - len(stripped) <= _SHORT_LINE_MAX - and stripped.strip() - and (_is_title_case(stripped.strip()) or _is_allcaps(stripped.strip())) - ) - if is_header and current: - sections.append("".join(current)) - current = [line] - else: - current.append(line) - if current: - sections.append("".join(current)) - return [s for s in sections if s.strip()] - - -def _split_at_blank_line(text: str) -> list[str]: - """Split *text* into sections at short lines surrounded by blank lines. - - A line qualifies as a section header when it: - - * Is non-empty and at most 80 characters. - * Contains no terminal punctuation (not a regular sentence). - * Has no more than 6 words (avoids treating body sentences as headers). - * Is preceded by a blank line (or is the very first non-blank line). - * Is followed by a blank line. - """ - lines = text.splitlines(keepends=True) - n = len(lines) - sections: list[str] = [] - current: list[str] = [] - - i = 0 - while i < n: - line = lines[i] - stripped = line.rstrip("\r\n").strip() - # A "blank-line boundary" header: non-empty, short, preceded by blank, - # followed by blank, few words, no terminal punctuation. - prev_blank = (i == 0) or (not lines[i - 1].strip()) - next_blank = (i + 1 >= n) or (not lines[i + 1].strip()) - word_count = len(stripped.split()) if stripped else 0 - is_header = ( - stripped - and len(stripped) <= _SHORT_LINE_MAX - and not _TERMINAL_PUNCT_RE.search(stripped) - and word_count <= 6 - and prev_blank - and next_blank - ) - if is_header and current: - sections.append("".join(current)) - current = [line] - else: - current.append(line) - i += 1 - - if current: - sections.append("".join(current)) - return [s for s in sections if s.strip()] - - -def _split_at_spacy_verbless(text: str, nlp: Any) -> list[str]: - """Split *text* into sections at SpaCy-detected verbless heading sentences. - - A sentence qualifies as a heading when it: - - * Has at most 80 characters. - * Contains no verb token (POS tag ``VERB`` or ``AUX``). - * Is immediately followed by a longer sentence (> 80 chars). - """ - doc = nlp(text) - sentences = list(doc.sents) - if not sentences: - return [text] if text.strip() else [] - - # Identify which sentences are "headers". - is_header = [False] * len(sentences) - for idx, sent in enumerate(sentences): - sent_text = sent.text.strip() - if len(sent_text) > _SHORT_LINE_MAX: - continue - has_verb = any(tok.pos_ in ("VERB", "AUX") for tok in sent) - if has_verb: - continue - # Must be followed by a longer sentence. - if ( - idx + 1 < len(sentences) - and len(sentences[idx + 1].text.strip()) > _SHORT_LINE_MAX - ): - is_header[idx] = True - - # Build sections: split before each header sentence (except the first). - sections: list[str] = [] - current_parts: list[str] = [] - for idx, sent in enumerate(sentences): - if is_header[idx] and current_parts: - sections.append(" ".join(current_parts)) - current_parts = [sent.text] - else: - current_parts.append(sent.text) - - if current_parts: - sections.append(" ".join(current_parts)) - - return [s for s in sections if s.strip()] - - -# --------------------------------------------------------------------------- -# Overlap-based character splitter (for sections larger than max_chunk_size) -# --------------------------------------------------------------------------- - - -def _split_with_overlap(text: str, max_size: int, overlap: int) -> list[str]: - """Split *text* into character-level chunks of at most *max_size* with - *overlap* characters carried over from the previous chunk. - - Splits are attempted at whitespace boundaries to avoid cutting words. - """ - if not text: - return [] - chunks: list[str] = [] - start = 0 - length = len(text) - step = max(1, max_size - overlap) - - while start < length: - end = min(start + max_size, length) - chunk = text[start:end] - - # Prefer to cut at a whitespace boundary when not at end of text. - if end < length: - # Walk backwards to find a space. - cut = end - while cut > start and not text[cut - 1].isspace(): - cut -= 1 - if cut > start: - end = cut - chunk = text[start:end] - - chunks.append(chunk) - - # Advance by step, ensuring we always make progress. - next_start = start + step - if next_start <= start: - next_start = start + 1 - start = next_start - - return chunks - - -# --------------------------------------------------------------------------- -# Verb-filter using SpaCy -# --------------------------------------------------------------------------- - - -def _drop_verbless_sentences(text: str, nlp: Any) -> str: - """Remove sentences with no verb token from *text* using SpaCy. - - A sentence is considered *verbless* when it contains no token whose - part-of-speech tag is ``VERB`` or ``AUX``. - """ - doc = nlp(text) - kept: list[str] = [] - for sent in doc.sents: - has_verb = any(tok.pos_ in ("VERB", "AUX") for tok in sent) - if has_verb: - kept.append(sent.text) - return " ".join(kept) - - -# --------------------------------------------------------------------------- -# Valid strategy names -# --------------------------------------------------------------------------- - -_VALID_STRATEGIES = frozenset( - {"markdown", "capitalization", "blank_line", "spacy_verbless"} -) - - -# --------------------------------------------------------------------------- -# Main component -# --------------------------------------------------------------------------- - - -class HierarchicalTextSplitter(TextSplitter): - """Splits text by first detecting section boundaries then chunking each section. - - Args: - max_chunk_size (int): Maximum number of characters per output chunk. - Defaults to 2048. - chunk_overlap (int): Characters of overlap between consecutive chunks - when a section must be further split. Must be less than - ``max_chunk_size``. Defaults to 200. - header_strategy (str): How to detect section boundaries. One of: - - * ``"markdown"`` — Markdown ATX header lines (``#``, ``##``, …). - * ``"capitalization"`` — short Title Case or ALL_CAPS lines without - terminal punctuation. - * ``"blank_line"`` — short lines surrounded by blank lines on both - sides. - * ``"spacy_verbless"`` — SpaCy-detected short verbless sentences - that precede a longer sentence. - - model (str): SpaCy model name loaded when *header_strategy* is - ``"spacy_verbless"`` or *drop_verbless_sentences* is ``True``. - Defaults to ``"en_core_web_sm"``. - drop_verbless_sentences (bool): When ``True`` (default), SpaCy is used - to remove verbless sentences from every emitted chunk. Note that - this default value causes SpaCy to be loaded at construction time - regardless of the chosen *header_strategy* — install - ``neo4j-graphrag[nlp]`` when using the default, or explicitly pass - ``drop_verbless_sentences=False`` to avoid the SpaCy dependency. - - Example: - - .. code-block:: python - - from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( - HierarchicalTextSplitter, - ) - from neo4j_graphrag.experimental.pipeline import Pipeline - - pipeline = Pipeline() - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=200, - header_strategy="markdown", - ) - pipeline.add_component(splitter, "text_splitter") - """ - - @validate_call - def __init__( - self, - max_chunk_size: int = 2048, - chunk_overlap: int = 200, - header_strategy: str = "markdown", - model: str = "en_core_web_sm", - drop_verbless_sentences: bool = True, - ) -> None: - if max_chunk_size <= 0: - raise ValueError("max_chunk_size must be strictly greater than 0") - if chunk_overlap < 0: - raise ValueError("chunk_overlap must be >= 0") - if chunk_overlap >= max_chunk_size: - raise ValueError("chunk_overlap must be strictly less than max_chunk_size") - if header_strategy not in _VALID_STRATEGIES: - raise ValueError( - f"header_strategy must be one of {sorted(_VALID_STRATEGIES)}, " - f"got {header_strategy!r}" - ) - - self.max_chunk_size = max_chunk_size - self.chunk_overlap = chunk_overlap - self.header_strategy = header_strategy - self.model = model - self.drop_verbless_sentences = drop_verbless_sentences - - # Pre-load SpaCy only when needed. - self._nlp: Optional[Any] = None - needs_spacy = header_strategy == "spacy_verbless" or drop_verbless_sentences - if needs_spacy: - self._nlp = self._load_spacy(model) - - @staticmethod - def _load_spacy(model: str) -> Any: - """Load a SpaCy model, raising a clear error when SpaCy is missing.""" - try: - import spacy # noqa: PLC0415 - except ImportError as exc: - raise ImportError( - "SpaCy is required for this configuration of HierarchicalTextSplitter. " - "Install it with: pip install 'neo4j-graphrag[nlp]'" - ) from exc - try: - return spacy.load(model) - except OSError as exc: - raise ValueError( - f"SpaCy model {model!r} is not installed. " - f"Download it with: python -m spacy download {model}" - ) from exc - - def _detect_sections(self, text: str) -> list[str]: - """Detect section boundaries and return a list of section strings.""" - strategy = self.header_strategy - if strategy == "markdown": - sections = _split_at_markdown_headers(text) - elif strategy == "capitalization": - sections = _split_at_capitalization(text) - elif strategy == "blank_line": - sections = _split_at_blank_line(text) - else: # "spacy_verbless" - if self._nlp is None: - raise RuntimeError( - "SpaCy model not loaded for 'spacy_verbless' strategy; this is a bug" - ) - sections = _split_at_spacy_verbless(text, self._nlp) - - # Fallback: if no sections were detected, treat the whole text as one. - if not sections: - sections = [text] if text.strip() else [] - return sections - - def _chunk_section(self, section_text: str) -> list[str]: - """Return one or more raw text chunks for a single *section_text*.""" - if len(section_text) <= self.max_chunk_size: - return [section_text] - return _split_with_overlap( - section_text, self.max_chunk_size, self.chunk_overlap - ) - - def _filter_verbless(self, text: str) -> str: - """Apply the verbless-sentence filter if enabled.""" - if not self.drop_verbless_sentences or self._nlp is None: - return text - filtered = _drop_verbless_sentences(text, self._nlp) - # Fall back to original text when filtering removes everything. - return filtered if filtered.strip() else text - - @validate_call - async def run(self, text: str) -> TextChunks: - """Split *text* into hierarchical chunks. - - Args: - text (str): The text to be split. - - Returns: - TextChunks: A list of chunks with sequential index values starting - from 0. - """ - if not text.strip(): - return TextChunks(chunks=[]) - - sections = self._detect_sections(text) - raw_chunks: list[str] = [] - for section in sections: - raw_chunks.extend(self._chunk_section(section)) - - chunks: list[TextChunk] = [] - for raw in raw_chunks: - filtered = self._filter_verbless(raw) - if filtered.strip(): - chunks.append(TextChunk(text=filtered, index=len(chunks))) - - return TextChunks(chunks=chunks) diff --git a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py deleted file mode 100644 index ba06e02d6..000000000 --- a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter.py +++ /dev/null @@ -1,524 +0,0 @@ -# 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. -"""Unit tests for HierarchicalTextSplitter. - -SpaCy is not downloaded in these tests. Where the splitter would normally -load a model (`drop_verbless_sentences=True` or `header_strategy="spacy_verbless"`), -the tests patch `spacy.load` with a lightweight fake nlp object so that no -network access or model installation is required. -""" - -from __future__ import annotations - -from typing import Any -from unittest.mock import MagicMock, patch - -import pytest - -# Skip the entire module when spaCy is not installed at all. -spacy = pytest.importorskip("spacy") - -from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( # noqa: E402 - HierarchicalTextSplitter, -) - - -# --------------------------------------------------------------------------- -# Helpers for building fake SpaCy objects without a real model -# --------------------------------------------------------------------------- - - -def _make_fake_token(text: str, pos: str) -> MagicMock: - """Return a MagicMock that looks like a spaCy Token.""" - tok = MagicMock() - tok.text = text - tok.pos_ = pos - return tok - - -def _make_fake_sent(text: str, tokens: list[MagicMock]) -> MagicMock: - """Return a MagicMock that looks like a spaCy Span (sentence). - - Uses ``side_effect`` instead of ``return_value`` so that each call to - ``iter(sent)`` creates a *fresh* iterator — important if the sentence is - iterated more than once (e.g. across multiple ``run()`` calls or in future - multi-pass tests). - """ - sent = MagicMock() - sent.text = text - sent.__iter__ = MagicMock(side_effect=lambda: iter(tokens)) - return sent - - -def _make_fake_doc(sentences: list[MagicMock]) -> MagicMock: - """Return a MagicMock that looks like a spaCy Doc with .sents.""" - doc = MagicMock() - doc.sents = sentences - doc.__iter__ = MagicMock(return_value=iter([])) - return doc - - -def _make_nlp_returning_doc(sentences: list[MagicMock]) -> MagicMock: - """Return a callable MagicMock that acts as a spaCy nlp() pipeline.""" - nlp = MagicMock() - nlp.return_value = _make_fake_doc(sentences) - return nlp - - -# --------------------------------------------------------------------------- -# Tests: header_strategy="markdown" -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_markdown_two_sections_produce_two_chunks() -> None: - """Two Markdown sections produce exactly two chunks.""" - text = "# Introduction\nThis is the intro section.\n# Conclusion\nThis is the conclusion." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 2 - assert "Introduction" in result.chunks[0].text - assert "Conclusion" in result.chunks[1].text - - -@pytest.mark.asyncio -async def test_markdown_single_section_produces_one_chunk() -> None: - """Text with no Markdown headers is treated as one section.""" - text = "No headers here. Just a single paragraph of text." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 1 - - -@pytest.mark.asyncio -async def test_markdown_three_sections_sequential_indices() -> None: - """Three Markdown sections produce chunks with sequential indices 0, 1, 2.""" - text = "# A\nSection A body.\n# B\nSection B body.\n# C\nSection C body." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 3 - for i, chunk in enumerate(result.chunks): - assert chunk.index == i - - -# --------------------------------------------------------------------------- -# Tests: header_strategy="capitalization" -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_capitalization_allcaps_header_splits() -> None: - """ALL_CAPS lines without terminal punctuation act as section headers.""" - text = "INTRODUCTION\nThis paragraph describes the introduction.\nCONCLUSION\nThis paragraph wraps things up." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="capitalization", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 2 - assert "INTRODUCTION" in result.chunks[0].text - assert "CONCLUSION" in result.chunks[1].text - - -@pytest.mark.asyncio -async def test_capitalization_title_case_header_splits() -> None: - """Title Case lines without terminal punctuation act as section headers.""" - text = "First Section Title\nContent for the first section.\nSecond Section Title\nContent for the second section." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="capitalization", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 2 - - -@pytest.mark.asyncio -async def test_capitalization_indices_sequential() -> None: - """Chunk indices are sequential starting from 0.""" - text = "PART ONE\nBody of part one.\nPART TWO\nBody of part two.\nPART THREE\nBody of part three." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="capitalization", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - for i, chunk in enumerate(result.chunks): - assert chunk.index == i, f"chunk {i} has index {chunk.index}" - - -# --------------------------------------------------------------------------- -# Tests: header_strategy="blank_line" -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_blank_line_short_surrounded_header_splits() -> None: - """Short lines surrounded by blank lines act as section headers.""" - text = "\nOverview\n\nThis section covers the overview of the system.\n\nDetails\n\nThis section covers the fine-grained details." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="blank_line", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 2 - assert "Overview" in result.chunks[0].text - assert "Details" in result.chunks[1].text - - -@pytest.mark.asyncio -async def test_blank_line_indices_sequential() -> None: - """Chunk indices are sequential starting from 0 for blank_line strategy.""" - text = "\nPart A\n\nContent for part A goes here.\n\nPart B\n\nContent for part B goes here." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="blank_line", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - for i, chunk in enumerate(result.chunks): - assert chunk.index == i - - -# --------------------------------------------------------------------------- -# Tests: overlap splitting for large sections -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_large_section_split_with_overlap() -> None: - """A section larger than max_chunk_size is split into multiple chunks - with the last N characters of chunk K equal to the first N characters - of chunk K+1. - - Strategy: use no header strategy (treat whole text as one section) via a - single Markdown section; make the body very long and max_chunk_size small - relative to the body so that we get several chunks all in the body content, - and the overlap check applies to non-header chunks. - """ - # The section body is pure 'x' characters with no whitespace. - # With max_size=50, overlap=20, step=30: - # chunk 0: body[0:50] = 'x'*50 - # chunk 1: body[30:80] = 'x'*50, so overlap = body[30:50] = 'x'*20 - # We avoid the Markdown header prefix eating into the overlap window by - # using a plain-text body (no Markdown header at all, so the whole text - # is one section) and using drop_verbless_sentences=False. - body = "x" * 500 - overlap = 20 - max_size = 50 - splitter = HierarchicalTextSplitter( - max_chunk_size=max_size, - chunk_overlap=overlap, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(body) - chunks = result.chunks - assert len(chunks) > 1 - - # Verify overlap: last `overlap` chars of chunk[k] == first `overlap` chars of chunk[k+1]. - for k in range(len(chunks) - 1): - tail = chunks[k].text[-overlap:] - head = chunks[k + 1].text[:overlap] - assert tail == head, ( - f"Overlap mismatch between chunk {k} and {k + 1}: " - f"tail={tail!r}, head={head!r}" - ) - - -@pytest.mark.asyncio -async def test_small_section_emitted_as_single_chunk() -> None: - """A section shorter than max_chunk_size is emitted as a single chunk.""" - text = "# Tiny\nShort body." - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 1 - assert "Tiny" in result.chunks[0].text - assert "Short body" in result.chunks[0].text - - -@pytest.mark.asyncio -async def test_overlap_chunk_indices_sequential() -> None: - """Indices remain sequential when a section is split due to size.""" - body = "y" * 300 - text = f"# Big Section\n{body}" - splitter = HierarchicalTextSplitter( - max_chunk_size=60, - chunk_overlap=10, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) > 1 - for i, chunk in enumerate(result.chunks): - assert chunk.index == i - - -# --------------------------------------------------------------------------- -# Tests: drop_verbless_sentences=True (SpaCy mocked) -# --------------------------------------------------------------------------- - - -def _make_spacy_nlp_with_verbless() -> MagicMock: - """Return a fake nlp() that drops verbless sentences when called. - - Sentence 1 (verbless): "No verb here" — tokens have no VERB/AUX. - Sentence 2 (with verb): "The dog runs fast" — one VERB token. - """ - # Tokens for sentence 1 (verbless: no VERB/AUX tags) - sent1_tokens = [ - _make_fake_token("No", "DET"), - _make_fake_token("verb", "NOUN"), - _make_fake_token("here", "ADV"), - ] - sent1 = _make_fake_sent("No verb here", sent1_tokens) - - # Tokens for sentence 2 (contains a VERB) - sent2_tokens = [ - _make_fake_token("The", "DET"), - _make_fake_token("dog", "NOUN"), - _make_fake_token("runs", "VERB"), - _make_fake_token("fast", "ADV"), - ] - sent2 = _make_fake_sent("The dog runs fast", sent2_tokens) - - return _make_nlp_returning_doc([sent1, sent2]) - - -@pytest.mark.asyncio -async def test_drop_verbless_removes_verbless_sentence() -> None: - """When drop_verbless_sentences=True, sentences with no verb are removed.""" - fake_nlp = _make_spacy_nlp_with_verbless() - - with patch("spacy.load", return_value=fake_nlp): - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=True, - model="en_core_web_sm", - ) - - text = "# Section\nNo verb here. The dog runs fast." - result = await splitter.run(text) - - assert len(result.chunks) == 1 - chunk_text = result.chunks[0].text - # Verbless sentence should be gone; verbal sentence should remain. - assert "The dog runs fast" in chunk_text - assert "No verb here" not in chunk_text - - -@pytest.mark.asyncio -async def test_drop_verbless_keeps_verbal_sentences() -> None: - """When all sentences contain a verb, no text is removed.""" - sent_tokens = [ - _make_fake_token("She", "PRON"), - _make_fake_token("walks", "VERB"), - _make_fake_token("home", "NOUN"), - ] - sent = _make_fake_sent("She walks home", sent_tokens) - fake_nlp = _make_nlp_returning_doc([sent]) - - with patch("spacy.load", return_value=fake_nlp): - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=True, - model="en_core_web_sm", - ) - - text = "# Section\nShe walks home." - result = await splitter.run(text) - - assert len(result.chunks) == 1 - assert "She walks home" in result.chunks[0].text - - -# --------------------------------------------------------------------------- -# Tests: chunk index sequencing across multiple sections -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_indices_sequential_across_multiple_sections() -> None: - """Indices are global — they continue from the last chunk of the previous section.""" - text = ( - "# Alpha\nFirst section body.\n" - "# Beta\nSecond section body.\n" - "# Gamma\nThird section body." - ) - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - assert len(result.chunks) == 3 - for expected_index, chunk in enumerate(result.chunks): - assert chunk.index == expected_index - - -# --------------------------------------------------------------------------- -# Tests: header_strategy="spacy_verbless" -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_spacy_verbless_strategy_splits_at_verbless_heading() -> None: - """header_strategy="spacy_verbless" uses SpaCy to detect verbless headings. - - A short sentence with no verb that precedes a longer sentence is treated as - a section heading and causes a split. This test mocks the nlp pipeline so - no model is downloaded. - """ - # Sentence 1 (verbless heading, short ≤ 80 chars): "Introduction" - sent1_tokens = [_make_fake_token("Introduction", "NOUN")] - sent1 = _make_fake_sent("Introduction", sent1_tokens) - - # Sentence 2 (long body sentence, > 80 chars, contains VERB): - long_body = "This section covers all the foundational concepts you need to understand before proceeding further." - sent2_tokens = [ - _make_fake_token("This", "DET"), - _make_fake_token("section", "NOUN"), - _make_fake_token("covers", "VERB"), - ] - sent2 = _make_fake_sent(long_body, sent2_tokens) - - # Sentence 3 (verbless heading): "Conclusion" - sent3_tokens = [_make_fake_token("Conclusion", "NOUN")] - sent3 = _make_fake_sent("Conclusion", sent3_tokens) - - # Sentence 4 (another long body): - long_body2 = "This final section wraps up all the topics discussed and provides closing remarks for the reader." - sent4_tokens = [ - _make_fake_token("This", "DET"), - _make_fake_token("section", "NOUN"), - _make_fake_token("wraps", "VERB"), - ] - sent4 = _make_fake_sent(long_body2, sent4_tokens) - - fake_nlp = _make_nlp_returning_doc([sent1, sent2, sent3, sent4]) - - with patch("spacy.load", return_value=fake_nlp): - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="spacy_verbless", - drop_verbless_sentences=False, - model="en_core_web_sm", - ) - - # Use arbitrary text — the nlp mock controls what sents are returned. - text = "Introduction. " + long_body + " Conclusion. " + long_body2 - result = await splitter.run(text) - - # Two verbless headings (sent1, sent3) each trigger a split → 2 sections. - assert len(result.chunks) == 2 - assert result.chunks[0].index == 0 - assert result.chunks[1].index == 1 - - -# --------------------------------------------------------------------------- -# Tests: edge cases -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_empty_text_returns_no_chunks() -> None: - """Empty input produces an empty chunk list.""" - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run("") - assert result.chunks == [] - - -@pytest.mark.asyncio -async def test_whitespace_only_returns_no_chunks() -> None: - """Whitespace-only input produces an empty chunk list.""" - splitter = HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - result = await splitter.run(" \n\t\n ") - assert result.chunks == [] - - -def test_invalid_header_strategy_raises() -> None: - """An unrecognised header_strategy raises ValueError at construction time.""" - with pytest.raises(ValueError, match="header_strategy must be one of"): - HierarchicalTextSplitter( - max_chunk_size=2048, - chunk_overlap=0, - header_strategy="unknown_strategy", - drop_verbless_sentences=False, - ) - - -def test_chunk_overlap_ge_max_chunk_size_raises() -> None: - """chunk_overlap >= max_chunk_size raises ValueError.""" - with pytest.raises(ValueError, match="chunk_overlap must be strictly less than max_chunk_size"): - HierarchicalTextSplitter( - max_chunk_size=100, - chunk_overlap=100, - header_strategy="markdown", - drop_verbless_sentences=False, - ) - - -def test_max_chunk_size_zero_raises() -> None: - """max_chunk_size=0 raises ValueError.""" - with pytest.raises(ValueError, match="max_chunk_size must be strictly greater than 0"): - HierarchicalTextSplitter( - max_chunk_size=0, - chunk_overlap=0, - header_strategy="markdown", - drop_verbless_sentences=False, - ) diff --git a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py b/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py deleted file mode 100644 index e2693880f..000000000 --- a/tests/unit/experimental/components/text_splitters/test_hierarchical_splitter_integration.py +++ /dev/null @@ -1,99 +0,0 @@ -# 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. -"""Integration tests for HierarchicalTextSplitter using the real en_core_web_sm model. - -All tests in this module are automatically skipped when either spaCy or the -``en_core_web_sm`` model is not installed. No mocking is performed — these -tests exercise the full NLP pipeline end-to-end. -""" - -from __future__ import annotations - -import pytest - -spacy = pytest.importorskip("spacy") - - -@pytest.fixture(scope="module") -def nlp(): # type: ignore[return] - try: - return spacy.load("en_core_web_sm") - except OSError: - pytest.skip( - "en_core_web_sm not installed — run: python -m spacy download en_core_web_sm" - ) - - -@pytest.mark.asyncio -async def test_markdown_split_real(nlp) -> None: # noqa: ANN001 - """HierarchicalTextSplitter with markdown strategy returns sequential chunks for a 3-section doc.""" - from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( - HierarchicalTextSplitter, - ) - - text = ( - "# Introduction\n" - "This section introduces the topic and provides background information.\n\n" - "# Methods\n" - "This section describes the experimental methods used in the study.\n\n" - "# Conclusion\n" - "This section summarises the findings and suggests future work.\n" - ) - - splitter = HierarchicalTextSplitter( - header_strategy="markdown", - max_chunk_size=200, - chunk_overlap=20, - drop_verbless_sentences=False, - ) - result = await splitter.run(text) - - assert len(result.chunks) >= 2, ( - f"Expected at least 2 chunks for a 3-section markdown doc, got {len(result.chunks)}" - ) - # Indices must be sequential starting from 0. - for i, chunk in enumerate(result.chunks): - assert chunk.index == i, f"chunk {i} has non-sequential index {chunk.index}" - - -@pytest.mark.asyncio -async def test_drop_verbless_sentences_real(nlp) -> None: # noqa: ANN001 - """drop_verbless_sentences=True drops verbless fragments using the real SpaCy model.""" - from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter import ( - HierarchicalTextSplitter, - ) - - # "Overview" is a single-word verbless fragment. - # The second sentence contains a real verb ("covers"). - text = "Overview\n\nThis section covers the main concepts of the system in detail." - - splitter = HierarchicalTextSplitter( - header_strategy="blank_line", - max_chunk_size=500, - chunk_overlap=0, - drop_verbless_sentences=True, - model="en_core_web_sm", - ) - result = await splitter.run(text) - - # At least one chunk must survive. - assert len(result.chunks) >= 1, "Expected at least one chunk after filtering" - - # The verbless word "Overview" should not appear as a standalone sentence - # in any chunk, while the main sentence content should be present. - all_text = " ".join(chunk.text for chunk in result.chunks) - assert "covers" in all_text, ( - "Expected the verbal sentence to survive verbless-sentence filtering" - ) From c1f0688b27abaea396dbf8c0d537332918aebf4d Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 12:17:51 +0200 Subject: [PATCH 11/12] ci: download en_core_web_sm in PR workflow to run SpaCy integration tests --- .github/workflows/pr.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index ccacc3987..a92563266 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -33,6 +33,9 @@ jobs: - name: Install dependencies if: steps.cached-venv.outputs.cache-hit != 'true' run: uv sync --frozen --all-extras --group dev + - name: Download SpaCy model + if: matrix.python-version != '3.14' + run: uv run --frozen python -m spacy download en_core_web_sm - name: Check format and linting run: | uv run --frozen ruff check . From f0d32a42b2082ebb6412e49c393f24c5c103314c Mon Sep 17 00:00:00 2001 From: Michael Hunger Date: Wed, 24 Jun 2026 12:18:35 +0200 Subject: [PATCH 12/12] Revert "ci: download en_core_web_sm in PR workflow to run SpaCy integration tests" This reverts commit 208d5694b57a6ed98f1c7001a455ad0382c42c38. --- .github/workflows/pr.yaml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a92563266..ccacc3987 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -33,9 +33,6 @@ jobs: - name: Install dependencies if: steps.cached-venv.outputs.cache-hit != 'true' run: uv sync --frozen --all-extras --group dev - - name: Download SpaCy model - if: matrix.python-version != '3.14' - run: uv run --frozen python -m spacy download en_core_web_sm - name: Check format and linting run: | uv run --frozen ruff check .