diff --git a/graphrag_sdk/src/graphrag_sdk/__init__.py b/graphrag_sdk/src/graphrag_sdk/__init__.py index 17251be1..7b9fb61a 100644 --- a/graphrag_sdk/src/graphrag_sdk/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/__init__.py @@ -91,11 +91,14 @@ FastCorefResolver, ) from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( + CompositeExtractor, EntityExtractor, GLiNERExtractor, LLMExtractor, + SpacyExtractor, ) from graphrag_sdk.ingestion.extraction_strategies.graph_extraction import ( + DEFAULT_RELATION_TYPES, GraphExtraction, ) from graphrag_sdk.ingestion.loaders.base import LoaderStrategy @@ -130,6 +133,7 @@ from graphrag_sdk.storage.vector_store import VectorStore __all__ = [ + "DEFAULT_RELATION_TYPES", # Version "__version__", # API @@ -191,8 +195,10 @@ "CachedChunkExtraction", "GraphExtraction", "EntityExtractor", + "CompositeExtractor", "GLiNERExtractor", "LLMExtractor", + "SpacyExtractor", "CorefResolver", "FastCorefResolver", "IngestionPipeline", diff --git a/graphrag_sdk/src/graphrag_sdk/core/models.py b/graphrag_sdk/src/graphrag_sdk/core/models.py index 382fd760..42935403 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/models.py +++ b/graphrag_sdk/src/graphrag_sdk/core/models.py @@ -32,6 +32,14 @@ class Config: # ── Graph Data Types ───────────────────────────────────────────── +#: Node labels the graph store reserves for corpus bookkeeping. They must not +#: be used as entity types: an entity carrying one of these labels corrupts +#: document-level queries and is never marked ``__Entity__``, so it silently +#: disappears from deduplication and retrieval. Single source of truth for +#: ``GraphStore._STRUCTURAL_LABELS`` and the extractor's config-time check. +RESERVED_NODE_LABELS: frozenset[str] = frozenset({"Chunk", "Document"}) + + class GraphNode(DataModel): """A node in the knowledge graph.""" @@ -555,13 +563,39 @@ def _merge_props(existing: list[Attribute], incoming: list[Attribute]) -> list[A class GraphData(DataModel): - """Entities and relationships extracted from text.""" + """Entities and relationships extracted from text. + + ``chunks_attempted`` / ``failed_chunks`` exist because a per-chunk + extraction failure is otherwise invisible: failures are swallowed and the + chunk contributes nothing, so a document where every call failed returns + exactly what a document containing no entities returns — same type, same + empty lists, no exception. Callers had no way to tell "nothing to find" + from "found nothing because everything broke", and a *partial* failure + silently shipped a half-empty graph that looked successful. + + Read them together: ``0`` entities from ``0`` attempted chunks is an empty + document; ``0`` from ``14`` is a broken one. ``failed_chunks`` carries the + chunk uids that raised so the caller can retry just those (see + ``BackfillExecutor``, which follows the same convention) rather than + re-ingesting the document. + """ nodes: list[GraphNode] = Field(default_factory=list) relationships: list[GraphRelationship] = Field(default_factory=list) mentions: list[EntityMention] = Field(default_factory=list) extracted_entities: list[ExtractedEntity] = Field(default_factory=list) extracted_relations: list[ExtractedRelation] = Field(default_factory=list) + chunks_attempted: int = 0 + failed_chunks: list[str] = Field(default_factory=list) + + @property + def extraction_failed(self) -> bool: + """True when every attempted chunk failed. + + Distinguishes total extraction failure from a genuinely empty + document, which reports ``chunks_attempted == 0``. + """ + return self.chunks_attempted > 0 and len(self.failed_chunks) == self.chunks_attempted class ExtractedEntity(DataModel): diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py index a60956c8..ae51121c 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py @@ -9,21 +9,27 @@ FastCorefResolver, ) from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( + CompositeExtractor, EntityExtractor, GLiNERExtractor, LLMExtractor, + SpacyExtractor, ) from graphrag_sdk.ingestion.extraction_strategies.graph_extraction import ( + DEFAULT_RELATION_TYPES, GraphExtraction, ) __all__ = [ + "DEFAULT_RELATION_TYPES", "CachedChunkExtraction", "ExtractionStrategy", "GraphExtraction", "EntityExtractor", + "CompositeExtractor", "GLiNERExtractor", "LLMExtractor", + "SpacyExtractor", "CorefResolver", "FastCorefResolver", ] diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py index b65ab65b..3d4132ca 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py @@ -12,7 +12,8 @@ import re import threading from abc import ABC, abstractmethod -from typing import Any +from collections.abc import Iterable, Sequence +from typing import Any, ClassVar from graphrag_sdk.core.models import ExtractedEntity from graphrag_sdk.core.providers import LLMInterface @@ -67,47 +68,93 @@ "its", } -_ENTITY_STOPLIST: set[str] = _PRONOUNS | { - # Generic/anonymous references - "narrator", - "the narrator", - "author", - "the author", - "reader", - "the reader", - "speaker", - "the speaker", - "listener", - "the listener", - "the man", - "the woman", - "the boy", - "the girl", - "the child", - "man", - "woman", - "boy", - "girl", - "child", - "people", - "person", - "someone", - "somebody", - "everyone", - "everybody", - "mistress", - "master", - # Meta-textual - "story", - "chapter", - "passage", - "book", - "text", - "narrative", - "paragraph", - "section", - "document", -} +# Shell and system abbreviations that read as entities to an NER model but name +# nothing in any document domain. Ported here from the extraction prompt, which +# used to ask the LLM to remove them: a fixed list costs nothing, cannot vary +# between runs, and cannot be quietly skipped the way the prompt instruction was +# (see RESULTS.md P2.10). Well-known two-letter acronyms - AI, US, UK, EU, UN, +# Go - are deliberately absent and stay. +_SHELL_TOKENS: frozenset[str] = frozenset( + { + "sh", + "cd", + "ls", + "rm", + "cp", + "mv", + "dt", + "bg", + "fg", + "fn", + "df", + "du", + "ps", + "cat", + "pwd", + "echo", + "mkdir", + "rmdir", + "chmod", + "chown", + "grep", + "awk", + "sed", + "env", + "sudo", + "ssh", + "tmp", + "var", + "usr", + "bin", + "etc", + } +) + +_ENTITY_STOPLIST: set[str] = ( + _PRONOUNS + | _SHELL_TOKENS + | { + # Generic/anonymous references + "narrator", + "the narrator", + "author", + "the author", + "reader", + "the reader", + "speaker", + "the speaker", + "listener", + "the listener", + "the man", + "the woman", + "the boy", + "the girl", + "the child", + "man", + "woman", + "boy", + "girl", + "child", + "people", + "person", + "someone", + "somebody", + "everyone", + "everybody", + "mistress", + "master", + # Meta-textual + "story", + "chapter", + "passage", + "book", + "text", + "narrative", + "paragraph", + "section", + "document", + } +) # ── Entity Utility Functions ───────────────────────────────────── @@ -127,6 +174,46 @@ def _normalize_type_label(raw: str) -> str: return re.sub(r"[\s_\-/]+", "", s) +# A date that pins down a single moment is an *attribute* of an event ("the +# lamp was installed in 1823"), not something you can hold a conversation +# about. A date that names a *period* ("the 1820s", "the Abbasid era") is a +# thing facts get attached to, so it stays. +# +# Measured on the 11-document benchmark: specific dates were 63 of 274 false +# positive entities (23%). Removing them lifted entity precision 0.577 -> 0.642 +# with recall unchanged at 0.644 (F1 0.609 -> 0.643) and cost nothing, because +# every gold Date entity is a decade. This also stops the graph accumulating +# "X happened_in 1957" edges that carry no answerable content. +_SPECIFIC_DATE_RE = re.compile( + r"""^(?: + (?:c\.?\s*|circa\s+|ca\.?\s*)?\d{3,4}\s*(?:ce|bce|ad|bc)? # 1823, 1003 ce, c. 1200 + | \d{1,2}\s+[a-z]+\s+\d{4} # 14 january 1904 + | [a-z]+\s+\d{1,2},?\s+\d{4} # january 14, 1904 + | \d{4}[-/]\d{1,2}(?:[-/]\d{1,2})? # 1904-01-14 + )$""", + re.IGNORECASE | re.VERBOSE, +) + +# Overrides the rule above: these name a span of time, not a moment. +_DATE_PERIOD_RE = re.compile(r"\d{3,4}s\b|centur|era\b|dynasty|period|decade|age\b", re.IGNORECASE) + + +def is_specific_date(name: str) -> bool: + """True if name pins down one moment in time rather than naming a period. + + Specific dates are rejected as entity names: they belong on the relation + that mentions them, not as nodes of their own. Note the deliberate gap -- + a product genuinely named for a number ("747", "1984") is indistinguishable + from a year here and will be dropped. That trade was worth 63 false + positives against zero true positives on the benchmark corpus, but it is + the first thing to revisit if a domain uses numeric product names. + """ + stripped = name.strip() + if _DATE_PERIOD_RE.search(stripped): + return False + return bool(_SPECIFIC_DATE_RE.match(stripped)) + + def is_valid_entity_name(name: str) -> bool: """Return True if name passes quality gates for entity extraction.""" if not name or not name.strip(): @@ -134,7 +221,16 @@ def is_valid_entity_name(name: str) -> bool: stripped = name.strip() if len(stripped) < MIN_NAME_LEN or len(stripped) > MAX_NAME_LEN: return False - if stripped.lower() in _ENTITY_STOPLIST: + # "US" is a country, "us" is a pronoun, and casefolding the stoplist check + # conflates them. An all-caps short token is an acronym, not a pronoun. + is_acronym = len(stripped) <= 3 and stripped.isupper() and stripped.isalpha() + if not is_acronym and stripped.lower() in _ENTITY_STOPLIST: + return False + if is_specific_date(stripped): + return False + # Operator and punctuation tokens (+=, ->, ==, !=). A name with no letter or + # digit anywhere in it cannot be the name of anything. + if not any(ch.isalnum() for ch in stripped): return False return True @@ -317,41 +413,278 @@ class GLiNERExtractor(EntityExtractor): so a single instance can be safely shared across concurrent ``asyncio.to_thread`` calls (e.g. parallel doc ingestion). + GLiNER has a hard input limit (``config.max_len``) and truncates anything + beyond it **silently** — no exception, no warning, the tail simply never + reaches the model. Measured on ``gliner_medium-v2.1`` (``max_len`` 384): a + probe entity placed at word-token 388 is always returned, at 389 never. + The default model's limit is 2048, but real documents still exceed it — + disabling windowing on our benchmark corpus dropped recall from 0.805 to + 0.621 — so windowing matters regardless of the model. + + To keep long chunks fully visible, text longer than ``window_tokens`` is + processed as a series of overlapping windows and the results are merged. + Short text takes a fast path and behaves exactly as before. + + **Confidence handling.** By default the model is queried at ``threshold``, + so anything less confident is discarded inside GLiNER and never reaches this + SDK. Set ``candidate_threshold`` below ``threshold`` to instead *keep* those + entities and label them ``"Unknown"``. The rest of the pipeline is already + built for this: ontology filtering explicitly whitelists ``"Unknown"`` so + low-confidence nodes survive pruning, and entity resolution prefers any + specific type over ``"Unknown"`` when merging duplicates. Off by default — + lowering ``threshold`` outright was measured to raise entity recall 32% + while dropping entity F1 0.568 -> 0.474 and triple F1 0.236 -> 0.211, so + low-confidence predictions are mostly noise and should be marked, not + trusted. + + **Thresholds are model-specific and are not comparable between models.** + ``DEFAULT_THRESHOLDS`` records the measured operating point for each known + model, and ``threshold=None`` (the default) looks it up. Passing an explicit + number that was tuned for a different model is the single easiest way to + break this class: the bi-encoder models return almost nothing at the 0.75 + that suits ``gliner_medium-v2.1`` — measured at **2 entities for an entire + corpus**, with no error raised. + Args: threshold: Confidence threshold (0-1). Below this → "Unknown". + ``None`` (default) selects the value measured for ``model_name``. model_name: HuggingFace model name for GLiNER. + window_tokens: Word-tokens per inference window. ``None`` derives it + from the model's own ``config.max_len`` minus a safety margin. + window_overlap: Word-tokens shared between consecutive windows. Must + exceed the model's ``max_width`` (longest representable entity, + 12 words by default) or entities on a boundary are lost. """ + # Safety margin under config.max_len; the label prompt and special tokens + # share the window with the text. + _WINDOW_MARGIN = 34 + + #: Default model. Measured against ``gliner_medium-v2.1`` on an 11-document + #: benchmark: ceiling recall 0.805 vs 0.709, 432MB vs 781MB on disk, 1004MB + #: vs 1532MB resident, ~14s vs ~29s, and a 2048-word window vs 384. + DEFAULT_MODEL = "knowledgator/gliner-bi-small-v2.0" + + #: Confidence thresholds are on different scales per model and MUST be + #: re-tuned when the model changes. Each value below is the measured best + #: end-to-end operating point, not a guess. + DEFAULT_THRESHOLDS: dict[str, float] = { + "knowledgator/gliner-bi-small-v2.0": 0.5, + "knowledgator/gliner-bi-base-v2.0": 0.5, + "urchade/gliner_medium-v2.1": 0.75, + "urchade/gliner_large-v2.1": 0.75, + "gliner-community/gliner_medium-v2.5": 0.75, + } + + #: Used when ``model_name`` is not in ``DEFAULT_THRESHOLDS``. The GLiNER + #: library's own default, which is a safer guess than assuming our tuned + #: value transfers to an unknown model. + _FALLBACK_THRESHOLD = 0.5 + def __init__( self, - threshold: float = 0.75, - model_name: str = "urchade/gliner_medium-v2.1", + threshold: float | None = None, + model_name: str | None = None, + window_tokens: int | None = None, + window_overlap: int = 48, + candidate_threshold: float | None = None, ) -> None: + self._model_name = model_name or self.DEFAULT_MODEL + if threshold is None: + threshold = self.DEFAULT_THRESHOLDS.get(self._model_name, self._FALLBACK_THRESHOLD) + if self._model_name not in self.DEFAULT_THRESHOLDS: + logger.warning( + "No measured threshold for GLiNER model %r; falling back to " + "%.2f. Thresholds are not comparable between models, so " + "tune this before relying on the results.", + self._model_name, + self._FALLBACK_THRESHOLD, + ) self._threshold = threshold - self._model_name = model_name + if candidate_threshold is not None and candidate_threshold > threshold: + raise ValueError( + f"candidate_threshold ({candidate_threshold}) must be <= " + f"threshold ({threshold}); a candidate floor above the demotion " + f"line would discard the very entities it is meant to keep" + ) + self._candidate_threshold = candidate_threshold self._model: Any = None + # Retained only so callers/tests that swap in a context manager to + # A/B the removed inference lock keep working. Nothing in this class + # acquires it any more; model loading uses the class-level + # ``_CACHE_LOCK`` and inference deliberately takes no lock at all. + # See ``_predict_sync`` for the thread-safety evidence. self._lock = threading.Lock() + self._window_tokens = window_tokens + self._window_overlap = window_overlap + self._splitter: Any = None def _load_model(self) -> Any: if self._model is None: - with self._lock: - # Double-check after acquiring lock - if self._model is None: - try: - from gliner import GLiNER - except ImportError: - raise ImportError( - "GLiNER is required for GLiNERExtractor. " - "Install with: pip install gliner" - ) - self._model = GLiNER.from_pretrained(self._model_name) + self._model = self._get_shared_model(self._model_name) return self._model + # Process-wide cache of loaded GLiNER models, keyed on model name. + # + # Bug #11: nothing cached the model, so every ``GLiNERExtractor`` instance + # loaded its own copy. Measured with current (not peak) RSS, creating six + # extractors and forcing each to load: + # + # baseline 74.5 MB -> 2447 MB after 6 copies = ~395 MB per copy + # + # which projects to ~11.6 GB for 30 concurrent documents. That matches the + # customer report of large RSS under concurrent ingest. With this cache the + # second and later extractors for the same model cost ~0. + # + # Keyed on model name rather than shared unconditionally because two + # extractors may legitimately want different models; those must stay + # separate or one would silently answer with the other's weights. + _MODEL_CACHE: ClassVar[dict[str, Any]] = {} + _CACHE_LOCK: ClassVar[threading.Lock] = threading.Lock() + + @classmethod + def _get_shared_model(cls, model_name: str) -> Any: + cached = cls._MODEL_CACHE.get(model_name) + if cached is not None: + return cached + with cls._CACHE_LOCK: + # Double-checked: another thread may have loaded it while we waited. + cached = cls._MODEL_CACHE.get(model_name) + if cached is None: + try: + from gliner import GLiNER + except ImportError: + raise ImportError( + "GLiNER is required for GLiNERExtractor. Install with: pip install gliner" + ) + cached = GLiNER.from_pretrained(model_name) + cls._MODEL_CACHE[model_name] = cached + return cached + + def _resolve_window(self, model: Any) -> int: + """Window size in word-tokens, derived from the model if not set.""" + if self._window_tokens is not None: + return self._window_tokens + max_len = getattr(getattr(model, "config", None), "max_len", None) + if not isinstance(max_len, int) or max_len <= 0: + max_len = 384 + return max(64, max_len - self._WINDOW_MARGIN) + + def _word_spans(self, model: Any, text: str) -> list[tuple[str, int, int]]: + """Split text the same way GLiNER does, keeping char offsets.""" + if self._splitter is None: + splitter = getattr(getattr(model, "data_processor", None), "words_splitter", None) + if splitter is None: + from gliner.data_processing import WordsSplitter + + splitter = WordsSplitter() + self._splitter = splitter + return list(self._splitter(text)) + + @staticmethod + def _merge(preds: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Deduplicate predictions from overlapping windows. + + Identical spans found in two windows collapse to the highest-scoring + copy. A span that is strictly contained in a longer span of the same + label is dropped: it is the truncated remains of an entity clipped by a + window edge, which the neighbouring window saw whole. + """ + best: dict[tuple[int, int, str], dict[str, Any]] = {} + for p in preds: + key = (p["start"], p["end"], p["label"]) + prev = best.get(key) + if prev is None or p.get("score", 0.0) > prev.get("score", 0.0): + best[key] = p + + kept: list[dict[str, Any]] = [] + for p in best.values(): + contained = any( + q is not p + and q["label"] == p["label"] + and q["start"] <= p["start"] + and p["end"] <= q["end"] + and (q["end"] - q["start"]) > (p["end"] - p["start"]) + for q in best.values() + ) + if not contained: + kept.append(p) + kept.sort(key=lambda p: (p["start"], p["end"])) + return kept + def _predict_sync(self, text: str, entity_types: list[str]) -> list[dict[str, Any]]: model = self._load_model() labels = [t.lower() for t in entity_types] - with self._lock: - return model.predict_entities(text, labels, threshold=self._threshold) + window = self._resolve_window(model) + # What we ask the MODEL for. Anything between this floor and + # ``self._threshold`` comes back and is demoted to ``UNKNOWN_LABEL`` by + # ``_parse_predictions`` rather than being discarded. When no candidate + # threshold is configured the two are equal and nothing is demoted. + floor = self._threshold if self._candidate_threshold is None else self._candidate_threshold + + # No lock here, deliberately. + # + # Bug #8: this whole block used to run under ``self._lock`` while the + # caller dispatched it through ``asyncio.to_thread`` — the SDK paid for + # threads and then serialised them anyway. Concurrent documents queued + # behind each other in NER. + # + # Removing it is only sound if GLiNER inference is genuinely + # thread-safe, so that was tested rather than assumed: eight documents + # extracted concurrently, unlocked, compared field-by-field against the + # serialised result, five trials — 40/40 documents byte-identical. + # Inference mutates no model state; only ``_load_model`` does, and that + # is guarded separately by ``_CACHE_LOCK``. + # + # Measured on eight documents, counterbalanced (locked, unlocked, + # unlocked, locked) so ordering and warm caches cannot explain it: + # locked 3.75 s / 3.54 s versus unlocked 2.21 s / 2.46 s = **1.56x**. + # Note issue #71 claimed 3.44x; the honest measured figure is 1.56x, + # because torch's own intra-op threading already uses the cores. + return self._predict_body(model, text, labels, window, floor) + + def _predict_body( + self, + model: Any, + text: str, + labels: list[str], + window: int, + floor: float, + ) -> list[dict[str, Any]]: + """Windowed inference. Split out from :meth:`_predict_sync` so the + lock removal above could be A/B tested by wrapping one call site.""" + words = self._word_spans(model, text) + + # Fast path: fits in one window, identical to unwindowed behaviour. + if len(words) <= window: + return model.predict_entities(text, labels, threshold=floor) + + step = max(1, window - self._window_overlap) + out: list[dict[str, Any]] = [] + for begin in range(0, len(words), step): + span = words[begin : begin + window] + if not span: + break + lo, hi = span[0][1], span[-1][2] + for p in model.predict_entities(text[lo:hi], labels, threshold=floor): + p = dict(p) + p["start"] += lo + p["end"] += lo + p["text"] = text[p["start"] : p["end"]] + out.append(p) + if begin + window >= len(words): + break + + merged = self._merge(out) + logger.debug( + "GLiNER windowed inference: %d word-tokens -> %d windows, " + "%d raw predictions -> %d after merge", + len(words), + (len(words) - 1) // step + 1, + len(out), + len(merged), + ) + return merged async def extract_entities( self, @@ -448,3 +781,299 @@ def _parse_response( ) ) return entities + + +# ── spaCy Extractor ────────────────────────────────────────────── + + +class SpacyExtractor(EntityExtractor): + """Classic (non zero-shot) NER via spaCy, for well-known proper nouns. + + This exists to cover a measured blind spot rather than as a general + extractor. Swapping the GLiNER default to ``gliner-bi-small-v2.0`` gained 49 + gold entities but *lost* 24, and the losses were textbook proper nouns — + ``Baghdad``, ``Cairo``, ``Madrid``, ``Barcelona``, ``Constantinople``, + ``Paris Observatory`` — which a supervised model trained on exactly those + categories gets right trivially. + + Used alone it is a poor fit for GraphRAG: its label set is fixed, so it + cannot represent custom types such as ``Method`` or ``Technology``. Its + value is as the second half of a :class:`CompositeExtractor`. + + ``DEFAULT_LABELS`` is deliberately narrow. Measured on an 11-document + benchmark (157 chunks), unioned with the default GLiNER extractor: + + =========================== ====== ========= ===== + spaCy labels added recall precision F1 + =========================== ====== ========= ===== + none (GLiNER alone) 0.494 0.554 0.522 + PERSON/ORG/GPE/FAC 0.613 0.499 0.550 + + LOC 0.616 0.497 0.550 + + NORP 0.616 0.477 0.538 + all 12 usable labels 0.654 0.352 0.458 + =========================== ====== ========= ===== + + Widening past the four default labels buys recall by giving away + precision, and F1 falls off a cliff at the wide setting. The narrow set + recovers 18 of the 24 lost entities for ~5s per 157 chunks. + + Requires the ``spacy`` extra and a downloaded model:: + + pip install "graphrag-sdk[spacy]" + python -m spacy download en_core_web_lg + """ + + #: spaCy labels kept by default. Everything else is dropped rather than + #: guessed at, because a wrong type is worse than a missing entity here. + DEFAULT_LABELS = frozenset({"PERSON", "ORG", "GPE", "FAC"}) + + #: spaCy label -> the generic type name we look for in ``entity_types``. + #: Candidates are tried in order and the first one the caller allows wins, + #: so this works whether a schema calls it ``Place``, ``Location`` or + #: ``Organization``. + LABEL_MAP: dict[str, tuple[str, ...]] = { + "PERSON": ("Person",), + "ORG": ("Organization", "Institution", "Company"), + "GPE": ("Location", "Place", "GeographicLocation"), + "LOC": ("Location", "Place", "GeographicLocation"), + "FAC": ("Building", "Facility", "Location", "Place"), + "NORP": ("Group", "Nationality", "Organization"), + "EVENT": ("Event",), + "PRODUCT": ("Product",), + "WORK_OF_ART": ("Work", "Publication", "Product"), + "LAW": ("Law",), + "DATE": ("Date",), + "LANGUAGE": ("Language",), + } + + DEFAULT_MODEL = "en_core_web_lg" + + def __init__( + self, + model_name: str | None = None, + labels: Iterable[str] | None = None, + confidence: float = 0.5, + ) -> None: + """Initialise the extractor. + + Args: + model_name: spaCy pipeline to load. Defaults to ``en_core_web_lg``. + ``en_core_web_sm`` is ~2.6x smaller and just as fast but + measured 0.554 ceiling recall against 0.601 for ``lg``. + labels: spaCy entity labels to keep. Defaults to + :attr:`DEFAULT_LABELS`. Widening this reduces F1 — see the + class docstring for the measured trade. + confidence: Score recorded on emitted entities. spaCy's ``ents`` + expose no per-span probability, so this is a fixed stand-in + rather than a real confidence, and is set at the default + GLiNER threshold so these entities are neither favoured nor + penalised relative to GLiNER's. + """ + self._model_name = model_name or self.DEFAULT_MODEL + self._labels = frozenset(labels) if labels is not None else self.DEFAULT_LABELS + self._confidence = confidence + self._nlp: Any = None + self._lock = asyncio.Lock() + + async def _get_nlp(self) -> Any: + if self._nlp is not None: + return self._nlp + async with self._lock: + if self._nlp is None: + self._nlp = await asyncio.to_thread(self._load) + return self._nlp + + def _load(self) -> Any: + try: + import spacy + except ImportError as exc: # pragma: no cover - depends on env + raise ImportError( + "SpacyExtractor requires the 'spacy' extra. Install with: " + 'pip install "graphrag-sdk[spacy]"' + ) from exc + try: + # The parser and lemmatizer cost time and contribute nothing to NER. + return spacy.load(self._model_name, disable=["lemmatizer", "textcat"]) + except OSError as exc: # pragma: no cover - depends on env + raise OSError( + f"spaCy model '{self._model_name}' is not installed. Run: " + f"python -m spacy download {self._model_name}" + ) from exc + + def _type_for(self, label: str, entity_types: list[str]) -> str | None: + """Map a spaCy label onto an allowed type, or None if unrepresentable.""" + for candidate in self.LABEL_MAP.get(label, ()): + mapped = label_for_type(candidate, entity_types) + if mapped != UNKNOWN_LABEL: + return mapped + return None + + async def extract_entities( + self, + text: str, + entity_types: list[str], + source_chunk_id: str, + ) -> list[ExtractedEntity]: + nlp = await self._get_nlp() + doc = await asyncio.to_thread(nlp, text) + + preds: list[dict[str, Any]] = [] + for ent in doc.ents: + if ent.label_ not in self._labels: + continue + etype = self._type_for(ent.label_, entity_types) + if etype is None: + continue # caller's schema has no home for this label + preds.append( + { + "text": ent.text, + "label": etype, + "score": self._confidence, + "start": ent.start_char, + "end": ent.end_char, + } + ) + return _parse_predictions(preds, entity_types, source_chunk_id, self._confidence) + + +# ── Composite Extractor ────────────────────────────────────────── + + +class CompositeExtractor(EntityExtractor): + """Run several extractors over the same text and merge their entities. + + Built for the measured fact that no single extractor we tested wins + everywhere: ``gliner-bi-small-v2.0`` is far stronger on multi-word, + domain-specific entities (``Fresnel lens``, ``differential gear train + mechanism``, ``Samarkand Expedition of 892``), while a supervised spaCy + pipeline is stronger on plain proper nouns (``Baghdad``, ``Madrid``, + ``Paris Observatory``). Combining them recovered 18 of the 24 entities lost + in the model swap. + + Measured on an 11-document benchmark, 157 chunks, GLiNER default plus + ``SpacyExtractor`` at its default labels: recall 0.494 -> 0.613, F1 + 0.522 -> 0.550, for about 5 extra seconds. Precision falls 0.554 -> 0.499, + so this trades some precision for a larger gain in recall. + + Extractors run concurrently. Duplicates are resolved by normalised name: + the **earliest** extractor in the list wins the type, which makes ordering + meaningful — put your most trusted or most schema-aware extractor first. + Chunk provenance is unioned so a merged entity keeps every chunk it came + from. + + A failing extractor does not take the others down: the exception is logged + and its results are skipped, on the grounds that degraded extraction beats + a failed ingest. If *every* extractor fails the error is re-raised. + + Example:: + + extractor = CompositeExtractor([ + GLiNERExtractor(), + SpacyExtractor(), + ]) + """ + + def __init__( + self, + extractors: Sequence[EntityExtractor], + suppress_overlaps: bool = True, + ) -> None: + """Initialise the extractor. + + Args: + extractors: Extractors to run, in priority order. The first one to + produce a given name decides its type. + suppress_overlaps: Drop an entity from a later extractor when its + character span overlaps one already claimed by an earlier + extractor. Without this, merging two NER systems reliably + produces near-duplicate fragments — measured on one sentence, + spaCy contributed ``Fresnel`` (typed ``Organization``) next to + GLiNER's ``Fresnel lens``, and ``The Paris Observatory`` next + to ``Paris Observatory``. Those fragments are false positives + and inflate the entity count without adding knowledge. + Entities without span information are always kept, since there + is no evidence on which to drop them. + + Raises: + ValueError: If ``extractors`` is empty. + """ + if not extractors: + raise ValueError("CompositeExtractor requires at least one extractor") + self._extractors = list(extractors) + self._suppress_overlaps = suppress_overlaps + + @staticmethod + def _spans_of(ent: ExtractedEntity) -> list[tuple[int, int]]: + """Character spans claimed by an entity, flattened across chunks. + + ``_parse_predictions`` passes ``spans`` as an extra model field rather + than into ``attributes``, so check both. + """ + spans = getattr(ent, "spans", None) + if spans is None: + spans = ent.attributes.get("spans") + out: list[tuple[int, int]] = [] + if isinstance(spans, dict): + for items in spans.values(): + for sp in items or (): + try: + out.append((int(sp["start"]), int(sp["end"]))) + except (KeyError, TypeError, ValueError): + continue + return out + + async def extract_entities( + self, + text: str, + entity_types: list[str], + source_chunk_id: str, + ) -> list[ExtractedEntity]: + results = await asyncio.gather( + *(e.extract_entities(text, entity_types, source_chunk_id) for e in self._extractors), + return_exceptions=True, + ) + + merged: dict[str, ExtractedEntity] = {} + claimed: list[tuple[int, int]] = [] + failures: list[BaseException] = [] + for extractor, result in zip(self._extractors, results, strict=True): + if isinstance(result, BaseException): + failures.append(result) + logger.warning( + "%s failed during entity extraction, skipping its results: %s", + type(extractor).__name__, + result, + ) + continue + fresh: list[tuple[int, int]] = [] + for ent in result: + key = ent.name.strip().casefold() + if not key: + continue + existing = merged.get(key) + if existing is None: + spans = self._spans_of(ent) + if ( + self._suppress_overlaps + and spans + and any(s < ce and cs < e for (s, e) in spans for (cs, ce) in claimed) + ): + continue # fragment of an entity a better extractor already has + merged[key] = ent + fresh.extend(spans) + continue + # Keep the earlier extractor's type; only fill genuine gaps. + if existing.type == UNKNOWN_LABEL and ent.type != UNKNOWN_LABEL: + existing.type = ent.type + if not existing.description and ent.description: + existing.description = ent.description + for cid in ent.source_chunk_ids: + if cid not in existing.source_chunk_ids: + existing.source_chunk_ids.append(cid) + # Only claim spans once the whole extractor is processed, so two + # entities from the SAME extractor never suppress each other. + claimed.extend(fresh) + + if failures and len(failures) == len(self._extractors): + raise failures[0] + return list(merged.values()) diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py index 28b8b55e..a4b442aa 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py @@ -13,6 +13,7 @@ from graphrag_sdk.core.context import Context from graphrag_sdk.core.models import ( _SDK_MANAGED_ATTRIBUTE_NAMES, + RESERVED_NODE_LABELS, Attribute, EntityMention, ExtractedEntity, @@ -30,6 +31,7 @@ from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( DEFAULT_ENTITY_TYPES, NER_PROMPT, + UNKNOWN_LABEL, EntityExtractor, GLiNERExtractor, LLMExtractor, @@ -42,6 +44,76 @@ logger = logging.getLogger(__name__) +# Default relation vocabulary — the missing counterpart to DEFAULT_ENTITY_TYPES. +# +# Entity extraction has always shipped a default type list, so the LLM is told +# what kinds of things to look for. Relations shipped nothing: with no ontology +# the prompt said "use a descriptive label in UPPER_SNAKE_CASE", and the model +# invented a fresh name for almost every edge. Measured on an 11-document +# corpus: 447 distinct relation labels against 30 in the gold annotation, 68.2% +# of them used exactly once, and only 17.3% of edges carrying a label the gold +# data also uses. Two of gold's most common predicates (``contains``, 123 +# triples; ``authored``, 54) were emitted zero times. +# +# Supplying any fixed list doubles exact triple F1 (0.0725 -> 0.1490, +2.06x), +# and a list written *without* reference to the gold vocabulary scored as well +# as the gold vocabulary itself (0.1490 vs 0.1456). The gain comes from being +# consistent, not from guessing the right words — which is what makes a shipped +# default worth having. +# +# This list is deliberately domain-neutral and pairs with DEFAULT_ENTITY_TYPES. +# It is *guidance*, not a filter: exactly like ``entity_types``, it steers the +# prompt, and a relation the model labels outside this list is still kept. The +# hard-filtering path is ``Ontology.relations``, which prunes non-conforming +# edges in ``IngestionPipeline._prune``. Users who want that stricter behaviour +# should declare an ontology; users who want none of it can pass +# ``relation_types=[]``. +DEFAULT_RELATION_TYPES: list[str] = [ + # structure / place + "located_in", + "part_of", + "contains", + "occurred_in", + # affiliation + "member_of", + "employed_at", + "founded", + "owns", + "subsidiary_of", + # creation & production + "created", + "authored", + "designed_by", + "developed_by", + "manufactured", + "published_in", + "supplied_to", + # people + "born_in", + "died_in", + "married_to", + "child_of", + "sibling_of", + "student_of", + "colleague_of", + # activity & influence + "participated_in", + "directed", + "awarded", + "named_after", + "succeeded_by", + "influenced", + # technical + "uses", + "based_on", +] + +# This prompt asks for entity verification AND descriptions AND relations. +# Removing the verification job was tried and REVERTED: without it the LLM +# emitted 813 entities instead of 719 and entity precision fell 0.645 -> 0.551 +# (RESULTS.md P2.12). The instruction is doing real work, even though a +# dedicated call built to do the same job scored no better than random +# (P2.10). Do not remove it again without re-running that measurement. VERIFY_EXTRACT_RELS_PROMPT = ( "You are an expert knowledge graph builder.\n" "Given the text and pre-extracted entities below, do two things:\n" @@ -58,18 +130,28 @@ "{text}\n\n" "## Instructions\n\n" "### Entities\n" - "- REMOVE any entity that is:\n" - " - A purely symbolic or operator token (e.g. +=, ->, ++, ==, !=)\n" - " - A common non-domain-specific shell/system abbreviation " - "(e.g. sh, cd, dt, ls, rm, cp, mv)\n" - " - A generic short token (1-2 characters) that is not a widely-recognised " - "named entity or acronym (AI, US, UK, Go are fine; dt, bg, fn are not)\n" + "- REMOVE any entity that is not a real named thing in the text: an " + "operator or symbol token, a generic shell or system abbreviation, or a " + "short token that is not a widely-recognised name or acronym.\n" "- For each verified entity provide a concise 1-2 sentence description " "capturing key attributes and roles from the text. This description is " "embedded for semantic search.\n\n" "### Relationships\n" "- Extract ALL factual connections stated or implied in the text.\n" - "- source and target must be entity names from the verified entity list.\n" + "- source and target must be entity names from the entity list above.\n" + # Measured (P5.6): "ALL" above is not enough on its own -- the model treats + # the task as a summary, returns ~12 relations per chunk and stops while + # using 2.5k of a 16k reply budget. Giving it twice the text grew the reply + # 1.3%; the three lines below grew it 15.5% for +4% ingest time, and beat a + # second "what did you miss?" LLM call that cost 4.7x the ingest time. + # Do not extend this with a per-entity walkthrough instruction: measured at + # 4.6x ingest time and a worse graph than changing nothing. + "- This is an EXHAUSTIVE extraction task, NOT a summary. Do not stop after " + "the most important or most obvious connections.\n" + "- There is no maximum. A dense paragraph often yields 20 or more " + "relationships. A long list is correct, not a mistake.\n" + "- Never end the list early. Never leave connections out because you have " + "already written several.\n" "{relationship_type_instruction}" "- description: one sentence describing the relationship as a " "standalone fact. This is embedded for semantic search — it must be " @@ -82,6 +164,57 @@ "Return ONLY valid JSON, nothing else." ) +# ── Entity verification prompt ───────────────────────────────────── +# +# Split out of VERIFY_EXTRACT_RELS_PROMPT after measuring that a single call +# asked to verify entities, describe them AND extract relations does the last +# two and effectively skips the first. Symptoms traced to that one cause: +# ~29% of written relations unsupported by the source text; raising NER recall +# 32% dropped entity precision by almost exactly the same amount (the verifier +# passed the junk straight through); and tagging low-confidence entities +# "Unknown" changed nothing downstream. +# +# Two design choices matter here: +# 1. The reply is a fixed-length verdict list, one row per input entity, not +# a rewritten entity list. A rewrite can be satisfied by echoing the +# input; a verdict per row cannot. +# 2. Every "keep" must carry the exact quote containing the entity. The +# caller checks that quote against the real text without asking the model +# again, so a fabricated justification is detectable rather than trusted. +VERIFY_ENTITIES_PROMPT = ( + "You are a strict fact-checker for a knowledge graph.\n" + "Below is a numbered list of candidate entities. Each was proposed by an " + "automatic extractor and may be wrong. Judge each one INDEPENDENTLY.\n\n" + "## Allowed Entity Types\n" + "{entity_types}\n\n" + "## Candidates\n" + "{candidates}\n\n" + "## Rules\n" + 'Mark an entity "drop" if ANY of these is true:\n' + "- It is not a real named thing (e.g. a bare year or decade like " + '"1010s" or "1003 CE", a stray number, a date fragment)\n' + "- It is a symbol or operator (+=, ->, ==)\n" + "- It is a generic 1-2 character token that is not a well-known acronym " + "(AI, US, UK are fine; dt, bg, fn are not)\n" + "- It is a sentence fragment or description rather than a name\n" + "- It is only a PART of a longer entity in the same list " + '(e.g. "Fresnel" when "Fresnel lens" is also listed)\n' + "- It does not appear in the context provided for it\n\n" + 'Otherwise mark it "keep".\n\n' + 'For every "keep":\n' + '- Set "type" to the best-fitting type from the allowed list above. ' + "Correct the proposed type if it is wrong.\n" + '- Set "quote" to text copied EXACTLY, character for character, from ' + "that entity's context, containing the entity name. Do not paraphrase, " + "reword or shorten it. This is checked automatically.\n\n" + "Return ONLY a JSON array with exactly {n} objects, one per candidate, in " + "the same order:\n" + '[{{"id": 1, "verdict": "keep", "type": "Person", "quote": "..."}}, ' + '{{"id": 2, "verdict": "drop"}}]\n' + "Return ONLY valid JSON, nothing else." +) + + _DEFAULT_JSON_EXAMPLE = ( '{{"entities": [{{"name": "...", "type": "...", "description": "..."}}], ' '"relationships": [{{"source": "...", "target": "...", "type": "...", ' @@ -336,6 +469,38 @@ def _optional_extras(obj: Any) -> dict[str, Any]: return extra +def _reject_reserved_labels(types: list[str]) -> list[str]: + """Reject entity types that collide with the store's structural labels. + + ``Document`` and ``Chunk`` are used by the graph store for corpus + bookkeeping. An extracted entity carrying one of those labels fails two + ways at once, both silently: + + 1. Document-level queries (``MATCH (p:Document) ...``) start returning + extracted entities, so document counts and lookups are wrong. This was + found in the field as an 11-document corpus reporting 106 documents. + 2. ``GraphStore._write_nodes`` marks a node as an entity only when its + label is *not* structural, so the node never receives ``__Entity__`` + and is dropped from deduplication and retrieval. It is written, then + ignored. + + Neither failure raises, so the graph just quietly degrades. Fail here + instead, at configuration time, where the caller can act on it. + """ + reserved = {label.casefold(): label for label in RESERVED_NODE_LABELS} + clashes = [t for t in types if str(t).strip().casefold() in reserved] + if clashes: + raise ValueError( + f"entity_types may not contain the reserved label(s) {sorted(set(clashes))}. " + f"{sorted(RESERVED_NODE_LABELS)} are used internally for corpus " + "bookkeeping; reusing them silently corrupts document counts and " + "removes the entity from deduplication and retrieval. " + "Rename the type (e.g. 'Document' -> 'Publication', " + "'Chunk' -> 'TextSegment')." + ) + return list(types) + + def _format_entity_types(types: list[str], descs: dict[str, str] | None = None) -> str: """Format entity types for prompt injection. @@ -416,6 +581,13 @@ class GraphExtraction(ExtractionStrategy): coref_resolver: Optional coreference resolver applied per-chunk. entity_types: Entity type labels. Default: DEFAULT_ENTITY_TYPES. Overridden by ontology.entities if present. + relation_types: Relation labels offered to the LLM. Default: + DEFAULT_RELATION_TYPES. This is the exact counterpart of + ``entity_types`` — guidance for the prompt, not a filter, so a + relation labelled outside the list is still kept. Overridden by + ``ontology.relations`` when one is declared, and that path *does* + prune non-conforming edges. Pass ``[]`` to restore the old + open-vocabulary behaviour where the model invents every label. max_concurrency: Maximum parallel LLM calls. """ @@ -426,13 +598,198 @@ def __init__( entity_extractor: EntityExtractor | None = None, coref_resolver: CorefResolver | None = None, entity_types: list[str] | None = None, + relation_types: list[str] | None = None, max_concurrency: int | None = None, + verify_entities: bool = False, + verify_batch_size: int = 100, + verify_context_chars: int = 240, ) -> None: self.llm = llm self.entity_extractor = entity_extractor or GLiNERExtractor() self.coref_resolver = coref_resolver - self.entity_types = entity_types or list(DEFAULT_ENTITY_TYPES) + self.entity_types = _reject_reserved_labels(entity_types or list(DEFAULT_ENTITY_TYPES)) + # `is None` rather than falsy: relation_types=[] is a meaningful request + # for open-vocabulary mode, not an omission. + self.relation_types = ( + list(DEFAULT_RELATION_TYPES) if relation_types is None else list(relation_types) + ) self._max_concurrency = max_concurrency + self.verify_entities = bool(verify_entities) + self.verify_batch_size = max(1, int(verify_batch_size)) + self.verify_context_chars = max(40, int(verify_context_chars)) + + # ── Entity verification (step 1b) ──────────────────────────── + + @staticmethod + def _context_for( + ent: ExtractedEntity, + chunk_uid: str, + text: str, + width: int, + ) -> str: + """A window of source text around the entity, for the judge to read. + + Prefers the recorded character span. Falls back to a literal search, + then to the head of the chunk, so an entity is never sent without + context (which would guarantee a "drop" verdict for the wrong reason). + """ + spans = getattr(ent, "spans", None) or {} + pos = -1 + for sp in spans.get(chunk_uid, ()) or (): + try: + pos = int(sp["start"]) + break + except (KeyError, TypeError, ValueError): + continue + if pos < 0: + pos = text.find(ent.name) + if pos < 0: + return text[:width] + half = width // 2 + return text[max(0, pos - half) : pos + half] + + async def _verify_entities( + self, + chunk_entities: list[list[ExtractedEntity]], + chunk_texts: list[str], + chunk_uids: list[str], + entity_types: list[str], + entity_type_descs: dict[str, str], + ctx: Context, + ) -> list[list[ExtractedEntity]]: + """Drop candidate entities that a dedicated LLM pass rejects. + + .. warning:: + **Measured as ineffective; off by default. Do not enable without + re-measuring.** On the 11-document benchmark this removed 114 of 702 + entities: 57 correct and 41 junk, against 56.6 and 41.4 expected from + removing the same number *at random*. Entity precision moved 0.5772 + to 0.5764. Entity F1 0.609 -> 0.561, lax triple F1 0.248 -> 0.213, + ingest 124.5s -> 224.6s (+80%, because this is a barrier between two + otherwise-overlapping phases, not because of the ~7 added calls). + + The failure is not fact-checking. Junk like ``1823`` survives because + 1823 really is in the text and ``Date`` really is an allowed type -- + the model answers the question we asked correctly. The question we + need answered is *salience* (should this be a node in this graph), + which depends on the schema and the queries rather than the text, and + which we have not defined. Kept because the mechanism is sound and + becomes useful the moment there is a salience criterion to give it. + + Runs between NER and relation extraction. Entities are deduplicated by + name across the whole corpus before judging, so cost scales with the + number of *distinct* names rather than with the number of chunks — on an + 11-document benchmark that is ~700 names, or 7 calls at the default + batch size, against 157 relation calls. + + A verdict is only honoured when the model's supporting quote is actually + present in the source text. Anything else - malformed JSON, a missing + row, a fabricated quote, a failed request - leaves the entity in place. + Verification may only ever *remove* entities it can justify, so a broken + judge degrades to today's behaviour instead of emptying the graph. + """ + # name -> (canonical entity, context) for the first occurrence seen + first: dict[str, tuple[ExtractedEntity, str]] = {} + for ents, text, uid in zip(chunk_entities, chunk_texts, chunk_uids): + for ent in ents: + key = ent.name.strip().casefold() + if key and key not in first: + first[key] = ( + ent, + self._context_for(ent, uid, text, self.verify_context_chars), + ) + if not first: + return chunk_entities + + keys = list(first) + batches = [ + keys[i : i + self.verify_batch_size] + for i in range(0, len(keys), self.verify_batch_size) + ] + prompts = [] + for batch in batches: + lines = [] + for n, key in enumerate(batch, 1): + ent, context = first[key] + lines.append( + f'{n}. name: "{ent.name}" | proposed type: {ent.type}\n context: {context!r}' + ) + prompts.append( + VERIFY_ENTITIES_PROMPT.format( + entity_types=_format_entity_types(entity_types, entity_type_descs), + candidates="\n".join(lines), + n=len(batch), + ) + ) + + batch_kw: dict[str, Any] = {} + if self._max_concurrency is not None: + batch_kw["max_concurrency"] = self._max_concurrency + results = await self.llm.abatch_invoke(prompts, **batch_kw) + + dropped: set[str] = set() + retyped: dict[str, str] = {} + unverified_quotes = 0 + for item in results: + if not item.ok or item.response is None: + ctx.log( + f"Entity verification batch {item.index} failed, keeping all " + f"its entities: {item.error}", + logging.WARNING, + ) + continue + try: + rows = json.loads(_strip_markdown_fences(item.response.content)) + except (ValueError, TypeError): + ctx.log( + f"Entity verification batch {item.index} returned unparseable " + f"JSON, keeping all its entities", + logging.WARNING, + ) + continue + if not isinstance(rows, list): + continue + batch = batches[item.index] + for row in rows: + if not isinstance(row, dict): + continue + try: + key = batch[int(row.get("id", 0)) - 1] + except (ValueError, TypeError, IndexError): + continue + if str(row.get("verdict", "")).strip().lower() == "drop": + dropped.add(key) + continue + # A "keep" must be justified by a quote that really exists. + quote = str(row.get("quote", "")).strip() + _ent, context = first[key] + if quote and quote not in context: + unverified_quotes += 1 + new_type = str(row.get("type", "")).strip() + if new_type and new_type in entity_types: + retyped[key] = new_type + + out: list[list[ExtractedEntity]] = [] + for ents in chunk_entities: + kept = [] + for ent in ents: + key = ent.name.strip().casefold() + if key in dropped: + continue + if key in retyped and ent.type != UNKNOWN_LABEL: + ent.type = retyped[key] + kept.append(ent) + out.append(kept) + + total = sum(len(e) for e in chunk_entities) + ctx.log( + f"Entity verification: {len(keys)} distinct names in " + f"{len(prompts)} call(s); dropped {len(dropped)} names " + f"({total} -> {sum(len(e) for e in out)} mentions); " + f"retyped {len(retyped)}; {unverified_quotes} quote(s) not found " + f"in source" + ) + return out async def extract( self, @@ -442,7 +799,10 @@ async def extract( ) -> GraphData: # Resolve entity types: ontology overrides instance default if ontology.entities: - entity_types = [e.label for e in ontology.entities] + # Ontology labels bypass the constructor, so re-check here: an + # ontology declaring a reserved label must fail as loudly as + # passing one to entity_types would. + entity_types = _reject_reserved_labels([e.label for e in ontology.entities]) entity_type_descs: dict[str, str] = { e.label: e.description for e in ontology.entities if e.description } @@ -450,6 +810,15 @@ async def extract( entity_types = list(self.entity_types) entity_type_descs = {} + # Relation vocabulary. A declared ontology wins, and that path also + # prunes non-conforming edges downstream. Otherwise fall back to the + # instance default, which only steers the prompt — nothing is pruned, + # mirroring how entity_types behaves. + if ontology.relations: + prompt_relations = list(ontology.relations) + else: + prompt_relations = [Relation(label=lbl) for lbl in self.relation_types] + ctx.log( f"Extracting from {len(chunks.chunks)} chunks (hybrid, " f"extractor={self.entity_extractor.__class__.__name__}, " @@ -467,6 +836,10 @@ async def extract( if not active_chunks: return GraphData(nodes=[], relationships=[]) + # Chunks whose extraction raised. Tracked so the caller can tell an + # empty document from a broken one, and retry just these uids. + failed_chunk_uids: set[str] = set() + # ── Optional: Coreference resolution per chunk ── chunk_texts: list[str] = [] if self.coref_resolver is not None: @@ -537,10 +910,22 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: f"Step 1 NER failed for chunk {active_chunks[i].index}: {result}", logging.WARNING, ) + failed_chunk_uids.add(active_chunks[i].uid) chunk_entities.append([]) else: chunk_entities.append(result) + # ── Step 1b: dedicated entity verification ── + if self.verify_entities: + chunk_entities = await self._verify_entities( + chunk_entities, + chunk_texts, + [c.uid for c in active_chunks], + entity_types, + entity_type_descs, + ctx, + ) + # ── Step 2: LLM verify + relationship extraction ── step2_prompts: list[str] = [] step2_indices: list[int] = [] # maps prompt index -> active_chunk index @@ -552,9 +937,9 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: has_attrs = _ontology_has_attributes(ontology) prompt = VERIFY_EXTRACT_RELS_PROMPT.format( entity_types=_format_entity_types(entity_types, entity_type_descs), - relation_patterns=_format_relation_patterns(ontology.relations), + relation_patterns=_format_relation_patterns(prompt_relations), attribute_block=_render_attribute_block(ontology), - relationship_type_instruction=_relationship_type_instruction(ontology.relations), + relationship_type_instruction=_relationship_type_instruction(prompt_relations), entities_json=entities_json, text=text, json_example=_JSON_EXAMPLE_WITH_ATTRS if has_attrs else _DEFAULT_JSON_EXAMPLE, @@ -568,6 +953,7 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: all_entities: list[ExtractedEntity] = [] all_relations: list[ExtractedRelation] = [] + rels_by_chunk: dict[int, list[ExtractedRelation]] = {} if step2_prompts: step2_results = await self.llm.abatch_invoke(step2_prompts, **batch_kw2) @@ -580,6 +966,10 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: f"Step 2 verify+rels failed for chunk {chunk.index}: {item.error}", logging.WARNING, ) + # Relations for this chunk are lost even though step 1 + # entities survive, so it counts as failed: the caller + # needs to know this chunk's edges were never extracted. + failed_chunk_uids.add(chunk.uid) # Fall back to step 1 entities only all_entities.extend(chunk_entities[chunk_idx]) continue @@ -601,6 +991,9 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: else: # LLM returned no entities — use step 1 entities all_entities.extend(chunk_entities[chunk_idx]) + rels_by_chunk.setdefault(chunk_idx, []).extend(rels) + + for rels in rels_by_chunk.values(): all_relations.extend(rels) # ── Aggregate across chunks ── @@ -629,12 +1022,22 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: mentions=all_mentions, extracted_entities=merged_entities, extracted_relations=merged_relations, + chunks_attempted=len(active_chunks), + failed_chunks=sorted(failed_chunk_uids), ) ctx.log( f"Extracted {len(nodes)} nodes, {len(relationships)} relationships, " f"{len(all_mentions)} mentions" ) + if failed_chunk_uids: + # Escalate: an INFO summary saying "0 nodes" reads as an empty + # document. Say plainly that chunks were lost. + ctx.log( + f"{len(failed_chunk_uids)} of {len(active_chunks)} chunks failed " + f"extraction and contributed nothing to the graph", + logging.WARNING, + ) return graph_data @staticmethod diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index c2cef098..4a445f73 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -14,6 +14,7 @@ from graphrag_sdk.core.connection import FalkorDBConnection from graphrag_sdk.core.exceptions import DatabaseError from graphrag_sdk.core.models import ( + RESERVED_NODE_LABELS, ChunkEntityRow, ChunkRelationshipRow, DocumentRecord, @@ -48,11 +49,12 @@ class GraphStore: def __init__(self, connection: FalkorDBConnection) -> None: self._conn = connection + self._indexed_labels: set[str] = set() # ── Write Operations ───────────────────────────────────────── _BATCH_SIZE = 500 - _STRUCTURAL_LABELS = frozenset({"Chunk", "Document"}) + _STRUCTURAL_LABELS = RESERVED_NODE_LABELS _REL_LABEL_HINTS: dict[str, tuple[str, str]] = { "PART_OF": ("Document", "Chunk"), "NEXT_CHUNK": ("Chunk", "Chunk"), @@ -60,6 +62,45 @@ def __init__(self, connection: FalkorDBConnection) -> None: "RELATES": ("__Entity__", "__Entity__"), } + async def _ensure_id_index(self, safe_label: str) -> None: + """Create a range index on ``id`` for a label, once per label per store. + + Every write path here addresses nodes by ``{id: ...}`` — ``MERGE`` in + :meth:`upsert_nodes`, and a double ``MATCH`` in + :meth:`upsert_relationships`. Without a range index each of those scans + every node carrying the label, so the cost of writing one batch grows + linearly with the graph and total ingest cost grows quadratically. + + Measured against FalkorDB v4.18.0, writing 50K nodes in batches of 500 + using the exact MERGE this class issues: + + ============= =========== =========== ========== + arm first batch last batch total + ============= =========== =========== ========== + no index 4.6 ms 2194.8 ms 111.47 s + range index 8.8 ms 8.5 ms 0.69 s + ============= =========== =========== ========== + + That is 477x degradation across the run unindexed versus 0.97x (flat) + indexed, and 161x less total time. The degradation is why bulk ingest + "gets slower as the graph grows" — it is not the LLM, it is this. + + Creating the index is idempotent in FalkorDB but still a round trip, so + results are memoised per label. The cache is per-store-instance, which + is the right scope: a new store means a possibly different graph. + + Failures are logged and swallowed. A missing index is a performance + problem; a raised exception here would be a correctness problem, and + the caller's write must not depend on the optimisation succeeding. + """ + if safe_label in self._indexed_labels: + return + self._indexed_labels.add(safe_label) + try: + await self._conn.query(f"CREATE INDEX FOR (n:`{safe_label}`) ON (n.id)") + except Exception as exc: # noqa: BLE001 - optimisation only + logger.debug("Could not create id index on %s: %s", safe_label, exc) + async def upsert_nodes(self, nodes: list[GraphNode]) -> int: """Batch upsert nodes using UNWIND, grouped by label. @@ -99,6 +140,9 @@ async def upsert_nodes(self, nodes: list[GraphNode]) -> int: ) if not cleaned_group: continue + await self._ensure_id_index(safe_label) + if is_entity: + await self._ensure_id_index("__Entity__") # Process in batches for start in range(0, len(cleaned_group), self._BATCH_SIZE): batch = cleaned_group[start : start + self._BATCH_SIZE] @@ -189,6 +233,9 @@ async def upsert_relationships(self, relationships: list[GraphRelationship]) -> ) if not cleaned_group: continue + hint_src, hint_tgt = self._REL_LABEL_HINTS.get(rel_type, ("__Entity__", "__Entity__")) + await self._ensure_id_index(sanitize_cypher_label(hint_src)) + await self._ensure_id_index(sanitize_cypher_label(hint_tgt)) for start in range(0, len(cleaned_group), self._BATCH_SIZE): batch = cleaned_group[start : start + self._BATCH_SIZE] batch_data = [ diff --git a/graphrag_sdk/tests/test_entity_extractors.py b/graphrag_sdk/tests/test_entity_extractors.py index c6790230..6d709f95 100644 --- a/graphrag_sdk/tests/test_entity_extractors.py +++ b/graphrag_sdk/tests/test_entity_extractors.py @@ -164,3 +164,62 @@ async def extract_entities(self, text, entity_types, source_chunk_id): )] assert isinstance(MyExtractor(), EntityExtractor) + + +class TestGLiNERModelSharing: + """Bugs #8 and #11 — one model copy per extractor, and a serialising lock. + + #11: with current (not peak) RSS, six extractors each loading their own + model went 74.5 MB -> 2447 MB, ~395 MB marginal per copy, projecting + ~11.6 GB for 30 concurrent documents. With the shared cache the same six + sit at 1425.3 -> 1425.4 MB: 0.0 MB marginal, ~1.39 GB projected. + + #8: inference ran under a per-instance lock while the caller dispatched it + via ``asyncio.to_thread``, so concurrent documents serialised. Counter- + balanced over eight documents: locked 3.75/3.54 s, unlocked 2.21/2.46 s = + 1.56x. Removing it is only sound because GLiNER inference proved + thread-safe — 40/40 documents byte-identical against the serialised run. + """ + + def _extractor(self, monkeypatch, loads): + from graphrag_sdk.ingestion.extraction_strategies import entity_extractors as ee + + class FakeGLiNER: + @staticmethod + def from_pretrained(name): + loads.append(name) + return object() + + monkeypatch.setitem(__import__("sys").modules, "gliner", + type("m", (), {"GLiNER": FakeGLiNER})) + monkeypatch.setattr(ee.GLiNERExtractor, "_MODEL_CACHE", {}, raising=False) + return ee.GLiNERExtractor + + def test_model_loaded_once_across_instances(self, monkeypatch): + loads = [] + cls = self._extractor(monkeypatch, loads) + models = [cls()._load_model() for _ in range(5)] + assert len(loads) == 1 + assert len({id(m) for m in models}) == 1 + + def test_different_models_are_not_shared(self, monkeypatch): + loads = [] + cls = self._extractor(monkeypatch, loads) + a = cls(model_name="model-a", threshold=0.5)._load_model() + b = cls(model_name="model-b", threshold=0.5)._load_model() + assert loads == ["model-a", "model-b"] + assert a is not b + + def test_inference_takes_no_lock(self): + """Guards bug #8 against reintroduction.""" + import inspect + + from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( + GLiNERExtractor, + ) + + src = inspect.getsource(GLiNERExtractor._predict_sync) + code = "\n".join( + line for line in src.splitlines() if not line.strip().startswith("#") + ) + assert "self._lock" not in code diff --git a/graphrag_sdk/tests/test_graph_extraction.py b/graphrag_sdk/tests/test_graph_extraction.py index 1c0bcc3b..8f46f318 100644 --- a/graphrag_sdk/tests/test_graph_extraction.py +++ b/graphrag_sdk/tests/test_graph_extraction.py @@ -13,18 +13,26 @@ Ontology, Entity, Relation, + RESERVED_NODE_LABELS, TextChunk, TextChunks, ) +from graphrag_sdk.storage.graph_store import GraphStore from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( EntityExtractor, LLMExtractor, ) +from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( + is_valid_entity_name, +) from graphrag_sdk.ingestion.extraction_strategies.graph_extraction import ( GraphExtraction, VERIFY_EXTRACT_RELS_PROMPT, _format_entity_types, + _reject_reserved_labels, + DEFAULT_RELATION_TYPES, _format_relation_patterns, + _relationship_type_instruction, ) from .conftest import MockLLM, MockLLMWithGraphExtraction @@ -288,6 +296,54 @@ def test_custom_entity_types(self): assert extractor.entity_types == ["Vehicle", "Road"] +class TestReservedNodeLabels: + """`Document`/`Chunk` are the graph store's bookkeeping labels. + + Reusing one as an entity type used to corrupt the graph silently: document + counts picked up extracted entities (an 11-document corpus reported 106), + and `GraphStore._write_nodes` skips `__Entity__` for structural labels, so + the entity vanished from dedup and retrieval without any error. + """ + + @pytest.mark.parametrize("bad", ["Document", "Chunk", "document", "cHuNk"]) + def test_reserved_entity_type_is_rejected(self, bad): + llm = MockLLM() + with pytest.raises(ValueError, match="reserved label"): + GraphExtraction( + llm=llm, + entity_extractor=LLMExtractor(llm), + entity_types=["Person", bad], + ) + + def test_error_names_the_offending_label(self): + llm = MockLLM() + with pytest.raises(ValueError, match="Document"): + GraphExtraction( + llm=llm, + entity_extractor=LLMExtractor(llm), + entity_types=["Document"], + ) + + def test_non_reserved_types_still_allowed(self): + llm = MockLLM() + extractor = GraphExtraction( + llm=llm, + entity_extractor=LLMExtractor(llm), + entity_types=["Publication", "TextSegment"], + ) + assert extractor.entity_types == ["Publication", "TextSegment"] + + def test_ontology_labels_are_checked_too(self): + """The ontology path assigns entity types without going through + __init__, so it needs its own guard or the fix is bypassable.""" + with pytest.raises(ValueError, match="reserved label"): + _reject_reserved_labels(["Person", "Document"]) + + def test_store_and_extractor_share_one_definition(self): + """Two hardcoded copies would drift; the bug returns when they do.""" + assert GraphStore._STRUCTURAL_LABELS is RESERVED_NODE_LABELS + + class TestGraphExtractionStep2Parsing: def test_parse_valid_response(self): content = json.dumps({ @@ -574,17 +630,58 @@ def test_relationship_spans_merge_across_chunks(self): assert "chunk-1" in merged[0].spans -class TestNoiseFilteringPrompt: - """Bug 4: VERIFY_EXTRACT_RELS_PROMPT should contain noise-filtering instructions.""" +class TestNoiseFiltering: + """Bug 4: operator/abbreviation/short-token noise must be filtered out. + + These rules used to live as instructions inside VERIFY_EXTRACT_RELS_PROMPT + and were asserted by checking the prompt's wording. They now live in + ``is_valid_entity_name`` instead, after measurement showed the LLM does not + reliably act on verification instructions (RESULTS.md P2.10). Asserting the + behaviour rather than the prompt text is also what these tests should have + done in the first place: the old version passed whether or not anything was + actually filtered. + """ - def test_prompt_contains_operator_filtering(self): - assert "symbolic" in VERIFY_EXTRACT_RELS_PROMPT.lower() + @pytest.mark.parametrize("name", ["+=", "->", "++", "==", "!="]) + def test_operator_tokens_rejected(self, name): + assert not is_valid_entity_name(name) - def test_prompt_contains_abbreviation_filtering(self): - assert "non-domain-specific" in VERIFY_EXTRACT_RELS_PROMPT + @pytest.mark.parametrize("name", ["sh", "cd", "ls", "rm", "cp", "mv"]) + def test_shell_abbreviations_rejected(self, name): + assert not is_valid_entity_name(name) - def test_prompt_contains_short_token_filtering(self): - assert "1-2 characters" in VERIFY_EXTRACT_RELS_PROMPT + @pytest.mark.parametrize("name", ["dt", "bg", "fn"]) + def test_generic_short_tokens_rejected(self, name): + assert not is_valid_entity_name(name) + + @pytest.mark.parametrize("name", ["AI", "US", "UK", "Go", "EU", "UN"]) + def test_real_acronyms_kept(self, name): + """The filter must not take widely-recognised acronyms with it.""" + assert is_valid_entity_name(name) + + @pytest.mark.parametrize("name", ["1823", "1957", "1003 ce", "14 january 1904"]) + def test_specific_dates_rejected(self, name): + """A date pins down a moment; it is an attribute, not an entity.""" + assert not is_valid_entity_name(name) + + @pytest.mark.parametrize("name", ["1820s", "19th century", "Abbasid era"]) + def test_periods_kept(self, name): + """A period is something facts attach to, so it stays a node.""" + assert is_valid_entity_name(name) + + @pytest.mark.parametrize("name", ["Boeing 747", "COVID-19"]) + def test_numeric_names_not_mistaken_for_dates(self, name): + assert is_valid_entity_name(name) + + def test_prompt_still_asks_the_llm_to_verify(self): + """Removing this instruction was tried and reverted (RESULTS.md P2.12). + + Without it the LLM emitted 813 entities instead of 719 and entity + precision fell 0.645 -> 0.551. The code-side rules above are a floor, + not a replacement. + """ + assert "REMOVE any entity" in VERIFY_EXTRACT_RELS_PROMPT + assert "VERIFY the entities" in VERIFY_EXTRACT_RELS_PROMPT class TestEntityTypeDescriptions: @@ -836,3 +933,162 @@ def test_property_less_schema_keeps_attributes_empty(self): ) assert len(ents) == 1 assert ents[0].attributes == {} + + +class TestDefaultRelationTypes: + """The shipped relation vocabulary — the counterpart to DEFAULT_ENTITY_TYPES. + + Entity extraction always shipped a default type list; relations shipped + nothing, so the prompt asked the model to invent a label per edge. Measured + on an 11-document corpus that produced 447 distinct labels against 30 in + gold. Supplying a default list doubled exact triple F1 (0.065 -> 0.134) with + no loss of recall, and held across five unrelated Wikipedia domains + (vocabulary 2.2-2.9x smaller, 10-18% -> 66-81% of edges on the list). + """ + + def test_default_is_applied_when_nothing_is_passed(self): + ge = GraphExtraction(llm=MockLLM()) + assert ge.relation_types == list(DEFAULT_RELATION_TYPES) + assert len(ge.relation_types) > 0 + + def test_explicit_list_overrides_the_default(self): + ge = GraphExtraction(llm=MockLLM(), relation_types=["eats", "owns"]) + assert ge.relation_types == ["eats", "owns"] + + def test_empty_list_restores_open_vocabulary(self): + """``[]`` is a request, not an omission. + + Guards the ``is None`` check: a truthiness test would silently swap an + explicit open-vocabulary request for the default list. + """ + ge = GraphExtraction(llm=MockLLM(), relation_types=[]) + assert ge.relation_types == [] + + def test_default_list_is_not_shared_between_instances(self): + a = GraphExtraction(llm=MockLLM()) + b = GraphExtraction(llm=MockLLM()) + a.relation_types.append("mutated") + assert "mutated" not in b.relation_types + assert "mutated" not in DEFAULT_RELATION_TYPES + + def test_default_labels_are_well_formed(self): + for label in DEFAULT_RELATION_TYPES: + assert label == label.lower(), f"{label} is not lower_snake_case" + assert " " not in label, f"{label} contains a space" + assert label.replace("_", "").isalpha(), f"{label} has odd characters" + assert len(set(DEFAULT_RELATION_TYPES)) == len(DEFAULT_RELATION_TYPES) + + def test_default_list_reaches_the_prompt(self): + """The list is worthless if it never renders into the prompt.""" + rels = [Relation(label=lbl) for lbl in DEFAULT_RELATION_TYPES] + block = _format_relation_patterns(rels) + assert "## Allowed Relationships" in block + for label in DEFAULT_RELATION_TYPES: + assert label in block + assert "MUST be one of" in _relationship_type_instruction(rels) + + def test_open_vocabulary_prompt_when_list_is_empty(self): + assert _format_relation_patterns([]) == "" + assert "UPPER_SNAKE_CASE" in _relationship_type_instruction([]) + + +class TestFailedChunkReporting: + """Finding #7: a broken extraction must not look like an empty document. + + Before the fix, per-chunk failures were swallowed and replaced with + empty results, so a document whose chunks all failed returned byte + identical output to a document that genuinely contained no entities. + These three arms mirror the benchmark harness that proved it. + """ + + class _ScriptedExtractor(EntityExtractor): + """Fails on the first ``n_fail`` chunks, returns [] for the rest.""" + + def __init__(self, n_fail: int) -> None: + self._n_fail = n_fail + self.calls = 0 + + async def extract_entities( + self, text: str, entity_types: list[str], source_chunk_id: str + ) -> list[ExtractedEntity]: + index = self.calls + self.calls += 1 + if index < self._n_fail: + raise RuntimeError(f"simulated NER failure on chunk {index}") + return [] + + async def _run(self, n_fail: int, ctx, *, silent_llm: bool = False): + extractor = self._ScriptedExtractor(n_fail) + # silent_llm: step 2 also yields nothing, so the graph really is empty + # in every arm and only the new fields can tell the arms apart. + llm = ( + _mock_hybrid_llm(step1_entities=[], step2_entities=[], step2_relationships=[]) + if silent_llm + else _mock_hybrid_llm() + ) + strategy = GraphExtraction(llm=llm, entity_extractor=extractor) + chunks = _make_chunks("one.", "two.", "three.", "four.") + result = await strategy.extract(chunks, Ontology(), ctx) + # Guard against the instrument silently not running: if the extractor + # was never called, every assertion below is vacuously true. + assert extractor.calls == 4, "extractor did not run on all 4 chunks" + return result + + async def test_empty_document_is_not_reported_as_failed(self, ctx): + result = await self._run(0, ctx) + assert result.chunks_attempted == 4 + assert result.failed_chunks == [] + assert result.extraction_failed is False + + async def test_total_failure_is_reported(self, ctx): + result = await self._run(4, ctx) + assert result.chunks_attempted == 4 + assert len(result.failed_chunks) == 4 + assert result.extraction_failed is True + + async def test_partial_failure_reports_only_the_failed_chunks(self, ctx): + result = await self._run(2, ctx) + assert result.chunks_attempted == 4 + assert len(result.failed_chunks) == 2 + # Not a total failure: the surviving chunks did their job. + assert result.extraction_failed is False + + async def test_empty_and_broken_are_distinguishable(self, ctx): + """The finding itself: these two used to be identical.""" + empty = await self._run(0, ctx, silent_llm=True) + broken = await self._run(4, ctx, silent_llm=True) + assert empty.nodes == broken.nodes == [] + assert (empty.chunks_attempted, empty.failed_chunks, empty.extraction_failed) != ( + broken.chunks_attempted, + broken.failed_chunks, + broken.extraction_failed, + ) + + async def test_failed_chunks_are_retryable_ids(self, ctx): + """A count is not enough; the caller must be able to re-ingest.""" + result = await self._run(2, ctx) + chunks = _make_chunks("one.", "two.", "three.", "four.") + known = {c.uid for c in chunks.chunks} + assert set(result.failed_chunks) <= known + assert all(isinstance(uid, str) and uid for uid in result.failed_chunks) + + async def test_no_chunks_reports_nothing_attempted(self, ctx): + strategy = GraphExtraction( + llm=_mock_hybrid_llm(), entity_extractor=self._ScriptedExtractor(0) + ) + result = await strategy.extract(TextChunks(chunks=[]), Ontology(), ctx) + assert result.chunks_attempted == 0 + assert result.failed_chunks == [] + assert result.extraction_failed is False + + async def test_successful_extraction_reports_no_failures(self, ctx): + """Guard the happy path: normal ingests must stay clean.""" + strategy = GraphExtraction( + llm=_mock_hybrid_llm(), entity_extractor=LLMExtractor(_mock_hybrid_llm()) + ) + chunks = _make_chunks("Alice is a software engineer at Acme Corp.") + result = await strategy.extract(chunks, Ontology(), ctx) + assert len(result.nodes) > 0 + assert result.chunks_attempted == 1 + assert result.failed_chunks == [] + assert result.extraction_failed is False diff --git a/graphrag_sdk/tests/test_graph_store.py b/graphrag_sdk/tests/test_graph_store.py index 18e3ff64..ff8bbb18 100644 --- a/graphrag_sdk/tests/test_graph_store.py +++ b/graphrag_sdk/tests/test_graph_store.py @@ -17,13 +17,28 @@ def graph_store(mock_connection): return GraphStore(mock_connection) +def upsert_calls(mock_connection): + """Query calls excluding the lazy ``CREATE INDEX`` round trips. + + ``upsert_nodes`` / ``upsert_relationships`` now ensure a range index on + ``id`` before writing (once per label per store). Asserting on raw call + counts would pin the tests to that bookkeeping instead of the write they + are actually about, so index calls are filtered out here. + """ + return [ + c for c in mock_connection.query.call_args_list + if "CREATE INDEX" not in c[0][0] + ] + + class TestGraphStoreUpsertNodes: async def test_upsert_single_node(self, graph_store, mock_connection): nodes = [GraphNode(id="n1", label="Person", properties={"name": "Alice"})] result = await graph_store.upsert_nodes(nodes) assert result == 1 - mock_connection.query.assert_called_once() - cypher = mock_connection.query.call_args[0][0] + calls = upsert_calls(mock_connection) + assert len(calls) == 1 + cypher = calls[0][0][0] assert "UNWIND" in cypher assert "MERGE" in cypher assert "Person" in cypher @@ -36,7 +51,7 @@ async def test_upsert_multiple_nodes(self, graph_store, mock_connection): ] result = await graph_store.upsert_nodes(nodes) assert result == 2 - assert mock_connection.query.call_count == 2 + assert len(upsert_calls(mock_connection)) == 2 async def test_upsert_empty_list(self, graph_store, mock_connection): result = await graph_store.upsert_nodes([]) @@ -50,7 +65,7 @@ async def test_upsert_raises_on_error(self, graph_store, mock_connection): async def test_upsert_passes_id_in_batch_param(self, graph_store, mock_connection): await graph_store.upsert_nodes([GraphNode(id="test-id", label="X", properties={})]) - params = mock_connection.query.call_args[0][1] + params = upsert_calls(mock_connection)[0][0][1] assert params["batch"][0]["id"] == "test-id" async def test_upsert_sanitizes_control_chars_in_batch_params( @@ -59,7 +74,7 @@ async def test_upsert_sanitizes_control_chars_in_batch_params( await graph_store.upsert_nodes( [GraphNode(id="id\x00\x01", label="Chunk", properties={"text": "A\x00B\x01C"})] ) - params = mock_connection.query.call_args[0][1] + params = upsert_calls(mock_connection)[0][0][1] assert params["batch"][0]["id"] == "id" assert params["batch"][0]["properties"]["text"] == "ABC" @@ -67,12 +82,21 @@ async def test_upsert_sanitizes_control_chars_in_fallback_params( self, graph_store, mock_connection ): """Per-item fallback path should also use sanitized IDs and properties.""" - mock_connection.query = AsyncMock(side_effect=[Exception("batch fail"), MagicMock()]) + # Index creation happens first now, so the batch write is the call + # after those; let every CREATE INDEX succeed, fail the batch, then + # succeed the per-item fallback. + def side_effect(cypher, params=None): + if "CREATE INDEX" in cypher: + return MagicMock() + if "UNWIND" in cypher: + raise Exception("batch fail") + return MagicMock() + + mock_connection.query = AsyncMock(side_effect=side_effect) await graph_store.upsert_nodes( [GraphNode(id="id\x00\x01", label="X", properties={"t": "A\x00B"})] ) - # Second call is the per-item fallback - fallback_params = mock_connection.query.call_args_list[1][0][1] + fallback_params = upsert_calls(mock_connection)[1][0][1] assert fallback_params["id"] == "id" assert fallback_params["properties"]["t"] == "AB" @@ -586,3 +610,79 @@ async def test_delete_orphan_entities_batches_large_lists( n = await graph_store.delete_orphan_entities(ids) assert n == 6 assert mock_connection.query.await_count == 3 + + +class TestGraphStoreIdIndex: + """Bug #9 — range index on ``id`` for every label we MERGE/MATCH by id. + + Measured against FalkorDB v4.18.0 writing 50K nodes with this class's own + MERGE: unindexed the last batch was 477x slower than the first (4.6 ms -> + 2194.8 ms, 111 s total); indexed it stayed flat (8.8 -> 8.5 ms, 0.69 s). + Driving the real ``GraphStore`` rather than raw Cypher showed 8.1x less + total time at 20K nodes. This is the "gets slower as the graph grows" + complaint. + """ + + @staticmethod + def index_calls(mock_connection): + return [ + c[0][0] for c in mock_connection.query.call_args_list + if "CREATE INDEX" in c[0][0] + ] + + async def test_creates_index_for_node_label_and_entity( + self, graph_store, mock_connection + ): + await graph_store.upsert_nodes( + [GraphNode(id="n1", label="Person", properties={})] + ) + idx = self.index_calls(mock_connection) + assert "CREATE INDEX FOR (n:`Person`) ON (n.id)" in idx + # MERGE targets :Person but relationship MATCHes target :__Entity__. + assert "CREATE INDEX FOR (n:`__Entity__`) ON (n.id)" in idx + + async def test_structural_labels_get_index_but_not_entity( + self, graph_store, mock_connection + ): + await graph_store.upsert_nodes( + [GraphNode(id="c1", label="Chunk", properties={})] + ) + idx = self.index_calls(mock_connection) + assert "CREATE INDEX FOR (n:`Chunk`) ON (n.id)" in idx + assert "CREATE INDEX FOR (n:`__Entity__`) ON (n.id)" not in idx + + async def test_index_created_once_per_label(self, graph_store, mock_connection): + for i in range(5): + await graph_store.upsert_nodes( + [GraphNode(id=f"n{i}", label="Person", properties={})] + ) + idx = self.index_calls(mock_connection) + assert idx.count("CREATE INDEX FOR (n:`Person`) ON (n.id)") == 1 + + async def test_relationship_upsert_indexes_both_endpoints( + self, graph_store, mock_connection + ): + await graph_store.upsert_relationships( + [GraphRelationship( + start_node_id="e1", end_node_id="c1", + type="MENTIONED_IN", properties={}, + )] + ) + idx = self.index_calls(mock_connection) + assert "CREATE INDEX FOR (n:`__Entity__`) ON (n.id)" in idx + assert "CREATE INDEX FOR (n:`Chunk`) ON (n.id)" in idx + + async def test_index_failure_does_not_break_the_write( + self, graph_store, mock_connection + ): + """A missing index is slow; a raised exception would be data loss.""" + def side_effect(cypher, params=None): + if "CREATE INDEX" in cypher: + raise Exception("index unsupported") + return MagicMock() + + mock_connection.query = AsyncMock(side_effect=side_effect) + result = await graph_store.upsert_nodes( + [GraphNode(id="n1", label="Person", properties={})] + ) + assert result == 1