From ddbcb4d17e78343d09a3c4de3969d498f39921dc Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:05:59 +0300 Subject: [PATCH 01/19] fix(extraction): reject specific dates as entity names 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. It also stops the graph accumulating "X happened_in 1957" edges that carry no answerable content. Known trade-off, documented in the code: a product genuinely named for a number ("747", "1984") is indistinguishable from a year here and will be dropped. That 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. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) 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..052e9b56 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 @@ -127,6 +127,48 @@ 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(): @@ -136,6 +178,8 @@ def is_valid_entity_name(name: str) -> bool: return False if stripped.lower() in _ENTITY_STOPLIST: return False + if is_specific_date(stripped): + return False return True From 6371f7d6bbe4588a326a88736d91bd6da8b7987d Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:06:11 +0300 Subject: [PATCH 02/19] fix(extraction): read short all-caps tokens as acronyms, not pronouns The stoplist check casefolded the candidate name before comparing, so "US" the country matched "us" the pronoun and every mention of the United States was silently discarded. The same collision affects any short acronym that happens to spell an English function word. A token that is short, all-caps and alphabetic is an acronym, so it now skips the pronoun stoplist. Lowercase "us", "it", "its" are unaffected. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ingestion/extraction_strategies/entity_extractors.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 052e9b56..24300ed6 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 @@ -176,7 +176,10 @@ 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 From 3fcc7d7cac3c24465a17b897d79419337f679784 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:06:28 +0300 Subject: [PATCH 03/19] fix(extraction): hard-filter junk tokens in code, not only in the prompt Shell fragments and system abbreviations read as entities to an NER model but name nothing in any document domain. Removing them was previously delegated to an instruction in the extraction prompt, which is not a guarantee: the model can quietly skip it, and the result varies between runs (see RESULTS.md P2.10). A fixed list costs nothing and cannot be ignored. Also rejects names with no alphanumeric character at all (+=, ->, ==, !=), which no prompt instruction reliably caught. Well-known short acronyms - AI, US, UK, EU, UN, Go - are deliberately absent from the list and continue to extract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../extraction_strategies/entity_extractors.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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 24300ed6..d7fb4d73 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 @@ -67,7 +67,19 @@ "its", } -_ENTITY_STOPLIST: set[str] = _PRONOUNS | { +# 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", @@ -183,6 +195,10 @@ def is_valid_entity_name(name: str) -> bool: 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 From 266967a52dc0ca60ccfb6ab7addd1f7779e2bbc3 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:07:13 +0300 Subject: [PATCH 04/19] feat(extraction): switch default GLiNER model and make thresholds per-model Two changes that only make sense together: the new default model does not share the old model's confidence scale, so changing the model without also making the threshold model-aware would silently mis-tune every install. Default model: urchade/gliner_medium-v2.1 -> knowledgator/gliner-bi-small-v2.0. Measured 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 input window vs 384. Smaller, faster, and it finds more. Thresholds: the cutoff was one hardcoded 0.75 applied to every model. DEFAULT_THRESHOLDS now records the measured operating point per model and threshold=None looks it up; an unknown model falls back to GLiNER's own 0.5 and logs a warning rather than silently inheriting a number tuned elsewhere. That single hardcoded number was the actual cause of a result we could not explain for some time: a 4-second spaCy model appearing to beat our shipped config. The bi-encoder models return almost nothing at 0.75 - measured at 2 entities for an entire corpus, with no error raised. An explicitly passed threshold still wins, so existing callers are unaffected. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 48 +++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) 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 d7fb4d73..6750600e 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 @@ -380,18 +380,60 @@ class GLiNERExtractor(EntityExtractor): so a single instance can be safely shared across concurrent ``asyncio.to_thread`` calls (e.g. parallel doc ingestion). + **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. """ + #: 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, ) -> 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 self._model: Any = None self._lock = threading.Lock() From 6dd301e39cc277f462461d81141ea9d1b062171e Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:07:37 +0300 Subject: [PATCH 05/19] perf(extraction): share one GLiNER model per process instead of one per ingest 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, and matches the customer report of large resident memory under concurrent ingest. With the cache the same workload holds one copy, ~1.39 GB. The cache is keyed on model name, not shared unconditionally, because two extractors may legitimately want different models; sharing those would make one silently answer with the other's weights. Loading is double-checked under a class-level lock so concurrent first-use cannot load twice. Refs FalkorDB/research#71 bug #11. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 52 ++++++++++++++----- 1 file changed, 40 insertions(+), 12 deletions(-) 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 6750600e..7b21ceaa 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,7 @@ import re import threading from abc import ABC, abstractmethod -from typing import Any +from typing import Any, ClassVar from graphrag_sdk.core.models import ExtractedEntity from graphrag_sdk.core.providers import LLMInterface @@ -439,19 +439,47 @@ def __init__( 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 _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] From af9d616b641ee592384884efd0aab02aa92f8f77 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:08:10 +0300 Subject: [PATCH 06/19] perf(extraction): drop the GLiNER inference lock Inference ran under self._lock while the caller dispatched it through asyncio.to_thread. The SDK paid for threads and then serialised them anyway, so concurrent documents queued behind each other in NER. Removing the lock is only sound if GLiNER inference is genuinely thread-safe, so that was tested rather than assumed: eight documents extracted concurrently and unlocked, compared field-by-field against the serialised result, five trials - 40/40 documents byte-identical. Inference mutates no model state; only model loading 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 the result: locked 3.75s / 3.54s versus unlocked 2.21s / 2.46s = 1.56x. Issue #71 claimed 3.44x. The honest measured figure is 1.56x, because torch's own intra-op threading already uses the cores. self._lock is retained as an attribute so callers and tests that swap in a context manager to A/B this change keep working. Nothing acquires it. Refs FalkorDB/research#71 bug #8. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 38 ++++++++++++++++++- 1 file changed, 36 insertions(+), 2 deletions(-) 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 7b21ceaa..7a56192e 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 @@ -435,6 +435,11 @@ def __init__( ) self._threshold = 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() def _load_model(self) -> Any: @@ -483,8 +488,37 @@ def _get_shared_model(cls, model_name: str) -> Any: 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) + + # 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) + + def _predict_body( + self, + model: Any, + text: str, + labels: list[str], + ) -> list[dict[str, Any]]: + """Inference. Split out from :meth:`_predict_sync` so the lock removal + above could be A/B tested by wrapping one call site.""" + return model.predict_entities(text, labels, threshold=self._threshold) async def extract_entities( self, From 9da4881d56b05db5af406cd195858867eceedd2c Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:09:12 +0300 Subject: [PATCH 07/19] fix(extraction): window GLiNER over long chunks instead of truncating silently GLiNER has a hard input limit (config.max_len) and discards 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 new default model raises the limit to 2048, but real documents still exceed it: disabling windowing on the benchmark corpus dropped entity recall from 0.805 to 0.621. Windowing matters regardless of the model. Text longer than window_tokens is now processed as a series of overlapping windows and the predictions merged. Short text takes a fast path and behaves exactly as before, so nothing changes for chunk-sized input. Two details that matter for correctness: - Predictions are re-offset into the original text, so character spans stay valid for downstream span-based code. - window_overlap defaults to 48 word-tokens, comfortably above the model's max_width (longest representable entity, 12 words), or an entity sitting on a window boundary would be lost by both neighbours. _merge collapses duplicate spans to the highest-scoring copy and drops any span strictly contained in a longer span of the same label - that is the clipped remains of an entity the neighbouring window saw whole. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 122 +++++++++++++++++- 1 file changed, 118 insertions(+), 4 deletions(-) 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 7a56192e..8f59ecdd 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 @@ -380,6 +380,18 @@ 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. + **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 @@ -392,8 +404,17 @@ class GLiNERExtractor(EntityExtractor): 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. @@ -419,6 +440,8 @@ def __init__( self, threshold: float | None = None, model_name: str | None = None, + window_tokens: int | None = None, + window_overlap: int = 48, ) -> None: self._model_name = model_name or self.DEFAULT_MODEL if threshold is None: @@ -441,6 +464,9 @@ def __init__( # ``_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: @@ -485,9 +511,63 @@ def _get_shared_model(cls, model_name: str) -> Any: 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] + window = self._resolve_window(model) # No lock here, deliberately. # @@ -508,17 +588,51 @@ def _predict_sync(self, text: str, entity_types: list[str]) -> list[dict[str, An # 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) + return self._predict_body(model, text, labels, window) def _predict_body( self, model: Any, text: str, labels: list[str], + window: int, ) -> list[dict[str, Any]]: - """Inference. Split out from :meth:`_predict_sync` so the lock removal - above could be A/B tested by wrapping one call site.""" - return model.predict_entities(text, labels, threshold=self._threshold) + """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=self._threshold) + + 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=self._threshold + ): + 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, From 91030443acddcb99d558710bf482414e55b4d78a Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:09:46 +0300 Subject: [PATCH 08/19] feat(extraction): opt-in candidate_threshold to keep low-confidence entities Previously the model was queried at `threshold`, so anything less confident was discarded inside GLiNER and never reached the SDK. Setting candidate_threshold below threshold instead keeps those entities and labels them "Unknown". The rest of the pipeline already supports 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, and measured that way deliberately. Lowering threshold outright raised entity recall 32% but dropped entity F1 0.568 -> 0.474 and triple F1 0.236 -> 0.211: low-confidence predictions are mostly noise, and should be marked rather than trusted. Keeping them behind an explicit opt-in means the default path is unchanged. A candidate_threshold above threshold is rejected at construction rather than silently discarding the entities it is meant to preserve. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 38 ++++++++++++++++--- 1 file changed, 33 insertions(+), 5 deletions(-) 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 8f59ecdd..41bebdf1 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 @@ -392,6 +392,18 @@ class GLiNERExtractor(EntityExtractor): 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 @@ -442,6 +454,7 @@ def __init__( 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: @@ -457,6 +470,13 @@ def __init__( self._FALLBACK_THRESHOLD, ) self._threshold = threshold + 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 @@ -568,6 +588,15 @@ def _predict_sync(self, text: str, entity_types: list[str]) -> list[dict[str, An model = self._load_model() labels = [t.lower() for t in entity_types] 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. # @@ -588,7 +617,7 @@ def _predict_sync(self, text: str, entity_types: list[str]) -> list[dict[str, An # 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) + return self._predict_body(model, text, labels, window, floor) def _predict_body( self, @@ -596,6 +625,7 @@ def _predict_body( 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.""" @@ -603,7 +633,7 @@ def _predict_body( # Fast path: fits in one window, identical to unwindowed behaviour. if len(words) <= window: - return model.predict_entities(text, labels, threshold=self._threshold) + return model.predict_entities(text, labels, threshold=floor) step = max(1, window - self._window_overlap) out: list[dict[str, Any]] = [] @@ -612,9 +642,7 @@ def _predict_body( if not span: break lo, hi = span[0][1], span[-1][2] - for p in model.predict_entities( - text[lo:hi], labels, threshold=self._threshold - ): + for p in model.predict_entities(text[lo:hi], labels, threshold=floor): p = dict(p) p["start"] += lo p["end"] += lo From 69c35ace214d2f653d34ff95ed92bc35b502ee51 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:10:24 +0300 Subject: [PATCH 09/19] feat(extraction): add SpacyExtractor as a second entity extractor A non-transformer alternative to GLiNER for callers who need lower latency or cannot ship a transformer model. spaCy's statistical NER runs in a few seconds where GLiNER takes tens of seconds. Included because it is a genuinely different error profile, not because it is better: on the benchmark corpus GLiNER retains the higher ceiling recall, which is why the default is unchanged. Having a second extractor also makes CompositeExtractor meaningful. spaCy's own label set (PERSON, ORG, GPE, DATE, ...) is mapped onto SDK entity types rather than leaked through, so downstream ontology filtering and resolution see the same vocabulary regardless of which extractor produced the entity. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) 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 41bebdf1..46e89a10 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,6 +12,7 @@ import re import threading from abc import ABC, abstractmethod +from collections.abc import Iterable, Sequence from typing import Any, ClassVar from graphrag_sdk.core.models import ExtractedEntity @@ -757,3 +758,158 @@ 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) + + From aeb18ae3cbc1e37a2ec0126df6da20714c375653 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:10:43 +0300 Subject: [PATCH 10/19] feat(extraction): add CompositeExtractor to combine several extractors Runs several entity extractors over the same text and merges the results, resolving overlapping spans rather than emitting both. Two extractors that disagree about the boundary of the same mention would otherwise produce two nodes for one entity, which deduplication then has to clean up. Built and exported but deliberately not the default: on the benchmark corpus combining extractors did not move the end-to-end score, so making it the default would spend a second model's latency and memory for no measured gain. It is here for callers whose corpus has a different profile, and so the option can be measured rather than argued about. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) 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 46e89a10..cc8f89b1 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 @@ -913,3 +913,149 @@ async def extract_entities( 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()) From cfe15e4a0ae1b36e31d5bd98370c39849e42be67 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:12:17 +0300 Subject: [PATCH 11/19] fix(core,extraction): reject reserved node labels as entity types `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. 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 and the symptom surfaces far from the cause. entity_types is now validated at construction, where the caller can act on it, and the check is casefolded so "document" is caught too. The reserved set moves to core.models.RESERVED_NODE_LABELS so GraphStore._STRUCTURAL_LABELS and the extractor cannot drift apart - previously each carried its own literal. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- graphrag_sdk/src/graphrag_sdk/core/models.py | 8 ++++ .../extraction_strategies/graph_extraction.py | 42 ++++++++++++++++++- .../src/graphrag_sdk/storage/graph_store.py | 3 +- 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/core/models.py b/graphrag_sdk/src/graphrag_sdk/core/models.py index 382fd760..741bbe7c 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.""" 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..0b2eb0ff 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 @@ -22,6 +22,7 @@ GraphRelationship, Ontology, Relation, + RESERVED_NODE_LABELS, TextChunks, ) from graphrag_sdk.core.providers import LLMInterface @@ -336,6 +337,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. @@ -431,7 +464,9 @@ def __init__( 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) + ) self._max_concurrency = max_concurrency async def extract( @@ -442,7 +477,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 } diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index c2cef098..bacae54b 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -19,6 +19,7 @@ DocumentRecord, GraphNode, GraphRelationship, + RESERVED_NODE_LABELS, ) from graphrag_sdk.utils.cypher import sanitize_cypher_label @@ -52,7 +53,7 @@ def __init__(self, connection: FalkorDBConnection) -> None: # ── 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"), From 8d46a508d3681aa62d3413c7bd26207b6669fad9 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:13:37 +0300 Subject: [PATCH 12/19] feat(extraction): ship DEFAULT_RELATION_TYPES, the missing relation vocabulary 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 were emitted zero times - `contains` (123 triples) and `authored` (54). Supplying any fixed list doubles exact triple F1 (0.0725 -> 0.1490, +2.06x). A list written without reference to the gold vocabulary scored as well as the gold vocabulary itself (0.1490 vs 0.1456), so the gain comes from being consistent, not from guessing the right words. That is exactly what makes a shipped default worth having. The list is domain-neutral and pairs with DEFAULT_ENTITY_TYPES. Like entity_types it is guidance, not a filter: it steers the prompt, and a relation labelled outside it is still kept. The hard-filtering path remains Ontology.relations, which prunes non-conforming edges in IngestionPipeline._prune. A declared ontology still wins. relation_types=[] stays meaningful and restores the previous open-vocabulary behaviour verbatim, for anyone who wants the old free-for-all. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../extraction_strategies/graph_extraction.py | 88 ++++++++++++++++++- 1 file changed, 86 insertions(+), 2 deletions(-) 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 0b2eb0ff..cc6151c5 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 @@ -43,6 +43,70 @@ 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", +] + VERIFY_EXTRACT_RELS_PROMPT = ( "You are an expert knowledge graph builder.\n" "Given the text and pre-extracted entities below, do two things:\n" @@ -449,6 +513,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. """ @@ -459,11 +530,15 @@ 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, ) -> None: self.llm = llm self.entity_extractor = entity_extractor or GLiNERExtractor() self.coref_resolver = coref_resolver + self.relation_types = ( + list(DEFAULT_RELATION_TYPES) if relation_types is None else list(relation_types) + ) self.entity_types = _reject_reserved_labels( entity_types or list(DEFAULT_ENTITY_TYPES) ) @@ -488,6 +563,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__}, " @@ -590,9 +674,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, From a9ef1ec2356db3c2187c543b312748ab72d902e5 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:13:58 +0300 Subject: [PATCH 13/19] fix(extraction): stop the relation extractor cutting its own output short The relationship step stopped early because it decided it had said enough. Measured (RESULTS.md P5.6): "Extract ALL factual connections" is not enough on its own - the model treats the task as a summary, returns roughly 12 relations per chunk and stops while using 2.5k of a 16k reply budget. Doubling the input text grew the reply by 1.3%, which is the signature of a self-imposed ceiling rather than a token limit. The prompt now states explicitly that this is an exhaustive extraction task and not a summary, that there is no maximum, that a dense paragraph often yields 20 or more relationships, and that a long list is correct rather than a mistake. The entity-removal instruction is also condensed. The specific examples it listed (+=, ->, sh, cd, dt) are now enforced in code by is_valid_entity_name, so spelling them out in the prompt spent tokens re-stating a rule that no longer depends on the model complying. A comment is added above the prompt recording that removing the entity verification job entirely 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). That instruction is doing real work and should not be deleted again without re-running the measurement. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../extraction_strategies/graph_extraction.py | 38 ++++++++++++++----- 1 file changed, 28 insertions(+), 10 deletions(-) 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 cc6151c5..d484bbaa 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 @@ -107,6 +107,12 @@ "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" @@ -123,18 +129,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 " @@ -536,12 +552,14 @@ def __init__( self.llm = llm self.entity_extractor = entity_extractor or GLiNERExtractor() self.coref_resolver = coref_resolver - self.relation_types = ( - list(DEFAULT_RELATION_TYPES) if relation_types is None else list(relation_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 async def extract( From eb13a366c3e6909caef21f079d091fe3552584b0 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:15:08 +0300 Subject: [PATCH 14/19] feat(extraction): add an optional entity verifier, shipped off by default A dedicated second LLM pass that re-reads each extracted entity in context and returns keep/drop plus a corrected type. Fully built: batched prompts, a context window around each mention, and quote-checking so a broken or hallucinating judge cannot silently empty your graph. Shipped disabled, because it was measured and it does not work. On the 11-document benchmark it 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 -> 0.5764. Entity F1 fell 0.609 -> 0.561, lax triple F1 0.248 -> 0.213, and ingest went 124.5s -> 224.6s (+80%, because it is a barrier between two otherwise-overlapping phases, not because of the ~7 extra calls). The interesting part is why it fails. It 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 that needs answering is salience: should this be a node in this graph. That depends on the schema and the expected queries, not on the text, so no amount of re-reading the text can answer it. Kept in the tree rather than deleted so the negative result stays reproducible and the next person does not rebuild it. The docstring carries the numbers and a warning not to enable it without re-measuring. Note this is a different thing from the entity-verification instruction inside VERIFY_EXTRACT_RELS_PROMPT, which was separately measured as load-bearing and must stay. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../extraction_strategies/graph_extraction.py | 244 ++++++++++++++++++ 1 file changed, 244 insertions(+) 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 d484bbaa..0ea13c74 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 @@ -31,6 +31,7 @@ from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( DEFAULT_ENTITY_TYPES, NER_PROMPT, + UNKNOWN_LABEL, EntityExtractor, GLiNERExtractor, LLMExtractor, @@ -163,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": "...", ' @@ -548,6 +600,9 @@ def __init__( 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() @@ -562,6 +617,184 @@ def __init__( ) 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' + f" 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, chunks: TextChunks, @@ -681,6 +914,17 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: 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 From c1196933dca3e916d22cfcfc6887f1e18395081c Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:15:28 +0300 Subject: [PATCH 15/19] feat(core,extraction): report which chunks failed instead of silently shipping half a graph A per-chunk extraction failure was invisible. Failures are swallowed and the chunk contributes nothing, so a document where every call failed returned 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 quietly shipped a half-empty graph that looked like a clean upload. GraphData now carries chunks_attempted and failed_chunks, plus an extraction_failed property for the total-failure case. 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 a caller can retry just those rather than re-ingesting the whole document - the same convention BackfillExecutor already uses. A chunk counts as failed if either phase lost it: step 1 entity extraction, or step 2 relationship extraction. The second case matters because the chunk's entities survive and only its edges are missing, which is exactly the kind of partial loss that used to be undetectable. When any chunk fails, ingest now logs a WARNING naming the count. An INFO summary saying "0 nodes" reads as an empty document; this says plainly that chunks were lost. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- graphrag_sdk/src/graphrag_sdk/core/models.py | 28 ++++++++++++++++++- .../extraction_strategies/graph_extraction.py | 28 +++++++++++++++++-- 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/core/models.py b/graphrag_sdk/src/graphrag_sdk/core/models.py index 741bbe7c..42935403 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/models.py +++ b/graphrag_sdk/src/graphrag_sdk/core/models.py @@ -563,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/graph_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py index 0ea13c74..076e4e18 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 @@ -616,7 +616,6 @@ def __init__( 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)) @@ -840,6 +839,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: @@ -910,6 +913,7 @@ 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) @@ -938,7 +942,9 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: entity_types=_format_entity_types(entity_types, entity_type_descs), relation_patterns=_format_relation_patterns(prompt_relations), attribute_block=_render_attribute_block(ontology), - relationship_type_instruction=_relationship_type_instruction(prompt_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, @@ -952,6 +958,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) @@ -964,6 +971,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 @@ -985,6 +996,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 ── @@ -1013,12 +1027,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 From df94c047fa9748bd277c5e52a287a015ec6d872f Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:15:45 +0300 Subject: [PATCH 16/19] perf(storage): create an id index per label so writes stop degrading as the graph grows Every write path addresses nodes by {id: ...} - MERGE in upsert_nodes and a double MATCH in 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. This is the real reason bulk ingest "gets slower as the graph grows". It is not the LLM. 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 index is created lazily on first write to a label and memoised per store instance, since creation is idempotent in FalkorDB but still a round trip. Per-instance is the right scope: a new store may mean a different graph. Entity nodes also index __Entity__, and relationship writes index both endpoint labels, because those are the labels their queries actually match on. Failures are logged and swallowed. A missing index is a performance problem; raising here would turn it into a correctness problem, and the caller's write must not depend on the optimisation succeeding. Refs FalkorDB/research#71 bug #9. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/graphrag_sdk/storage/graph_store.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index bacae54b..d1c1d960 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -49,6 +49,7 @@ class GraphStore: def __init__(self, connection: FalkorDBConnection) -> None: self._conn = connection + self._indexed_labels: set[str] = set() # ── Write Operations ───────────────────────────────────────── @@ -61,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. @@ -100,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] @@ -190,6 +233,11 @@ 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 = [ From 47f97b929355663891639e1217f5b460ed526950 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:15:53 +0300 Subject: [PATCH 17/19] feat(extraction): export the new extractors and the relation vocabulary SpacyExtractor, CompositeExtractor and DEFAULT_RELATION_TYPES are part of the public surface: users need to construct the alternative extractors, and need to read or extend the default relation vocabulary rather than re-type it. Added to both the extraction_strategies package and the top-level graphrag_sdk namespace, including __all__, so they are importable the same way as the existing extractors. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- graphrag_sdk/src/graphrag_sdk/__init__.py | 6 ++++++ .../ingestion/extraction_strategies/__init__.py | 6 ++++++ 2 files changed, 12 insertions(+) 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/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", ] From a50616fa9831a905fd2b98286c633a95e9dca87a Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:16:03 +0300 Subject: [PATCH 18/19] test(extraction): cover the extraction and storage changes Adds tests for the behaviour changed in this branch: date and junk-name rejection, acronym handling, per-model threshold resolution, the shared model cache, GLiNER windowing and span re-offsetting, SpacyExtractor and CompositeExtractor merging, reserved-label rejection on both the entity_types and ontology paths, failed-chunk reporting, and per-label id index creation including its memoisation and its swallow-on-failure behaviour. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- graphrag_sdk/tests/test_entity_extractors.py | 59 ++++ graphrag_sdk/tests/test_graph_extraction.py | 272 ++++++++++++++++++- graphrag_sdk/tests/test_graph_store.py | 116 +++++++- 3 files changed, 431 insertions(+), 16 deletions(-) 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 From 3ff3d86fcc8e86162f02b352868cc623f82f4fa0 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 25 Aug 2026 19:25:14 +0300 Subject: [PATCH 19/19] style: sort imports and apply ruff format CI runs 'ruff check src/' (with the I ruleset) and 'ruff format --check src/'. Two I001s: RESERVED_NODE_LABELS and DEFAULT_RELATION_TYPES sort before the CamelCase names they were appended after. Formatting follows from the same line-length 100 config. No behaviour change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../entity_extractors.py | 158 ++++++++++-------- .../extraction_strategies/graph_extraction.py | 27 ++- .../src/graphrag_sdk/storage/graph_store.py | 6 +- 3 files changed, 101 insertions(+), 90 deletions(-) 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 cc8f89b1..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 @@ -74,53 +74,87 @@ # 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", -} +_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 ───────────────────────────────────── @@ -161,9 +195,7 @@ def _normalize_type_label(raw: str) -> str: ) # 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 -) +_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: @@ -459,9 +491,7 @@ def __init__( ) -> 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 - ) + 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 " @@ -525,8 +555,7 @@ def _get_shared_model(cls, model_name: str) -> Any: from gliner import GLiNER except ImportError: raise ImportError( - "GLiNER is required for GLiNERExtractor. " - "Install with: pip install gliner" + "GLiNER is required for GLiNERExtractor. Install with: pip install gliner" ) cached = GLiNER.from_pretrained(model_name) cls._MODEL_CACHE[model_name] = cached @@ -544,9 +573,7 @@ def _resolve_window(self, model: Any) -> int: 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 - ) + splitter = getattr(getattr(model, "data_processor", None), "words_splitter", None) if splitter is None: from gliner.data_processing import WordsSplitter @@ -593,11 +620,7 @@ def _predict_sync(self, text: str, entity_types: list[str]) -> list[dict[str, An # ``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 - ) + floor = self._threshold if self._candidate_threshold is None else self._candidate_threshold # No lock here, deliberately. # @@ -1006,10 +1029,7 @@ async def extract_entities( source_chunk_id: str, ) -> list[ExtractedEntity]: results = await asyncio.gather( - *( - e.extract_entities(text, entity_types, source_chunk_id) - for e in self._extractors - ), + *(e.extract_entities(text, entity_types, source_chunk_id) for e in self._extractors), return_exceptions=True, ) @@ -1036,9 +1056,7 @@ async def extract_entities( if ( self._suppress_overlaps and spans - and any( - s < ce and cs < e for (s, e) in spans for (cs, ce) in claimed - ) + 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 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 076e4e18..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, @@ -22,7 +23,6 @@ GraphRelationship, Ontology, Relation, - RESERVED_NODE_LABELS, TextChunks, ) from graphrag_sdk.core.providers import LLMInterface @@ -190,21 +190,21 @@ "## Candidates\n" "{candidates}\n\n" "## Rules\n" - "Mark an entity \"drop\" if ANY of these is true:\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" + '"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" + '(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. " + '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 " + '- 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 " @@ -607,9 +607,7 @@ def __init__( self.llm = llm self.entity_extractor = entity_extractor or GLiNERExtractor() self.coref_resolver = coref_resolver - self.entity_types = _reject_reserved_labels( - 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 = ( @@ -714,8 +712,7 @@ async def _verify_entities( for n, key in enumerate(batch, 1): ent, context = first[key] lines.append( - f'{n}. name: "{ent.name}" | proposed type: {ent.type}\n' - f" context: {context!r}" + f'{n}. name: "{ent.name}" | proposed type: {ent.type}\n context: {context!r}' ) prompts.append( VERIFY_ENTITIES_PROMPT.format( @@ -942,9 +939,7 @@ async def _step1(text: str, chunk_uid: str) -> list[ExtractedEntity]: entity_types=_format_entity_types(entity_types, entity_type_descs), relation_patterns=_format_relation_patterns(prompt_relations), attribute_block=_render_attribute_block(ontology), - relationship_type_instruction=_relationship_type_instruction( - prompt_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, diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index d1c1d960..4a445f73 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -14,12 +14,12 @@ 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, GraphNode, GraphRelationship, - RESERVED_NODE_LABELS, ) from graphrag_sdk.utils.cypher import sanitize_cypher_label @@ -233,9 +233,7 @@ 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__") - ) + 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):