diff --git a/.plans/progress-spacy-entity-relation-extractor.txt b/.plans/progress-spacy-entity-relation-extractor.txt new file mode 100644 index 000000000..445fb024e --- /dev/null +++ b/.plans/progress-spacy-entity-relation-extractor.txt @@ -0,0 +1,46 @@ +## 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 +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 +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 +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. + +## 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 + +## 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 new file mode 100644 index 000000000..87fc752ca --- /dev/null +++ b/.plans/tasks-spacy-entity-relation-extractor.yml @@ -0,0 +1,224 @@ +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: completed + completed_at: "2026-06-24T10:00:00Z" + 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)" + + - 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: completed + completed_at: "2026-06-24T10:00:00Z" + 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" + 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" + 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) 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: completed + completed_at: "2026-06-24T15:00:00Z" + 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'" + - "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" + + - 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: completed + 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: "2026-06-24T16:00:00Z" + + - id: task-006 + title: "Implement HierarchicalTextSplitter component" + description: | + Create `src/neo4j_graphrag/experimental/components/text_splitters/hierarchical_splitter.py` + 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: completed + completed_at: "2026-06-24T12:00:00Z" + dependencies: + - task-001 + acceptance_criteria: + - "Class is importable from neo4j_graphrag.experimental.components.text_splitters.hierarchical_splitter" + - "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 (SpaCy mocked in unit tests)" + - "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: completed + 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: "2026-06-24T14:00:00Z" + + - 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: completed + 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: "2026-06-24T17:00:00Z" + + - 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: completed + completed_at: "2026-06-24T16:00:00Z" + 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" + + - 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: completed + 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: "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..2d91a8b59 --- /dev/null +++ b/examples/customize/build_graph/components/spacy_kg_pipeline.py @@ -0,0 +1,115 @@ +# 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 (see PR #548) + → 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()) 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/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 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/src/neo4j_graphrag/experimental/components/spacy_utils.py b/src/neo4j_graphrag/experimental/components/spacy_utils.py new file mode 100644 index 000000000..a84390294 --- /dev/null +++ b/src/neo4j_graphrag/experimental/components/spacy_utils.py @@ -0,0 +1,341 @@ +# 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 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 +# --------------------------------------------------------------------------- + +_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() + + +# --------------------------------------------------------------------------- +# 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 +# --------------------------------------------------------------------------- + +#: 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_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_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) 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/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 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..4497012a9 --- /dev/null +++ b/tests/unit/experimental/components/test_spacy_utils.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. +"""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 + +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] 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" },