From cb24885bf79797e9f94f7cd0542fd2ae6e4a4806 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:57:39 +0300 Subject: [PATCH 01/21] feat(models): add ChunkEntityRow and ChunkRelationshipRow Typed rows for graph-backed cache reads: one entity per (chunk, entity) mention pair and one RELATES edge per (chunk, edge) provenance pair. Consumed by the upcoming chunk-level extraction cache. No behavior change. --- graphrag_sdk/src/graphrag_sdk/core/models.py | 38 ++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/graphrag_sdk/src/graphrag_sdk/core/models.py b/graphrag_sdk/src/graphrag_sdk/core/models.py index a0199334..382fd760 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/models.py +++ b/graphrag_sdk/src/graphrag_sdk/core/models.py @@ -133,6 +133,44 @@ class DocumentRecord(DataModel): content_hash: str | None = None +class ChunkEntityRow(DataModel): + """One entity mentioned by a specific chunk, as read back from the graph. + + Returned by ``GraphStore.get_entities_mentioned_in_chunks()``. Consumed + by ``CachedChunkExtraction`` to rebuild ``GraphData`` for unchanged + chunks without re-running LLM extraction. + + ``label`` is the entity's concrete label (the non-``__Entity__`` one); + ``None`` if the node somehow carries no usable label (tampered graph). + """ + + chunk_id: str + entity_id: str + label: str | None = None + name: str | None = None + type: str | None = None + description: str | None = None + source_chunk_ids: list[str] = Field(default_factory=list) + + +class ChunkRelationshipRow(DataModel): + """One RELATES edge whose provenance includes a specific chunk. + + Returned by ``GraphStore.get_relationships_for_chunks()``. Consumed by + ``CachedChunkExtraction`` to re-emit relationship facts for unchanged + chunks during a cached document update. + """ + + chunk_id: str + start_entity_id: str + end_entity_id: str + rel_type: str | None = None + description: str | None = None + fact: str | None = None + src_name: str | None = None + tgt_name: str | None = None + + # ── Schema Types ───────────────────────────────────────────────── From 5cfd34e4b785a78dcdcd62949c80b039201a8336 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:57:39 +0300 Subject: [PATCH 02/21] feat(storage): add chunk-level cache read accessors to GraphStore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three schema-owning read methods backing the chunk-level extraction cache: - get_document_chunk_texts(document_id) — (chunk_id, text) snapshot of a document's live chunks, for hashing before the update cutover. - get_entities_mentioned_in_chunks(chunk_ids) — previously extracted entities per chunk. No e.type label fallback: a MERGE on a label the node doesn't carry would mint a duplicate node, so label-less rows return None and the consumer skips the node write. - get_relationships_for_chunks(chunk_ids) — RELATES edges whose source_chunk_ids provenance includes the chunks. Single edge scan per batch (any(c IN r.source_chunk_ids WHERE c IN $cids)) rather than a per-chunk re-scan; chunk intersection is client-side. All batched by _BATCH_SIZE with parameterized Cypher. --- .../src/graphrag_sdk/storage/graph_store.py | 131 +++++++++++++++++- 1 file changed, 130 insertions(+), 1 deletion(-) diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index 457ce2b8..a425d9b5 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -13,7 +13,13 @@ from graphrag_sdk.core.connection import FalkorDBConnection from graphrag_sdk.core.exceptions import DatabaseError -from graphrag_sdk.core.models import DocumentRecord, GraphNode, GraphRelationship +from graphrag_sdk.core.models import ( + ChunkEntityRow, + ChunkRelationshipRow, + DocumentRecord, + GraphNode, + GraphRelationship, +) from graphrag_sdk.utils.cypher import sanitize_cypher_label logger = logging.getLogger(__name__) @@ -680,6 +686,129 @@ async def get_document_chunk_ids(self, document_id: str) -> list[str]: ) return [row[0] for row in result.result_set] if result.result_set else [] + async def get_document_chunk_texts(self, document_id: str) -> list[tuple[str, str]]: + """Snapshot ``(chunk_id, text)`` for every chunk of a document. + + Used by ``CachedChunkExtraction`` to hash existing chunk texts + before an update's cutover deletes them. Rows with a missing id + or non-string text (tampered/partial graphs) are skipped rather + than propagated — the caller treats absent rows as cache misses, + which degrades to full extraction, never to data loss. + + Like the other pre-update snapshots, this returns the full set in + one round trip with no LIMIT — documents with millions of chunks + are out of scope for incremental update. + """ + result = await self._conn.query( + "MATCH (:Document {id: $id})-[:PART_OF]->(c:Chunk) " + "RETURN c.id AS cid, c.text AS text", + {"id": document_id}, + ) + out: list[tuple[str, str]] = [] + for row in result.result_set or []: + cid, text = row[0], row[1] + if cid and isinstance(text, str): + out.append((cid, text)) + return out + + async def get_entities_mentioned_in_chunks( + self, chunk_ids: list[str] + ) -> list[ChunkEntityRow]: + """Read every entity with a ``MENTIONED_IN`` edge to any of + ``chunk_ids``, one row per (chunk, entity) pair. + + Feeds ``CachedChunkExtraction``: for unchanged chunks, the graph + itself is the extraction cache, and these rows are what was + previously extracted. ``label`` is the first non-``__Entity__`` + label (matching how ``upsert_nodes`` writes concrete label + + ``__Entity__`` marker). ``None`` when no concrete label exists — + deliberately NO fallback to ``e.type``: a MERGE on a label the + node doesn't carry would mint a duplicate node with the same id, + so the consumer must skip the node write instead + (``CachedChunkExtraction`` keeps the mention, which is enough to + survive orphan cleanup). + + Batched by ``_BATCH_SIZE`` to keep parameter sizes bounded. + """ + rows: list[ChunkEntityRow] = [] + for start in range(0, len(chunk_ids), self._BATCH_SIZE): + batch = chunk_ids[start : start + self._BATCH_SIZE] + result = await self._conn.query( + "UNWIND $cids AS cid " + "MATCH (e:__Entity__)-[:MENTIONED_IN]->(c:Chunk {id: cid}) " + "RETURN cid, e.id, labels(e), e.name, e.type, e.description, " + "e.source_chunk_ids", + {"cids": batch}, + ) + for row in result.result_set or []: + cid, eid, labels, name, etype, description, source_chunk_ids = row + if not cid or not eid: + continue + label = next((lb for lb in (labels or []) if lb != "__Entity__"), None) + rows.append( + ChunkEntityRow( + chunk_id=cid, + entity_id=eid, + label=label, + name=name if isinstance(name, str) else None, + type=etype if isinstance(etype, str) else None, + description=description if isinstance(description, str) else None, + source_chunk_ids=[ + s for s in (source_chunk_ids or []) if isinstance(s, str) + ], + ) + ) + return rows + + async def get_relationships_for_chunks( + self, chunk_ids: list[str] + ) -> list[ChunkRelationshipRow]: + """Read every ``RELATES`` edge whose ``source_chunk_ids`` + provenance includes any of ``chunk_ids``, one row per + (chunk, edge) pair. + + Counterpart of ``get_entities_mentioned_in_chunks`` for edges; + feeds ``CachedChunkExtraction``'s rebuild of relationship facts + for unchanged chunks. Batched by ``_BATCH_SIZE``. + + Query shape: ``source_chunk_ids`` is a plain list property, so no + index can serve the membership test — but the edge scan runs ONCE + per batch (``any(c IN r.source_chunk_ids WHERE c IN $cids)``), not + once per chunk id. The chunk-id intersection is computed client-side + from the returned provenance list. + """ + rows: list[ChunkRelationshipRow] = [] + for start in range(0, len(chunk_ids), self._BATCH_SIZE): + batch = chunk_ids[start : start + self._BATCH_SIZE] + batch_set = set(batch) + result = await self._conn.query( + "MATCH (a:__Entity__)-[r:RELATES]->(b:__Entity__) " + "WHERE any(c IN r.source_chunk_ids WHERE c IN $cids) " + "RETURN a.id, b.id, r.rel_type, r.description, r.fact, " + "r.src_name, r.tgt_name, r.source_chunk_ids", + {"cids": batch}, + ) + for row in result.result_set or []: + start_id, end_id, rel_type, description, fact, src_name, tgt_name, srcs = row + if not start_id or not end_id: + continue + for cid in srcs or []: + if cid not in batch_set: + continue + rows.append( + ChunkRelationshipRow( + chunk_id=cid, + start_entity_id=start_id, + end_entity_id=end_id, + rel_type=rel_type if isinstance(rel_type, str) else None, + description=description if isinstance(description, str) else None, + fact=fact if isinstance(fact, str) else None, + src_name=src_name if isinstance(src_name, str) else None, + tgt_name=tgt_name if isinstance(tgt_name, str) else None, + ) + ) + return rows + async def set_pending_cleanup_state( self, document_id: str, From dc29560da845962864a253ad163423737022f26f Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:58:00 +0300 Subject: [PATCH 03/21] feat(ingestion): add CachedChunkExtraction strategy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decorator ExtractionStrategy for document updates: chunks whose text is byte-identical (SHA-256) to a chunk already stored for the same document skip LLM extraction — their entities, relationships, and mentions are rebuilt from the live graph and remapped onto the new chunk uids. Only genuinely new/changed chunks reach the inner extractor. The graph itself is the cache; nothing new is stored. Safe by construction: - deterministic entity ids MERGE onto existing nodes (SET n += preserves embeddings; only non-empty props are emitted) - entity source_chunk_ids remapped: this doc's old ids -> new uid(s) or dropped at cutover; other documents' ids pass through untouched - duplicate identical chunks each get their own mentions and provenance - label-less entities skip the node write (MERGE would mint a duplicate) but keep the mention, surviving orphan cleanup - fail-open: any cache lookup/rebuild error falls back to full extraction — worst case is skippable LLM spend, never data loss cached_chunk_count / extracted_chunk_count expose the split for reporting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cached_chunk_extraction.py | 309 ++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100644 graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py new file mode 100644 index 00000000..1589ee85 --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -0,0 +1,309 @@ +# GraphRAG SDK — Ingestion: Chunk-Level Extraction Cache +# Pattern: Decorator — wraps any ExtractionStrategy with a graph-backed cache. + +"""Chunk-level extraction cache for document updates. + +Wraps any :class:`ExtractionStrategy` so that chunks whose text is +byte-identical to a chunk already stored for the same document skip LLM +extraction entirely: their entities, relationships, and mentions are rebuilt +from the live graph and remapped onto the new chunk uid. Only genuinely +new/changed chunks are passed to the inner extractor. + +The graph itself is the cache — nothing is stored anywhere new. Reads go +through :class:`~graphrag_sdk.storage.graph_store.GraphStore` accessors +(``get_document_chunk_texts``, ``get_entities_mentioned_in_chunks``, +``get_relationships_for_chunks``), so the persistence schema stays owned by +the storage layer. + +Safe by construction: + +- Entity ids are deterministic, so cached nodes MERGE onto the existing + entities (``SET n += props`` preserves embeddings and enrichment). +- Entity ``source_chunk_ids`` are re-emitted with this document's old chunk + ids remapped to the new chunk uids (old ids die in the cutover; other + documents' ids pass through), so citation resolution keeps working. +- Relationship upserts UNION ``source_chunk_ids`` (see + ``GraphStore.upsert_relationships``); the ``update()`` stale-edge cleanup + then strips the old chunk ids. +- Mentions are re-emitted against the new chunk uid, so the ``update()`` + orphan cleanup keeps every entity that unchanged chunks still support. +- Every cache failure falls back to real extraction (fail-open): the worst + case is paying for LLM calls that could have been skipped, never data loss. + +Semantics to be aware of (documented on ``GraphRAG.update()``): + +- The cache mirrors the live graph. If entities were manually deleted from + the graph, an update with caching enabled will NOT resurrect them from + unchanged chunks (full re-extraction would). Pass + ``cache_unchanged_chunks=False`` to force a full rebuild. +- If the ontology, prompts, or model changed since ingest, unchanged chunks + keep their previously extracted data. Disable the cache to re-extract + under the new configuration. +""" + +from __future__ import annotations + +import hashlib +import logging +from typing import TYPE_CHECKING, Any + +from graphrag_sdk.core.context import Context +from graphrag_sdk.core.models import ( + EntityMention, + GraphData, + GraphNode, + GraphRelationship, + Ontology, + TextChunk, + TextChunks, +) +from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy + +if TYPE_CHECKING: + from graphrag_sdk.storage.graph_store import GraphStore + +logger = logging.getLogger(__name__) + + +def _sha256(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +class CachedChunkExtraction(ExtractionStrategy): + """Extraction strategy that reuses graph data for unchanged chunks. + + Decorates a real extraction strategy: chunks whose text hash matches an + existing chunk of ``document_id`` are rebuilt from the graph; the rest + are forwarded to ``inner``. Designed for ``GraphRAG.update()`` — the + facade wires it automatically when ``cache_unchanged_chunks=True``. + + After ``extract()`` returns, ``cached_chunk_count`` and + ``extracted_chunk_count`` report the split so callers can surface cache + effectiveness (``update()`` copies them into ``UpdateResult.metadata``). + + Args: + inner: The real extraction strategy for new/changed chunks. + graph_store: Graph store to read previously extracted data from. + document_id: Id of the Document node being updated. + """ + + def __init__( + self, + inner: ExtractionStrategy, + graph_store: GraphStore, + document_id: str, + ) -> None: + if not document_id or not document_id.strip(): + raise ValueError("'document_id' must be a non-empty string") + self._inner = inner + self._graph_store = graph_store + self._document_id = document_id + # Filled during extract() so the caller can report cache stats. + self.cached_chunk_count = 0 + self.extracted_chunk_count = 0 + + async def _old_chunks_by_hash(self) -> tuple[dict[str, str], set[str]]: + """Return (sha256(chunk text) -> old chunk id, all old chunk ids).""" + by_hash: dict[str, str] = {} + all_ids: set[str] = set() + for cid, text in await self._graph_store.get_document_chunk_texts(self._document_id): + all_ids.add(cid) + # First occurrence wins; duplicates map to the same content anyway. + by_hash.setdefault(_sha256(text), cid) + return by_hash, all_ids + + async def _graph_data_from_cache( + self, pairs: list[tuple[str, str]], all_old_ids: set[str] + ) -> GraphData: + """Rebuild GraphData for all unchanged chunks from the live graph. + + Args: + pairs: ``(old_chunk_id, new_chunk_uid)`` for every cached chunk. + Multiple new chunks may map to the same old chunk (duplicate + text within the document) — every new uid gets its own + mentions and provenance. + all_old_ids: every pre-update chunk id of this document — used to + drop provenance pointing at chunks that die in the cutover. + """ + # old chunk id -> ALL new chunk uids that reuse its extraction. + id_map: dict[str, list[str]] = {} + for old_id, new_uid in pairs: + id_map.setdefault(old_id, []).append(new_uid) + old_ids = list(id_map) + + mentions: list[EntityMention] = [] + ent_props: dict[str, dict[str, Any]] = {} + ent_labels: dict[str, str | None] = {} + ent_sources: dict[str, list[str]] = {} + + for row in await self._graph_store.get_entities_mentioned_in_chunks(old_ids): + new_uids = id_map.get(row.chunk_id) + if not new_uids: + continue + for new_uid in new_uids: + mentions.append(EntityMention(chunk_id=new_uid, entity_id=row.entity_id)) + srcs = ent_sources.setdefault(row.entity_id, []) + if row.entity_id not in ent_props: + # Only emit non-empty fields — SET n += props would null-out + # existing values otherwise. + props: dict[str, Any] = {} + if row.name: + props["name"] = row.name + if row.type: + props["type"] = row.type + if row.description: + props["description"] = row.description + ent_props[row.entity_id] = props + ent_labels[row.entity_id] = row.label + # Remap provenance onto post-update chunk ids: this doc's + # old ids map to their replacement uid(s) or die in the + # cutover; other documents' chunk ids pass through. + for old in row.source_chunk_ids: + if old in all_old_ids: + mapped_uids = id_map.get(old, []) + else: + mapped_uids = [old] + for mapped in mapped_uids: + if mapped not in srcs: + srcs.append(mapped) + for new_uid in new_uids: + if new_uid not in srcs: + srcs.append(new_uid) + + nodes: list[GraphNode] = [] + for eid, props in ent_props.items(): + label = ent_labels[eid] + if not label: + # No usable label: upsert_nodes would MERGE on a fallback + # label and mint a duplicate node instead of matching this + # one. Skip the node write — the mention above still keeps + # the entity alive through orphan cleanup. + logger.warning("Skipping cache node write for label-less entity %s", eid) + continue + props["source_chunk_ids"] = ent_sources[eid] + nodes.append(GraphNode(id=eid, label=label, properties=props)) + + relationships: list[GraphRelationship] = [] + if ent_props: + rel_by_pair: dict[tuple[str, str], dict[str, Any]] = {} + for rel in await self._graph_store.get_relationships_for_chunks(old_ids): + new_uids = id_map.get(rel.chunk_id) + if not new_uids: + continue + key = (rel.start_entity_id, rel.end_entity_id) + props = rel_by_pair.get(key) + if props is None: + props = {"source_chunk_ids": []} + if rel.rel_type: + props["rel_type"] = rel.rel_type + if rel.description: + props["description"] = rel.description + if rel.fact: + props["fact"] = rel.fact + if rel.src_name: + props["src_name"] = rel.src_name + if rel.tgt_name: + props["tgt_name"] = rel.tgt_name + rel_by_pair[key] = props + for new_uid in new_uids: + if new_uid not in props["source_chunk_ids"]: + props["source_chunk_ids"].append(new_uid) + relationships = [ + GraphRelationship(start_node_id=s, end_node_id=e, type="RELATES", properties=p) + for (s, e), p in rel_by_pair.items() + ] + + return GraphData(nodes=nodes, relationships=relationships, mentions=mentions) + + @staticmethod + def _merge(parts: list[GraphData]) -> GraphData: + """Merge cached and freshly extracted GraphData. + + Later parts win on conflicting node properties (fresh extraction + follows cached parts, so updated descriptions take precedence) — + EXCEPT ``source_chunk_ids``, which is a union: an entity present in + both a cached and an extracted chunk must keep both provenances. + """ + nodes_by_id: dict[str, GraphNode] = {} + relationships: list[GraphRelationship] = [] + mentions: list[EntityMention] = [] + extracted_entities = [] + extracted_relations = [] + for part in parts: + for node in part.nodes: + existing = nodes_by_id.get(node.id) + if existing is None: + nodes_by_id[node.id] = node + else: + old_props = existing.properties or {} + new_props = node.properties or {} + merged = {**old_props, **new_props} + old_src = old_props.get("source_chunk_ids") or [] + new_src = new_props.get("source_chunk_ids") or [] + if old_src or new_src: + merged["source_chunk_ids"] = old_src + [ + c for c in new_src if c not in old_src + ] + label = node.label or existing.label + nodes_by_id[node.id] = GraphNode(id=node.id, label=label, properties=merged) + relationships.extend(part.relationships) + mentions.extend(part.mentions) + extracted_entities.extend(part.extracted_entities) + extracted_relations.extend(part.extracted_relations) + return GraphData( + nodes=list(nodes_by_id.values()), + relationships=relationships, + mentions=mentions, + extracted_entities=extracted_entities, + extracted_relations=extracted_relations, + ) + + async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> GraphData: + """Split chunks into cached vs. new, rebuild the former from the + graph, extract the latter with the inner strategy, and merge.""" + try: + old_by_hash, all_old_ids = await self._old_chunks_by_hash() + except Exception as exc: + logger.warning("Chunk cache lookup failed, extracting everything: %s", exc) + old_by_hash, all_old_ids = {}, set() + + cached: list[tuple[TextChunk, str]] = [] # (new chunk, old chunk id) + to_extract: list[TextChunk] = [] + for chunk in chunks.chunks: + old_id = old_by_hash.get(_sha256(chunk.text)) + if old_id is not None: + cached.append((chunk, old_id)) + else: + to_extract.append(chunk) + + self.cached_chunk_count = len(cached) + self.extracted_chunk_count = len(to_extract) + if ctx: + ctx.log( + f"update cache: {len(cached)} unchanged chunk(s) reused, " + f"{len(to_extract)} chunk(s) sent to extraction" + ) + + parts: list[GraphData] = [] + if cached: + try: + parts.append( + await self._graph_data_from_cache( + [(old_id, chunk.uid) for chunk, old_id in cached], all_old_ids + ) + ) + except Exception as exc: + # Cache miss must never lose data — fall back to real extraction. + logger.warning( + "Cache rebuild failed, extracting %d cached chunk(s) instead: %s", + len(cached), + exc, + ) + to_extract.extend(chunk for chunk, _ in cached) + self.extracted_chunk_count += self.cached_chunk_count + self.cached_chunk_count = 0 + + if to_extract: + parts.append(await self._inner.extract(TextChunks(chunks=to_extract), ontology, ctx)) + + return self._merge(parts) From 139594597fc3d6beb611b2f5edda7e97b2088c35 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:58:00 +0300 Subject: [PATCH 04/21] feat(api): wire cache_unchanged_chunks into update/apply_changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New opt-in flag (default False — existing behavior unchanged) on update(), update_sync(), apply_changes(), and apply_changes_sync(). When set, update() wraps the effective extractor in CachedChunkExtraction during Phase 3 — old chunks still exist until the Phase 5 cutover, so the cache read is safe by construction and the crash-safe pending/commit/rollforward state machine is untouched. Cache effectiveness is surfaced in UpdateResult.metadata["cache_stats"] (cached_chunks / extracted_chunks). The if_missing="ingest" fresh-ingest fallthrough is not wrapped (a new document has no chunks to reuse). Caveats documented on the docstring: the graph is the cache, so manual entity deletions are not resurrected from unchanged chunks, and ontology/prompt/model changes do not re-extract unchanged chunks — pass False to force a full rebuild. --- graphrag_sdk/src/graphrag_sdk/api/main.py | 57 ++++++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/api/main.py b/graphrag_sdk/src/graphrag_sdk/api/main.py index c774e96b..20da5409 100644 --- a/graphrag_sdk/src/graphrag_sdk/api/main.py +++ b/graphrag_sdk/src/graphrag_sdk/api/main.py @@ -54,6 +54,9 @@ SentenceTokenCapChunking, ) from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction import ( + CachedChunkExtraction, +) from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( DEFAULT_ENTITY_TYPES, _strip_markdown_fences, @@ -1923,6 +1926,7 @@ async def update( chunker: ChunkingStrategy | None = None, extractor: ExtractionStrategy | None = None, resolver: ResolutionStrategy | None = None, + cache_unchanged_chunks: bool = False, if_missing: Literal["error", "ingest"] = "error", ctx: Context | None = None, ) -> UpdateResult: @@ -1984,6 +1988,23 @@ async def update( with no extra plumbing. Required in text mode. loader / chunker / extractor / resolver: Per-call strategy overrides, identical to ``ingest()``. + cache_unchanged_chunks: When ``True``, wrap the extractor in + :class:`~graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction.CachedChunkExtraction`: + new chunks whose text is byte-identical to an existing + chunk of this document skip LLM extraction — their + entities/relations/mentions are rebuilt from the live + graph and remapped onto the new chunk uids. Only changed + chunks pay for extraction. Cache effectiveness is + reported in ``UpdateResult.metadata["cache_stats"]`` + (``cached_chunks`` / ``extracted_chunks``). Two semantic + caveats: (1) the graph is the cache, so manually deleted + entities are NOT resurrected from unchanged chunks; (2) + if the ontology, prompts, or model changed since ingest, + unchanged chunks keep their previously extracted data. + Leave ``False`` (the default) to force full re-extraction + in those situations. No effect on the + ``if_missing="ingest"`` fresh-ingest fallthrough (a new + document has no chunks to reuse). if_missing: ``"error"`` (default) raises ``DocumentNotFoundError`` when the id is unknown. ``"ingest"`` falls through to ``ingest()`` for upsert semantics. @@ -2109,10 +2130,24 @@ async def update( # Contradictions surface here rather than mid-extraction. await self._ensure_ontology_initialized() + # Chunk-level extraction cache (opt-in). Wraps the effective + # extractor so byte-identical chunks are rebuilt from the live + # graph instead of re-extracted. Must read the OLD chunks, which + # still exist until the Phase 5 cutover — safe by construction. + active_extractor = extractor or self._default_extractor() + cache_wrapper: CachedChunkExtraction | None = None + if cache_unchanged_chunks: + cache_wrapper = CachedChunkExtraction( + inner=active_extractor, + graph_store=self._graph_store, + document_id=resolved_id, + ) + active_extractor = cache_wrapper + pipeline = IngestionPipeline( loader=loader or TextLoader(), # unused (text is provided below) chunker=chunker or SentenceTokenCapChunking(), - extractor=extractor or self._default_extractor(), + extractor=active_extractor, resolver=resolver or ExactMatchResolution(), graph_store=self._graph_store, vector_store=self._vector_store, @@ -2190,12 +2225,19 @@ async def update( f"wrote {pipeline_result.chunks_indexed} new chunks" ) + result_metadata = dict(pipeline_result.metadata) + if cache_wrapper is not None: + result_metadata["cache_stats"] = { + "cached_chunks": cache_wrapper.cached_chunk_count, + "extracted_chunks": cache_wrapper.extracted_chunk_count, + } + return UpdateResult( document_info=DocumentInfo(uid=resolved_id, path=doc_path, metadata=loaded_metadata), nodes_created=pipeline_result.nodes_created, relationships_created=pipeline_result.relationships_created, chunks_indexed=pipeline_result.chunks_indexed, - metadata=pipeline_result.metadata, + metadata=result_metadata, chunks_deleted=chunks_deleted, entities_deleted=entities_deleted, replaced_existing=True, @@ -2324,6 +2366,7 @@ async def apply_changes( chunker: ChunkingStrategy | None = None, extractor: ExtractionStrategy | None = None, resolver: ResolutionStrategy | None = None, + cache_unchanged_chunks: bool = False, max_concurrency: int = 3, update_concurrency: int = 1, ctx: Context | None = None, @@ -2377,6 +2420,11 @@ async def apply_changes( ``added``/``modified``. ``deleted`` ignores this. resolver: Override the resolution strategy for ``added``/ ``modified``. ``deleted`` ignores this. + cache_unchanged_chunks: Forwarded to ``update()`` for the + ``modified`` list — unchanged chunks skip LLM extraction + and are rebuilt from the live graph. See + :meth:`update` for semantics and caveats. ``added`` and + ``deleted`` ignore this. max_concurrency: Parallelism cap for ``ingest()`` of the ``added`` list. Matches ``ingest()``'s own knob and the ``add`` step is pure ingestion with no orphan-cleanup @@ -2469,6 +2517,7 @@ async def _update_one(path: str) -> BatchEntry[UpdateResult]: chunker=chunker, extractor=extractor, resolver=resolver, + cache_unchanged_chunks=cache_unchanged_chunks, if_missing="ingest", ctx=ctx.child(), ) @@ -3138,6 +3187,7 @@ def update_sync( chunker: ChunkingStrategy | None = None, extractor: ExtractionStrategy | None = None, resolver: ResolutionStrategy | None = None, + cache_unchanged_chunks: bool = False, if_missing: Literal["error", "ingest"] = "error", ctx: Context | None = None, ) -> UpdateResult: @@ -3161,6 +3211,7 @@ def update_sync( chunker=chunker, extractor=extractor, resolver=resolver, + cache_unchanged_chunks=cache_unchanged_chunks, if_missing=if_missing, ctx=ctx, ) @@ -3193,6 +3244,7 @@ def apply_changes_sync( chunker: ChunkingStrategy | None = None, extractor: ExtractionStrategy | None = None, resolver: ResolutionStrategy | None = None, + cache_unchanged_chunks: bool = False, max_concurrency: int = 3, update_concurrency: int = 1, ctx: Context | None = None, @@ -3215,6 +3267,7 @@ def apply_changes_sync( chunker=chunker, extractor=extractor, resolver=resolver, + cache_unchanged_chunks=cache_unchanged_chunks, max_concurrency=max_concurrency, update_concurrency=update_concurrency, ctx=ctx, From a90d1acf37cd7abbecec9adcc2dcd6c64f9f14ea Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:58:15 +0300 Subject: [PATCH 05/21] test: unit + FalkorDB integration coverage for chunk cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 22 unit tests (fake graph store + recording inner extractor) covering the full matrix: all-cached / all-new / mixed splits, provenance remap (own old ids remapped, foreign ids pass through, dropped ids die), duplicate identical chunks, label-less entity guard, relationship rebuild + same-pair dedup, fail-open on lookup and rebuild failures, cached+extracted merge (fresh props win, provenance unions), stats reset, and the three GraphStore accessors (row mapping, bad-row filtering, single-scan query shape). 4 integration tests (RUN_INTEGRATION=1, real FalkorDB) with a strict scripted LLM — an LLM call for an unchanged chunk raises, so passing tests are the proof that caching skips extraction: - unchanged chunk skips LLM; cache_stats reported; provenance and mentions point only at live chunks; replaced-chunk entity cleaned - manually deleted entity is not resurrected by a cached update - default off: both chunks re-extracted, no cache_stats - whole-document no_op short-circuit unaffected --- .../tests/test_cached_chunk_extraction.py | 799 ++++++++++++++++++ 1 file changed, 799 insertions(+) create mode 100644 graphrag_sdk/tests/test_cached_chunk_extraction.py diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py new file mode 100644 index 00000000..264d6d91 --- /dev/null +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -0,0 +1,799 @@ +"""Tests for ingestion/extraction_strategies/cached_chunk_extraction.py. + +Unit tests use a fake graph store + recording inner extractor and cover the +full split/remap/merge/fallback matrix. Integration tests (RUN_INTEGRATION=1, +real FalkorDB) prove end-to-end that unchanged chunks skip LLM extraction — +the scripted LLM is strict, so any unexpected extraction raises loudly. +""" + +from __future__ import annotations + +import hashlib +from unittest.mock import AsyncMock + +import pytest + +from graphrag_sdk.core.context import Context +from graphrag_sdk.core.models import ( + ChunkEntityRow, + ChunkRelationshipRow, + EntityMention, + GraphData, + GraphNode, + Ontology, + TextChunk, + TextChunks, +) +from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction import ( + CachedChunkExtraction, +) +from graphrag_sdk.storage.graph_store import GraphStore + + +def _sha(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +class FakeGraphStore: + """Minimal in-memory stand-in for the three cache read accessors.""" + + def __init__( + self, + chunk_texts: list[tuple[str, str]] | None = None, + entity_rows: list[ChunkEntityRow] | None = None, + rel_rows: list[ChunkRelationshipRow] | None = None, + ) -> None: + self.chunk_texts = chunk_texts or [] + self.entity_rows = entity_rows or [] + self.rel_rows = rel_rows or [] + + async def get_document_chunk_texts(self, document_id: str) -> list[tuple[str, str]]: + return self.chunk_texts + + async def get_entities_mentioned_in_chunks( + self, chunk_ids: list[str] + ) -> list[ChunkEntityRow]: + wanted = set(chunk_ids) + return [r for r in self.entity_rows if r.chunk_id in wanted] + + async def get_relationships_for_chunks( + self, chunk_ids: list[str] + ) -> list[ChunkRelationshipRow]: + wanted = set(chunk_ids) + return [r for r in self.rel_rows if r.chunk_id in wanted] + + +class RecordingExtractor(ExtractionStrategy): + """Inner extractor that records every call and returns a canned result.""" + + def __init__(self, result: GraphData | None = None) -> None: + self.calls: list[TextChunks] = [] + self.result = result or GraphData() + + async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> GraphData: + self.calls.append(chunks) + return self.result + + +@pytest.fixture +def ontology() -> Ontology: + return Ontology() + + +@pytest.fixture +def ctx() -> Context: + return Context() + + +def _chunk(text: str, index: int = 0, uid: str | None = None) -> TextChunk: + kwargs = {"text": text, "index": index} + if uid is not None: + kwargs["uid"] = uid + return TextChunk(**kwargs) + + +class TestConstructor: + def test_rejects_empty_document_id(self): + with pytest.raises(ValueError, match="document_id"): + CachedChunkExtraction(RecordingExtractor(), FakeGraphStore(), "") + + def test_rejects_whitespace_document_id(self): + with pytest.raises(ValueError, match="document_id"): + CachedChunkExtraction(RecordingExtractor(), FakeGraphStore(), " ") + + +class TestSplit: + async def test_all_chunks_cached_inner_never_called(self, ontology, ctx): + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha"), ("old-2", "beta")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="alice__person", + label="Person", + name="Alice", + type="Person", + description="Engineer", + source_chunk_ids=["old-1"], + ), + ChunkEntityRow( + chunk_id="old-2", + entity_id="acme__organization", + label="Organization", + name="Acme", + type="Organization", + source_chunk_ids=["old-2"], + ), + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + new_chunks = TextChunks( + chunks=[_chunk("alpha", 0, "new-1"), _chunk("beta", 1, "new-2")] + ) + result = await strategy.extract(new_chunks, ontology, ctx) + + assert inner.calls == [] + assert strategy.cached_chunk_count == 2 + assert strategy.extracted_chunk_count == 0 + assert {m.chunk_id for m in result.mentions} == {"new-1", "new-2"} + assert {m.entity_id for m in result.mentions} == { + "alice__person", + "acme__organization", + } + by_id = {n.id: n for n in result.nodes} + assert by_id["alice__person"].label == "Person" + assert by_id["alice__person"].properties["source_chunk_ids"] == ["new-1"] + assert by_id["alice__person"].properties["description"] == "Engineer" + assert by_id["acme__organization"].properties["source_chunk_ids"] == ["new-2"] + + async def test_all_chunks_new_everything_extracted(self, ontology, ctx): + store = FakeGraphStore(chunk_texts=[("old-1", "something else entirely")]) + node = GraphNode(id="e1", label="Person", properties={"name": "Bob"}) + inner = RecordingExtractor(GraphData(nodes=[node])) + strategy = CachedChunkExtraction(inner, store, "doc-1") + + new_chunks = TextChunks(chunks=[_chunk("alpha"), _chunk("beta", 1)]) + result = await strategy.extract(new_chunks, ontology, ctx) + + assert len(inner.calls) == 1 + assert [c.text for c in inner.calls[0].chunks] == ["alpha", "beta"] + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 2 + assert result.nodes == [node] + + async def test_mixed_split_only_changed_chunk_extracted(self, ontology, ctx): + store = FakeGraphStore( + chunk_texts=[("old-1", "unchanged text")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="e-cached", + label="Person", + name="Cached", + source_chunk_ids=["old-1"], + ) + ], + ) + inner = RecordingExtractor( + GraphData(nodes=[GraphNode(id="e-new", label="Person", properties={})]) + ) + strategy = CachedChunkExtraction(inner, store, "doc-1") + + new_chunks = TextChunks( + chunks=[_chunk("unchanged text", 0, "new-1"), _chunk("brand new", 1, "new-2")] + ) + result = await strategy.extract(new_chunks, ontology, ctx) + + assert len(inner.calls) == 1 + assert [c.text for c in inner.calls[0].chunks] == ["brand new"] + assert strategy.cached_chunk_count == 1 + assert strategy.extracted_chunk_count == 1 + ids = {n.id for n in result.nodes} + assert ids == {"e-cached", "e-new"} + + async def test_empty_input_returns_empty_graph_data(self, ontology, ctx): + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, FakeGraphStore(), "doc-1") + result = await strategy.extract(TextChunks(chunks=[]), ontology, ctx) + assert inner.calls == [] + assert result.nodes == [] and result.mentions == [] and result.relationships == [] + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 0 + + +class TestProvenanceRemap: + async def test_own_old_ids_remap_other_docs_pass_through(self, ontology, ctx): + """source_chunk_ids: this doc's cached old id → new uid; this doc's + dropped old id → removed; other documents' chunk ids → unchanged.""" + store = FakeGraphStore( + chunk_texts=[("old-1", "kept"), ("old-2", "dropped content")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="e1", + label="Person", + name="Alice", + source_chunk_ids=["old-1", "old-2", "other-doc-chunk"], + ) + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + new_chunks = TextChunks( + chunks=[_chunk("kept", 0, "new-1"), _chunk("replacement", 1, "new-2")] + ) + result = await strategy.extract(new_chunks, ontology, ctx) + + node = next(n for n in result.nodes if n.id == "e1") + srcs = node.properties["source_chunk_ids"] + assert "new-1" in srcs # remapped + assert "other-doc-chunk" in srcs # passes through + assert "old-1" not in srcs and "old-2" not in srcs # old ids die + + async def test_duplicate_chunk_texts_all_get_mentions(self, ontology, ctx): + """Two identical new chunks map to one old chunk — every new uid must + get its own mention and appear in provenance (the server-side + original collapsed these; the SDK port must not).""" + store = FakeGraphStore( + chunk_texts=[("old-1", "same text")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="e1", + label="Person", + name="Alice", + source_chunk_ids=["old-1"], + ) + ], + ) + strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + + new_chunks = TextChunks( + chunks=[_chunk("same text", 0, "new-1"), _chunk("same text", 1, "new-2")] + ) + result = await strategy.extract(new_chunks, ontology, ctx) + + assert strategy.cached_chunk_count == 2 + assert {m.chunk_id for m in result.mentions} == {"new-1", "new-2"} + node = next(n for n in result.nodes if n.id == "e1") + assert set(node.properties["source_chunk_ids"]) == {"new-1", "new-2"} + + +class TestRelationshipRebuild: + async def test_relationships_reemitted_with_remapped_provenance(self, ontology, ctx): + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="a", + label="Person", + name="A", + source_chunk_ids=["old-1"], + ), + ChunkEntityRow( + chunk_id="old-1", + entity_id="b", + label="Person", + name="B", + source_chunk_ids=["old-1"], + ), + ], + rel_rows=[ + ChunkRelationshipRow( + chunk_id="old-1", + start_entity_id="a", + end_entity_id="b", + rel_type="KNOWS", + description="A knows B", + fact="(A, KNOWS, B)", + src_name="A", + tgt_name="B", + ) + ], + ) + strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx + ) + + assert len(result.relationships) == 1 + rel = result.relationships[0] + assert rel.type == "RELATES" + assert rel.start_node_id == "a" and rel.end_node_id == "b" + assert rel.properties["rel_type"] == "KNOWS" + assert rel.properties["fact"] == "(A, KNOWS, B)" + assert rel.properties["source_chunk_ids"] == ["new-1"] + + async def test_relationship_between_same_pair_deduped(self, ontology, ctx): + """Two cached chunks supporting the same (a, b) edge produce ONE + relationship whose provenance lists both new uids.""" + rows = [ + ChunkEntityRow( + chunk_id=cid, + entity_id=eid, + label="Person", + name=eid, + source_chunk_ids=[cid], + ) + for cid in ("old-1", "old-2") + for eid in ("a", "b") + ] + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha"), ("old-2", "beta")], + entity_rows=rows, + rel_rows=[ + ChunkRelationshipRow( + chunk_id="old-1", start_entity_id="a", end_entity_id="b", rel_type="KNOWS" + ), + ChunkRelationshipRow( + chunk_id="old-2", start_entity_id="a", end_entity_id="b", rel_type="KNOWS" + ), + ], + ) + strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0, "new-1"), _chunk("beta", 1, "new-2")]), + ontology, + ctx, + ) + + assert len(result.relationships) == 1 + assert set(result.relationships[0].properties["source_chunk_ids"]) == { + "new-1", + "new-2", + } + + +class TestLabelLessEntity: + async def test_label_less_entity_skips_node_but_keeps_mention(self, ontology, ctx): + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="ghost", + label=None, + name="Ghost", + source_chunk_ids=["old-1"], + ) + ], + ) + strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx + ) + + assert result.nodes == [] # no node write — would mint a duplicate + assert result.mentions == [EntityMention(chunk_id="new-1", entity_id="ghost")] + + +class TestFailOpen: + async def test_chunk_lookup_failure_extracts_everything(self, ontology, ctx): + store = FakeGraphStore() + store.get_document_chunk_texts = AsyncMock(side_effect=RuntimeError("boom")) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract( + TextChunks(chunks=[_chunk("alpha"), _chunk("beta", 1)]), ontology, ctx + ) + + assert len(inner.calls) == 1 + assert len(inner.calls[0].chunks) == 2 + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 2 + + async def test_cache_rebuild_failure_falls_back_to_extraction(self, ontology, ctx): + store = FakeGraphStore(chunk_texts=[("old-1", "alpha")]) + store.get_entities_mentioned_in_chunks = AsyncMock(side_effect=RuntimeError("boom")) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0), _chunk("beta", 1)]), ontology, ctx + ) + + # Both the changed chunk AND the failed cached chunk get extracted. + assert len(inner.calls) == 1 + assert sorted(c.text for c in inner.calls[0].chunks) == ["alpha", "beta"] + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 2 + + async def test_stats_reset_between_calls(self, ontology, ctx): + store = FakeGraphStore(chunk_texts=[("old-1", "alpha")]) + strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("alpha")]), ontology, ctx) + assert (strategy.cached_chunk_count, strategy.extracted_chunk_count) == (1, 0) + + await strategy.extract(TextChunks(chunks=[_chunk("other")]), ontology, ctx) + assert (strategy.cached_chunk_count, strategy.extracted_chunk_count) == (0, 1) + + +class TestMerge: + async def test_entity_in_cached_and_extracted_chunk_unions_provenance( + self, ontology, ctx + ): + """Same entity in an unchanged chunk (cached) and a changed chunk + (extracted): fresh properties win, provenance is the union.""" + store = FakeGraphStore( + chunk_texts=[("old-1", "unchanged")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="e1", + label="Person", + name="Alice", + description="Old description", + source_chunk_ids=["old-1"], + ) + ], + ) + fresh = GraphData( + nodes=[ + GraphNode( + id="e1", + label="Person", + properties={ + "name": "Alice", + "description": "New description", + "source_chunk_ids": ["new-2"], + }, + ) + ], + mentions=[EntityMention(chunk_id="new-2", entity_id="e1")], + ) + strategy = CachedChunkExtraction(RecordingExtractor(fresh), store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("unchanged", 0, "new-1"), _chunk("changed", 1, "new-2")]), + ontology, + ctx, + ) + + assert len(result.nodes) == 1 + node = result.nodes[0] + assert node.properties["description"] == "New description" # fresh wins + assert set(node.properties["source_chunk_ids"]) == {"new-1", "new-2"} # union + assert {m.chunk_id for m in result.mentions} == {"new-1", "new-2"} + + async def test_extracted_entities_and_relations_pass_through(self, ontology, ctx): + from graphrag_sdk.core.models import ExtractedEntity + + fresh = GraphData( + extracted_entities=[ExtractedEntity(name="Alice", type="Person")], + ) + strategy = CachedChunkExtraction( + RecordingExtractor(fresh), FakeGraphStore(), "doc-1" + ) + result = await strategy.extract( + TextChunks(chunks=[_chunk("anything")]), ontology, ctx + ) + assert len(result.extracted_entities) == 1 + + +class TestGraphStoreCacheAccessors: + """Unit tests for the three GraphStore read accessors backing the cache.""" + + @pytest.fixture + def graph_store(self, mock_connection): + return GraphStore(mock_connection) + + async def test_get_document_chunk_texts_filters_bad_rows( + self, graph_store, mock_connection + ): + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[ + ["c1", "hello"], + [None, "orphan text"], # missing id → skipped + ["c2", None], # missing text → skipped + ["c3", 42], # non-string text → skipped + ] + ) + ) + rows = await graph_store.get_document_chunk_texts("doc-1") + assert rows == [("c1", "hello")] + cypher = mock_connection.query.call_args[0][0] + assert "PART_OF" in cypher and "c.text" in cypher + + async def test_get_entities_mentioned_in_chunks_maps_rows( + self, graph_store, mock_connection + ): + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[ + [ + "c1", + "alice__person", + ["Person", "__Entity__"], + "Alice", + "Person", + "Engineer", + ["c1", "x"], + ], + ["c1", None, ["Person"], "NoId", "Person", None, []], # skipped + ] + ) + ) + rows = await graph_store.get_entities_mentioned_in_chunks(["c1"]) + assert len(rows) == 1 + row = rows[0] + assert row.entity_id == "alice__person" + assert row.label == "Person" # __Entity__ filtered out + assert row.source_chunk_ids == ["c1", "x"] + + async def test_get_entities_label_none_when_no_concrete_label( + self, graph_store, mock_connection + ): + """No fallback to e.type: a MERGE on a label the node doesn't carry + would mint a duplicate node — the consumer must skip the node write + (label=None triggers CachedChunkExtraction's skip guard).""" + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[["c1", "e1", ["__Entity__"], "X", "Widget", None, None]] + ) + ) + rows = await graph_store.get_entities_mentioned_in_chunks(["c1"]) + assert rows[0].label is None + + async def test_get_relationships_for_chunks_maps_rows( + self, graph_store, mock_connection + ): + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[ + # Edge supported by c1 and an out-of-batch chunk: one row + # per matching cid, foreign provenance ignored. + ["a", "b", "KNOWS", "desc", "fact", "A", "B", ["c1", "other"]], + [None, "b", "KNOWS", None, None, None, None, ["c1"]], # skipped + ] + ) + ) + rows = await graph_store.get_relationships_for_chunks(["c1"]) + assert len(rows) == 1 + assert rows[0].chunk_id == "c1" + assert rows[0].start_entity_id == "a" + assert rows[0].rel_type == "KNOWS" + cypher = mock_connection.query.call_args[0][0] + assert "RELATES" in cypher and "source_chunk_ids" in cypher + # Single edge scan per batch — no per-chunk UNWIND re-scan. + assert "UNWIND" not in cypher + + async def test_get_relationships_one_row_per_matching_chunk( + self, graph_store, mock_connection + ): + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[ + ["a", "b", "KNOWS", None, None, None, None, ["c1", "c2", "foreign"]], + ] + ) + ) + rows = await graph_store.get_relationships_for_chunks(["c1", "c2"]) + assert {(r.chunk_id, r.start_entity_id, r.end_entity_id) for r in rows} == { + ("c1", "a", "b"), + ("c2", "a", "b"), + } + + async def test_empty_chunk_ids_no_query(self, graph_store, mock_connection): + assert await graph_store.get_entities_mentioned_in_chunks([]) == [] + assert await graph_store.get_relationships_for_chunks([]) == [] + mock_connection.query.assert_not_called() + + +# ── Integration: real FalkorDB (RUN_INTEGRATION=1) ─────────────── + + +def _first_n_chars_stable(part1: str, size: int) -> str: + """Sanity helper: ensure part1 is exactly ``size`` chars so FixedSize + chunking yields a byte-identical first chunk across updates.""" + assert len(part1) == size, f"part1 must be exactly {size} chars, got {len(part1)}" + return part1 + + +@pytest.mark.asyncio +@pytest.mark.integration +class TestCachedUpdateIntegration: + """End-to-end proof against real FalkorDB. The scripted LLM is strict: + an LLM call for an unchanged chunk raises, so passing tests ARE the + proof that caching skips extraction.""" + + CHUNK = 64 + + def _texts(self): + # part1 is padded to exactly CHUNK chars → chunk 1 is byte-identical + # across both versions; only chunk 2 changes. FixedSizeChunking cuts + # pure character windows without stripping, so padding is stable. + part1 = _first_n_chars_stable( + "Alice works at Acme Corporation with her colleague Bob today.".ljust( + self.CHUNK, "|" + ), + self.CHUNK, + ) + v1 = part1 + "Carol manages the Berlin office of Acme Corporation." + v2 = part1 + "Dave now manages the Munich office of Acme Corporation." + return v1, v2 + + def _chunker(self): + from graphrag_sdk.ingestion.chunking_strategies.fixed_size import ( + FixedSizeChunking, + ) + + return FixedSizeChunking(chunk_size=self.CHUNK, chunk_overlap=0) + + async def test_cached_update_skips_llm_for_unchanged_chunk( + self, real_falkordb_rag_factory, scripted_llm + ): + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( + ExactMatchResolution, + ) + + v1, v2 = self._texts() + # Scripted responses: 2 for ingest (2 chunks), 1 for update (only + # the changed chunk). strict=True → a 4th call raises. + llm = scripted_llm( + [ + ("Alice", "Person", "Engineer at Acme"), + ("Acme Corporation", "Organization", "Tech company"), + ], + [("Carol", "Person", "Manager in Berlin")], + [("Dave", "Person", "Manager in Munich")], + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=ExactMatchResolution()) + + await rag.ingest(text=v1, document_id="doc-cache", chunker=self._chunker()) + + result = await rag.update( + text=v2, + document_id="doc-cache", + chunker=self._chunker(), + cache_unchanged_chunks=True, + ) + + assert result.no_op is False + assert result.metadata["cache_stats"] == { + "cached_chunks": 1, + "extracted_chunks": 1, + } + + # Entity from the unchanged chunk survives; entity from the replaced + # chunk is orphan-cleaned; new entity present. + async def count(name): + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = $n RETURN count(e)", {"n": name} + ) + return r.result_set[0][0] if r.result_set else 0 + + assert await count("Alice") == 1, "unchanged-chunk entity must survive" + assert await count("Carol") == 0, "replaced-chunk entity must be orphan-cleaned" + assert await count("Dave") == 1, "new-chunk entity must be extracted" + + # Provenance: Alice's source_chunk_ids must point only at LIVE chunks. + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = 'Alice' RETURN e.source_chunk_ids" + ) + alice_srcs = r.result_set[0][0] or [] + r = await rag._graph_store.query_raw( + "MATCH (:Document {id: 'doc-cache'})-[:PART_OF]->(c:Chunk) RETURN collect(c.id)" + ) + live_chunk_ids = set(r.result_set[0][0] or []) + assert alice_srcs, "Alice must retain chunk provenance" + assert set(alice_srcs) <= live_chunk_ids, ( + f"stale provenance survived the update: {set(alice_srcs) - live_chunk_ids}" + ) + + # Alice's MENTIONED_IN edge must target a live chunk. + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__ {name: 'Alice'})-[:MENTIONED_IN]->(c:Chunk) " + "RETURN collect(c.id)" + ) + assert set(r.result_set[0][0] or []) <= live_chunk_ids + + async def test_manually_deleted_entity_not_resurrected( + self, real_falkordb_rag_factory, scripted_llm + ): + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( + ExactMatchResolution, + ) + + v1, v2 = self._texts() + llm = scripted_llm( + [("Alice", "Person", "Engineer"), ("Bob", "Person", "Colleague")], + [("Carol", "Person", "Manager")], + [("Dave", "Person", "Manager")], + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=ExactMatchResolution()) + await rag.ingest(text=v1, document_id="doc-del", chunker=self._chunker()) + + # Curate: manually delete Bob from the graph. + await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = 'Bob' DETACH DELETE e" + ) + + await rag.update( + text=v2, + document_id="doc-del", + chunker=self._chunker(), + cache_unchanged_chunks=True, + ) + + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = 'Bob' RETURN count(e)" + ) + assert r.result_set[0][0] == 0, ( + "cached update must respect manual deletion, not resurrect Bob" + ) + # But Alice (also from the unchanged chunk, not deleted) survives. + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = 'Alice' RETURN count(e)" + ) + assert r.result_set[0][0] == 1 + + async def test_default_off_extracts_all_chunks( + self, real_falkordb_rag_factory, scripted_llm + ): + """cache_unchanged_chunks defaults to False → both chunks hit the + LLM on update (4 scripted responses, all consumed).""" + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( + ExactMatchResolution, + ) + + v1, v2 = self._texts() + llm = scripted_llm( + [("Alice", "Person", "Engineer")], + [("Carol", "Person", "Manager")], + [("Alice", "Person", "Engineer")], + [("Dave", "Person", "Manager")], + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=ExactMatchResolution()) + await rag.ingest(text=v1, document_id="doc-off", chunker=self._chunker()) + + result = await rag.update(text=v2, document_id="doc-off", chunker=self._chunker()) + + assert "cache_stats" not in result.metadata + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.name = 'Dave' RETURN count(e)" + ) + assert r.result_set[0][0] == 1 + + async def test_no_op_short_circuit_unaffected( + self, real_falkordb_rag_factory, scripted_llm + ): + from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( + ExactMatchResolution, + ) + + v1, _ = self._texts() + llm = scripted_llm( + [("Alice", "Person", "Engineer")], + [("Carol", "Person", "Manager")], + ) + rag = real_falkordb_rag_factory(llm=llm, resolver=ExactMatchResolution()) + await rag.ingest(text=v1, document_id="doc-noop", chunker=self._chunker()) + + result = await rag.update( + text=v1, + document_id="doc-noop", + chunker=self._chunker(), + cache_unchanged_chunks=True, + ) + assert result.no_op is True From 7cd0c85bf7cf5892e947f64d25244a098cc8d70b Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:58:16 +0300 Subject: [PATCH 06/21] docs: export CachedChunkExtraction and add CHANGELOG entry Export CachedChunkExtraction at top level, from graphrag_sdk.ingestion, and from graphrag_sdk.ingestion.extraction_strategies, mirroring GraphExtraction. Document the feature, its cache_stats reporting, and its caveats under Unreleased in the CHANGELOG. --- CHANGELOG.md | 32 +++++++++++++++++++ graphrag_sdk/src/graphrag_sdk/__init__.py | 4 +++ .../src/graphrag_sdk/ingestion/__init__.py | 4 +++ .../extraction_strategies/__init__.py | 4 +++ 4 files changed, 44 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ff0a9a56..7bc14b87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,38 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +#### Chunk-level extraction cache for `update()` (`cache_unchanged_chunks`) + +- **`GraphRAG.update(..., cache_unchanged_chunks=True)`** — opt-in + chunk-level extraction cache. New chunks whose text is byte-identical + to an existing chunk of the same document skip LLM extraction + entirely: their entities, relationships, and mentions are rebuilt + from the live graph and remapped onto the new chunk uids. Only + genuinely new/changed chunks are sent to the extractor. Editing one + paragraph of a 50-chunk document now costs ~1 extraction instead + of 50. Cache effectiveness is reported in + `UpdateResult.metadata["cache_stats"]` + (`cached_chunks` / `extracted_chunks`). Also available on + `update_sync()`, `apply_changes()` (applies to the `modified` list), + and `apply_changes_sync()`. Default `False` — existing behavior is + unchanged. +- **`CachedChunkExtraction(inner, graph_store, document_id)`** — the + underlying decorator `ExtractionStrategy`, exported at top level for + advanced pipelines. Fail-open by construction: any cache lookup or + rebuild failure falls back to full extraction (worst case is paying + for skippable LLM calls, never data loss). Caveats: the graph is the + cache, so manually deleted entities are not resurrected from + unchanged chunks, and ontology/prompt/model changes do not + re-extract unchanged chunks — pass `cache_unchanged_chunks=False` + (the default) to force a full rebuild. +- **`GraphStore.get_document_chunk_texts()`**, + **`get_entities_mentioned_in_chunks()`**, + **`get_relationships_for_chunks()`** — new schema-owning read + accessors backing the cache (with `ChunkEntityRow` / + `ChunkRelationshipRow` typed rows in `core.models`). + ## [1.3.0] - 2026-06-04 Ontology discovery (#271): bootstrap an ontology straight from a diff --git a/graphrag_sdk/src/graphrag_sdk/__init__.py b/graphrag_sdk/src/graphrag_sdk/__init__.py index 2f8d8ee4..ef2bf131 100644 --- a/graphrag_sdk/src/graphrag_sdk/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/__init__.py @@ -83,6 +83,9 @@ SentenceTokenCapChunking, ) from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction import ( + CachedChunkExtraction, +) from graphrag_sdk.ingestion.extraction_strategies.coref_resolvers import ( CorefResolver, FastCorefResolver, @@ -185,6 +188,7 @@ "FixedSizeChunking", "SentenceTokenCapChunking", "ExtractionStrategy", + "CachedChunkExtraction", "GraphExtraction", "EntityExtractor", "GLiNERExtractor", diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/__init__.py b/graphrag_sdk/src/graphrag_sdk/ingestion/__init__.py index 4b263d98..84c71652 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/__init__.py @@ -3,12 +3,16 @@ from graphrag_sdk.ingestion.chunking_strategies.base import ChunkingStrategy from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction import ( + CachedChunkExtraction, +) from graphrag_sdk.ingestion.extraction_strategies.graph_extraction import GraphExtraction from graphrag_sdk.ingestion.loaders.base import LoaderStrategy from graphrag_sdk.ingestion.pipeline import IngestionPipeline from graphrag_sdk.ingestion.resolution_strategies.base import ResolutionStrategy __all__ = [ + "CachedChunkExtraction", "ChunkingStrategy", "ExtractionStrategy", "GraphExtraction", 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 1bfa541e..a60956c8 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.py @@ -1,6 +1,9 @@ # GraphRAG SDK — Ingestion: Extraction Strategies from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.cached_chunk_extraction import ( + CachedChunkExtraction, +) from graphrag_sdk.ingestion.extraction_strategies.coref_resolvers import ( CorefResolver, FastCorefResolver, @@ -15,6 +18,7 @@ ) __all__ = [ + "CachedChunkExtraction", "ExtractionStrategy", "GraphExtraction", "EntityExtractor", From 07eae991ca4f4a7e8563b936e4ea62ae8704b770 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:00:13 +0300 Subject: [PATCH 07/21] style: apply ruff format to graph_store cache accessors --- graphrag_sdk/src/graphrag_sdk/storage/graph_store.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index a425d9b5..0445252d 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -700,8 +700,7 @@ async def get_document_chunk_texts(self, document_id: str) -> list[tuple[str, st are out of scope for incremental update. """ result = await self._conn.query( - "MATCH (:Document {id: $id})-[:PART_OF]->(c:Chunk) " - "RETURN c.id AS cid, c.text AS text", + "MATCH (:Document {id: $id})-[:PART_OF]->(c:Chunk) RETURN c.id AS cid, c.text AS text", {"id": document_id}, ) out: list[tuple[str, str]] = [] @@ -711,9 +710,7 @@ async def get_document_chunk_texts(self, document_id: str) -> list[tuple[str, st out.append((cid, text)) return out - async def get_entities_mentioned_in_chunks( - self, chunk_ids: list[str] - ) -> list[ChunkEntityRow]: + async def get_entities_mentioned_in_chunks(self, chunk_ids: list[str]) -> list[ChunkEntityRow]: """Read every entity with a ``MENTIONED_IN`` edge to any of ``chunk_ids``, one row per (chunk, entity) pair. From d6ddcf4d0722106b46da64549a3d43dfc98045ac Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:04:39 +0300 Subject: [PATCH 08/21] fix(storage): guard cache accessors against scalar list properties Tampered/partial graphs may hold a scalar (e.g. a string) where a list is expected in labels(e) or source_chunk_ids. Iterating it would yield characters as bogus labels/provenance instead of treating the row as a cache miss. Treat non-list values as empty (entities) or skip the row (relationships), per the documented skip-bad-rows contract. Addresses Copilot review feedback on PR #288. --- .../src/graphrag_sdk/storage/graph_store.py | 19 ++++++++--- .../tests/test_cached_chunk_extraction.py | 34 +++++++++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index 0445252d..d33916d0 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -741,7 +741,14 @@ async def get_entities_mentioned_in_chunks(self, chunk_ids: list[str]) -> list[C cid, eid, labels, name, etype, description, source_chunk_ids = row if not cid or not eid: continue - label = next((lb for lb in (labels or []) if lb != "__Entity__"), None) + # Tampered/partial graphs may hold a scalar where a list is + # expected — iterating a string would yield characters as + # bogus labels/provenance, so non-lists are treated as empty. + if not isinstance(labels, list): + labels = [] + if not isinstance(source_chunk_ids, list): + source_chunk_ids = [] + label = next((lb for lb in labels if lb != "__Entity__"), None) rows.append( ChunkEntityRow( chunk_id=cid, @@ -750,9 +757,7 @@ async def get_entities_mentioned_in_chunks(self, chunk_ids: list[str]) -> list[C name=name if isinstance(name, str) else None, type=etype if isinstance(etype, str) else None, description=description if isinstance(description, str) else None, - source_chunk_ids=[ - s for s in (source_chunk_ids or []) if isinstance(s, str) - ], + source_chunk_ids=[s for s in source_chunk_ids if isinstance(s, str)], ) ) return rows @@ -789,7 +794,11 @@ async def get_relationships_for_chunks( start_id, end_id, rel_type, description, fact, src_name, tgt_name, srcs = row if not start_id or not end_id: continue - for cid in srcs or []: + # Same scalar-vs-list guard as the entity accessor: a + # tampered string provenance must not iterate as characters. + if not isinstance(srcs, list): + continue + for cid in srcs: if cid not in batch_set: continue rows.append( diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py index 264d6d91..836c6f2c 100644 --- a/graphrag_sdk/tests/test_cached_chunk_extraction.py +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -594,6 +594,40 @@ async def test_get_relationships_one_row_per_matching_chunk( ("c2", "a", "b"), } + async def test_scalar_where_list_expected_treated_as_empty( + self, graph_store, mock_connection + ): + """Tampered graphs may hold a scalar where a list is expected — + it must not iterate as characters into bogus labels/provenance.""" + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[["c1", "e1", "Person", "Alice", "Person", None, "abc"]] + ) + ) + rows = await graph_store.get_entities_mentioned_in_chunks(["c1"]) + assert len(rows) == 1 + assert rows[0].label is None # scalar labels → no concrete label + assert rows[0].source_chunk_ids == [] # "abc" must not become ['a','b','c'] + + async def test_relationship_scalar_provenance_skipped( + self, graph_store, mock_connection + ): + from unittest.mock import MagicMock + + mock_connection.query = AsyncMock( + return_value=MagicMock( + result_set=[ + ["a", "b", "KNOWS", None, None, None, None, "c1"], # scalar → skipped + ["a", "b", "KNOWS", None, None, None, None, ["c1"]], + ] + ) + ) + rows = await graph_store.get_relationships_for_chunks(["c1"]) + assert len(rows) == 1 + assert rows[0].chunk_id == "c1" + async def test_empty_chunk_ids_no_query(self, graph_store, mock_connection): assert await graph_store.get_entities_mentioned_in_chunks([]) == [] assert await graph_store.get_relationships_for_chunks([]) == [] From ba003086c94284a6bae9d1748cab32a45e409754 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:34:49 +0300 Subject: [PATCH 09/21] fix: compare bare model names in graph config validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `EmbedderConfig.to_embedder()` prepends the provider to the model name when `provider` is set (e.g. `azure/text-embedding-3-large`), but the name persisted on the graph's config node is unprefixed. `_validate_graph_config` compared the two raw strings, so simply routing an existing graph through a provider — same model, same dimensions — raised ConfigError and made every subsequent update fail. Normalize both sides before comparing: strip a leading `/` segment when the prefix is a known litellm provider. The provider list is read from litellm at first use and cached, with a static fallback for when that import is unavailable. The check itself is deliberately kept: a genuine model swap silently corrupts vector search, since old and new embeddings stay mutually incomparable without ever raising. --- graphrag_sdk/src/graphrag_sdk/api/main.py | 72 ++++++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) diff --git a/graphrag_sdk/src/graphrag_sdk/api/main.py b/graphrag_sdk/src/graphrag_sdk/api/main.py index 20da5409..cc8bf47a 100644 --- a/graphrag_sdk/src/graphrag_sdk/api/main.py +++ b/graphrag_sdk/src/graphrag_sdk/api/main.py @@ -10,6 +10,7 @@ import logging import os import re +from functools import lru_cache from typing import Any, Literal, overload from uuid import uuid4 @@ -147,6 +148,73 @@ def _neutralize_context_close_tag(text: str) -> str: return _CONTEXT_CLOSE_RE.sub("", text) +# Fallback provider prefixes for :py:func:`_bare_model_name`, used when +# litellm isn't installed (it is an optional extra). Covers the providers +# that actually serve embedding models; the full list is pulled from +# litellm at runtime when available. +_FALLBACK_LITELLM_PROVIDERS = frozenset( + { + "azure", + "azure_ai", + "bedrock", + "cohere", + "databricks", + "deepinfra", + "gemini", + "huggingface", + "mistral", + "nvidia_nim", + "ollama", + "openai", + "openrouter", + "together_ai", + "vertex_ai", + "voyage", + "watsonx", + "xinference", + } +) + + +@lru_cache(maxsize=1) +def _litellm_provider_prefixes() -> frozenset[str]: + """Known litellm routing prefixes, lowercased. + + Prefers litellm's own ``provider_list`` so the set stays correct as + litellm adds providers; falls back to a static set when litellm isn't + installed. Cached — the list is static for the process lifetime. + """ + try: + import litellm + + providers = {str(getattr(p, "value", p)).lower() for p in litellm.provider_list} + if providers: + return frozenset(providers | _FALLBACK_LITELLM_PROVIDERS) + except Exception: # pragma: no cover - depends on optional extra + logger.debug("litellm provider_list unavailable", exc_info=True) + return _FALLBACK_LITELLM_PROVIDERS + + +def _bare_model_name(name: str) -> str: + """Strip litellm's provider routing prefix from a model identifier. + + ``"azure/text-embedding-3-large"`` and ``"text-embedding-3-large"`` name + the *same* model reached through different endpoints — the prefix tells + litellm where to send the request, it is not part of the model identity. + Comparing the raw strings therefore reports a mismatch for a graph that + merely changed provider, even though the vectors are unchanged. + + Only the first segment is removed, and only when it is a recognised + provider, so multi-segment identifiers such as + ``"openrouter/anthropic/claude-3"`` keep the part that distinguishes the + model. Names without a known prefix are returned unchanged. + """ + head, sep, tail = name.partition("/") + if sep and tail and head.lower() in _litellm_provider_prefixes(): + return tail + return name + + def _strip_and_load_json(text: str) -> Any: """Parse LLM-emitted JSON, tolerating optional markdown fences. @@ -2878,7 +2946,9 @@ async def _validate_graph_config(self, *, ctx: Context | None = None) -> None: stored_dim = result.result_set[0][1] current_model = self.embedder.model_name - if stored_model and stored_model != current_model: + if stored_model and _bare_model_name(stored_model) != _bare_model_name( + current_model + ): raise ConfigError( f"Embedding model mismatch: graph was built with " f"'{stored_model}' but current embedder is " From d8e46b88575e937c5f8921ca05fd15990b1ba1f5 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:34:56 +0300 Subject: [PATCH 10/21] test: cover provider-prefix normalization in config validation Twelve tests over the fix in the previous commit. `TestBareModelName` pins the helper's boundaries: known prefixes are stripped, unknown ones are not, and names that merely contain a slash (bare model names with a `/` in them, e.g. HuggingFace repo ids) are left intact. Case is deliberately significant. `TestConfigProviderPrefix` drives the validator end to end, asserting the prod scenario now passes (`azure/X` vs stored `X`) while a real model change still raises. Both mutants were checked: reverting to the raw comparison fails four tests, and stripping any leading path segment fails two. --- graphrag_sdk/tests/test_facade.py | 155 +++++++++++++++++++++++++++++- 1 file changed, 154 insertions(+), 1 deletion(-) diff --git a/graphrag_sdk/tests/test_facade.py b/graphrag_sdk/tests/test_facade.py index 36108731..3e640908 100644 --- a/graphrag_sdk/tests/test_facade.py +++ b/graphrag_sdk/tests/test_facade.py @@ -27,7 +27,7 @@ ) from graphrag_sdk.retrieval.strategies.base import RetrievalStrategy -from .conftest import MockLLM +from .conftest import MockEmbedder, MockLLM # ── Fixtures ──────────────────────────────────────────────────── @@ -1174,6 +1174,159 @@ async def test_no_config_node_passes(self, mock_conn, embedder): assert isinstance(result, RetrieverResult) +class TestBareModelName: + """Provider-prefix normalisation used by the config guard.""" + + @pytest.mark.parametrize( + "name,expected", + [ + ("azure/text-embedding-3-large", "text-embedding-3-large"), + ("openai/text-embedding-3-large", "text-embedding-3-large"), + ("AZURE/text-embedding-3-large", "text-embedding-3-large"), + ("vertex_ai/textembedding-gecko", "textembedding-gecko"), + ("text-embedding-3-large", "text-embedding-3-large"), + ], + ) + def test_known_prefixes_are_stripped(self, name, expected): + from graphrag_sdk.api.main import _bare_model_name + + assert _bare_model_name(name) == expected + + @pytest.mark.parametrize( + "name", + [ + # Not a provider — the first segment is part of the model id and + # dropping it would collapse two distinct models onto one name. + "my-org/custom-embedder", + # A bare name that merely contains a slashless provider word. + "azure-lookalike-model", + # Trailing slash leaves nothing to keep. + "azure/", + "", + ], + ) + def test_unknown_or_degenerate_names_unchanged(self, name): + from graphrag_sdk.api.main import _bare_model_name + + assert _bare_model_name(name) == name + + def test_only_first_segment_stripped(self): + """Multi-segment ids keep the part that identifies the model.""" + from graphrag_sdk.api.main import _bare_model_name + + assert ( + _bare_model_name("openrouter/anthropic/claude-3") + == "anthropic/claude-3" + ) + + def test_falls_back_when_litellm_missing(self, monkeypatch): + """litellm is an optional extra; the guard must still work without it.""" + import builtins + + from graphrag_sdk.api.main import ( + _bare_model_name, + _litellm_provider_prefixes, + ) + + _litellm_provider_prefixes.cache_clear() + real_import = builtins.__import__ + + def _no_litellm(name, *args, **kwargs): + if name == "litellm": + raise ImportError("litellm not installed") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _no_litellm) + try: + assert _bare_model_name("azure/text-embedding-3-large") == "text-embedding-3-large" + finally: + _litellm_provider_prefixes.cache_clear() + + +class TestConfigProviderPrefix: + """The stored model name may carry a litellm routing prefix that the + live embedder doesn't (or vice versa) when a graph is moved between + providers. Same model, same vectors — must not raise.""" + + @staticmethod + def _embedder_named(name: str, dimension: int = 8): + class _Named(MockEmbedder): + @property + def model_name(self) -> str: + return name + + return _Named(dimension=dimension) + + async def _retrieve_with(self, mock_conn, embedder, stored_model): + llm = MockLLM(responses=["unused"]) + g = GraphRAG(connection=mock_conn, llm=llm, embedder=embedder, embedding_dimension=8) + config_result = MagicMock() + config_result.result_set = [[stored_model, 8]] + g._graph_store.query_raw = AsyncMock(return_value=config_result) + mock_strategy = MagicMock(spec=RetrievalStrategy) + mock_strategy.search = AsyncMock(return_value=RetrieverResult(items=[])) + g._retrieval_strategy = mock_strategy + return await g.retrieve("test?") + + async def test_stored_prefixed_current_bare_passes(self, mock_conn): + """Graph built through Azure, now reached directly.""" + result = await self._retrieve_with( + mock_conn, self._embedder_named("mock-embedder"), "azure/mock-embedder" + ) + assert isinstance(result, RetrieverResult) + + async def test_stored_bare_current_prefixed_passes(self, mock_conn): + """The prod case: graph stored the bare name, config moved to Azure.""" + result = await self._retrieve_with( + mock_conn, self._embedder_named("azure/mock-embedder"), "mock-embedder" + ) + assert isinstance(result, RetrieverResult) + + async def test_differing_providers_same_model_passes(self, mock_conn): + result = await self._retrieve_with( + mock_conn, + self._embedder_named("openai/mock-embedder"), + "azure/mock-embedder", + ) + assert isinstance(result, RetrieverResult) + + async def test_genuine_model_change_still_raises(self, mock_conn): + """The guard must keep catching a real swap — a prefix must not be + able to mask two different models.""" + with pytest.raises(ConfigError, match="Embedding model mismatch"): + await self._retrieve_with( + mock_conn, + self._embedder_named("azure/text-embedding-ada-002"), + "azure/text-embedding-3-large", + ) + + async def test_unknown_prefix_still_raises(self, mock_conn): + """A non-provider first segment is part of the model identity, so + two such names remain distinct.""" + with pytest.raises(ConfigError, match="Embedding model mismatch"): + await self._retrieve_with( + mock_conn, + self._embedder_named("team-a/mock-embedder"), + "team-b/mock-embedder", + ) + + async def test_dimension_mismatch_still_raises_under_prefix(self, mock_conn): + """Prefix normalisation must not weaken the dimension check.""" + llm = MockLLM(responses=["unused"]) + g = GraphRAG( + connection=mock_conn, + llm=llm, + embedder=self._embedder_named("azure/mock-embedder"), + embedding_dimension=8, + ) + config_result = MagicMock() + config_result.result_set = [["mock-embedder", 1536]] + g._graph_store.query_raw = AsyncMock(return_value=config_result) + + with pytest.raises(ConfigError, match="Embedding dimension mismatch"): + await g.retrieve("test?") + + class TestGraphRAGEmbedderProbe: """A5: probe embedder dimension at validation time.""" From 160fd8f49bd48df8af073abee5d47c4456f48ae2 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:35:06 +0300 Subject: [PATCH 11/21] test: add chunk-cache edge coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eight tests against a real FalkorDB, exercising cases the existing suite does not reach. The most valuable is `test_cached_and_uncached_updates_converge`: it runs the same edit twice, once with the cache and once without, and asserts the resulting graphs are identical. That makes the cache's correctness property explicit rather than checking hand-written expectations, and it catches faults in the rebuild Cypher itself — a mutation swapping two returned columns is invisible to fixture-backed tests but fails here. The rest cover reordered chunks, duplicate identical chunks, provenance isolation across documents, orphan cleanup when a chunk disappears, stability across repeated cached updates, and the no-op short circuit reporting no cache stats. Uses a local `ScriptedExtractor` that emits relationships; the shared scripted fixture always returns an empty relationship list, so the rebuild path was previously never executed. --- .../test_cached_chunk_extraction_edges.py | 455 ++++++++++++++++++ 1 file changed, 455 insertions(+) create mode 100644 graphrag_sdk/tests/test_cached_chunk_extraction_edges.py diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction_edges.py b/graphrag_sdk/tests/test_cached_chunk_extraction_edges.py new file mode 100644 index 00000000..dc217021 --- /dev/null +++ b/graphrag_sdk/tests/test_cached_chunk_extraction_edges.py @@ -0,0 +1,455 @@ +# Local-only adversarial edge coverage for the chunk-level extraction cache. +# +# The PR's own integration tests script the LLM with `relationships: []`, so +# the relationship-rebuild path is never exercised against a real FalkorDB. +# These tests close that gap and add the equivalence property that matters +# most in production: a cached update and a full-extraction update must +# converge on the SAME graph. +# +# Run: RUN_INTEGRATION=1 FALKOR_PORT=6399 pytest tests/test_cached_chunk_extraction_edges.py + +from __future__ import annotations + +from typing import Any + +import pytest + +from graphrag_sdk.core.context import Context +from graphrag_sdk.core.models import ( + EntityMention, + GraphData, + GraphNode, + GraphRelationship, + TextChunk, + TextChunks, +) +from graphrag_sdk.ingestion.chunking_strategies.base import ChunkingStrategy +from graphrag_sdk.ingestion.extraction_strategies.base import ExtractionStrategy +from graphrag_sdk.ingestion.extraction_strategies.entity_extractors import ( + compute_entity_id, +) +from graphrag_sdk.ingestion.resolution_strategies.exact_match import ExactMatchResolution + +SEP = "\n@@@\n" + + +class SepChunking(ChunkingStrategy): + """Split on an explicit separator so chunk boundaries are exact. + + FixedSizeChunking makes "which chunks are byte-identical" a function of + character arithmetic; that hides real cache behaviour behind padding + tricks. Here the test author states it directly. + """ + + async def chunk(self, text: str, ctx: Context) -> TextChunks: + parts = [p for p in text.split(SEP) if p.strip()] + return TextChunks( + chunks=[TextChunk(text=p, index=i, metadata={}) for i, p in enumerate(parts)] + ) + + +class ScriptedExtractor(ExtractionStrategy): + """Content-keyed extractor: the chunk text *is* the script. + + Line grammar (one directive per line): + E||| + R|||| + + Deterministic and side-effect free, so the same chunk text always yields + the same GraphData — which is what makes the cached-vs-uncached + equivalence assertion meaningful. Mirrors GraphExtraction's id/property + construction exactly (compute_entity_id, label=type, `fact` string) so + cached rows MERGE onto the same nodes. + """ + + def __init__(self) -> None: + self.seen_chunk_texts: list[str] = [] + + async def extract(self, chunks: TextChunks, ontology, ctx) -> GraphData: + nodes: dict[str, GraphNode] = {} + rels: list[GraphRelationship] = [] + mentions: list[EntityMention] = [] + types: dict[str, str] = {} + + for chunk in chunks.chunks: + self.seen_chunk_texts.append(chunk.text) + for line in chunk.text.splitlines(): + line = line.strip() + if line.startswith("E|"): + _, name, etype, desc = line.split("|", 3) + types[name.strip().lower()] = etype + eid = compute_entity_id(name, etype) + existing = nodes.get(eid) + srcs = list(existing.properties["source_chunk_ids"]) if existing else [] + if chunk.uid not in srcs: + srcs.append(chunk.uid) + nodes[eid] = GraphNode( + id=eid, + label=etype, + properties={ + "name": name, + "type": etype, + "description": desc, + "source_chunk_ids": srcs, + }, + ) + mentions.append(EntityMention(chunk_id=chunk.uid, entity_id=eid)) + + # Second pass so relationship endpoints can reference entities + # declared in any chunk of this batch (matches the real extractor, + # which resolves types across the whole extraction). + for chunk in chunks.chunks: + for line in chunk.text.splitlines(): + line = line.strip() + if line.startswith("R|"): + _, src, tgt, rtype, desc = line.split("|", 4) + fact = f"({src}, {rtype}, {tgt}): {desc}" if desc else f"({src}, {rtype}, {tgt})" + rels.append( + GraphRelationship( + start_node_id=compute_entity_id(src, types.get(src.strip().lower(), "")), + end_node_id=compute_entity_id(tgt, types.get(tgt.strip().lower(), "")), + type="RELATES", + properties={ + "rel_type": rtype, + "fact": fact, + "description": desc, + "source_chunk_ids": [chunk.uid], + "src_name": src, + "tgt_name": tgt, + }, + ) + ) + + return GraphData(nodes=list(nodes.values()), relationships=rels, mentions=mentions) + + +# ── Snapshot helpers ──────────────────────────────────────────────── + + +async def _snapshot(rag) -> dict[str, Any]: + """Content-addressed view of the graph, free of volatile chunk uids. + + Chunk uids are regenerated on every update, so a raw dump can never be + compared across runs. Everything here is keyed by chunk *text* instead, + which is exactly the equivalence we care about. + """ + q = rag._graph_store.query_raw + + r = await q( + "MATCH (e:__Entity__) RETURN e.name, e.type, e.description ORDER BY e.name, e.type" + ) + entities = sorted(tuple(row) for row in (r.result_set or [])) + + r = await q( + "MATCH (a:__Entity__)-[rel:RELATES]->(b:__Entity__) " + "RETURN a.name, b.name, rel.rel_type, rel.fact ORDER BY a.name, b.name" + ) + relationships = sorted(tuple(row) for row in (r.result_set or [])) + + r = await q( + "MATCH (e:__Entity__)-[:MENTIONED_IN]->(c:Chunk) RETURN e.name, c.text" + ) + mentions = sorted(tuple(row) for row in (r.result_set or [])) + + r = await q("MATCH (:Document)-[:PART_OF]->(c:Chunk) RETURN c.text ORDER BY c.index") + chunk_texts = sorted(row[0] for row in (r.result_set or [])) + + # Entity provenance, resolved from volatile uids back to chunk text. + r = await q( + "MATCH (e:__Entity__) WHERE e.source_chunk_ids IS NOT NULL " + "UNWIND e.source_chunk_ids AS cid " + "OPTIONAL MATCH (c:Chunk {id: cid}) RETURN e.name, c.text" + ) + provenance = sorted(tuple(row) for row in (r.result_set or [])) + + return { + "entities": entities, + "relationships": relationships, + "mentions": mentions, + "chunks": chunk_texts, + "provenance": provenance, + } + + +async def _live_chunk_ids(rag, document_id: str) -> set[str]: + r = await rag._graph_store.query_raw( + "MATCH (:Document {id: $d})-[:PART_OF]->(c:Chunk) RETURN collect(c.id)", + {"d": document_id}, + ) + return set((r.result_set or [[[]]])[0][0] or []) + + +async def _assert_no_dangling_provenance(rag) -> None: + """No entity or RELATES edge may cite a chunk id that no longer exists. + + Dangling provenance is the classic silent corruption from an id remap + bug: retrieval still "works" but citations resolve to nothing. + """ + r = await rag._graph_store.query_raw("MATCH (c:Chunk) RETURN collect(c.id)") + live = set((r.result_set or [[[]]])[0][0] or []) + + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__) WHERE e.source_chunk_ids IS NOT NULL " + "RETURN e.name, e.source_chunk_ids" + ) + for name, srcs in r.result_set or []: + stale = set(srcs or []) - live + assert not stale, f"entity {name!r} cites dead chunks: {stale}" + + r = await rag._graph_store.query_raw( + "MATCH ()-[rel:RELATES]->() WHERE rel.source_chunk_ids IS NOT NULL " + "RETURN rel.fact, rel.source_chunk_ids" + ) + for fact, srcs in r.result_set or []: + stale = set(srcs or []) - live + assert not stale, f"RELATES {fact!r} cites dead chunks: {stale}" + + +# ── Fixtures ──────────────────────────────────────────────────────── + + +@pytest.fixture +def make_rag(real_falkordb_rag_factory): + from tests.conftest import MockLLM + + def _make(): + return real_falkordb_rag_factory(llm=MockLLM(), resolver=ExactMatchResolution()) + + return _make + + +# ── Document versions ─────────────────────────────────────────────── + +C_ALICE = "E|Alice|Person|Engineer at Acme\nE|Acme|Organization|Tech company\nR|Alice|Acme|WORKS_AT|Alice is employed by Acme" +C_CAROL = "E|Carol|Person|Manager in Berlin\nE|Berlin|Location|German city\nR|Carol|Berlin|BASED_IN|Carol runs the Berlin office" +C_DAVE = "E|Dave|Person|Manager in Munich\nE|Munich|Location|German city\nR|Dave|Munich|BASED_IN|Dave runs the Munich office" +C_EVE = "E|Eve|Person|Designer at Acme\nE|Acme|Organization|Tech company\nR|Eve|Acme|WORKS_AT|Eve is employed by Acme" + + +@pytest.mark.asyncio +@pytest.mark.integration +class TestChunkCacheEdges: + async def test_cached_and_uncached_updates_converge(self, make_rag): + """THE property test: caching is an optimisation, not a semantic change. + + Same ingest, same update, run twice on two isolated graphs — once + with the cache and once without. Every entity, relationship, + mention and provenance edge must match. Any divergence here is a + correctness bug in the remap, not a tuning issue. + """ + v1 = SEP.join([C_ALICE, C_CAROL]) + v2 = SEP.join([C_ALICE, C_DAVE]) + + snapshots = {} + for cached in (True, False): + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest( + text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), + ) + await rag.update( + text=v2, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=cached, + ) + snapshots[cached] = await _snapshot(rag) + await _assert_no_dangling_provenance(rag) + + assert snapshots[True] == snapshots[False], ( + "cached update diverged from full re-extraction:\n" + f"cached : {snapshots[True]}\n" + f"uncached: {snapshots[False]}" + ) + + async def test_relationship_rebuilt_from_cache_survives_update(self, make_rag): + """Relationship rebuild end-to-end — the gap in the PR's own tests. + + The unchanged chunk carries a RELATES edge. After a cached update + that edge must still exist, keep its rel_type/fact, and cite only + live chunks. + """ + v1 = SEP.join([C_ALICE, C_CAROL]) + v2 = SEP.join([C_ALICE, C_DAVE]) + + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution()) + + before = len(ex.seen_chunk_texts) + result = await rag.update( + text=v2, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True, + ) + after = ex.seen_chunk_texts[before:] + + assert result.metadata["cache_stats"] == {"cached_chunks": 1, "extracted_chunks": 1} + assert after == [C_DAVE], f"only the changed chunk may be extracted, got {after}" + + r = await rag._graph_store.query_raw( + "MATCH (a:__Entity__ {name:'Alice'})-[rel:RELATES]->(b:__Entity__ {name:'Acme'}) " + "RETURN rel.rel_type, rel.fact, rel.source_chunk_ids" + ) + assert r.result_set, "cached relationship was lost by the update" + rel_type, fact, srcs = r.result_set[0] + assert rel_type == "WORKS_AT" + assert "Alice is employed by Acme" in fact + live = await _live_chunk_ids(rag, "doc") + assert srcs and set(srcs) <= live, f"relationship provenance is stale: {set(srcs) - live}" + + # Replaced chunk's relationship must be gone. + r = await rag._graph_store.query_raw( + "MATCH (:__Entity__ {name:'Carol'})-[rel:RELATES]->() RETURN count(rel)" + ) + assert (r.result_set[0][0] if r.result_set else 0) == 0 + await _assert_no_dangling_provenance(rag) + + async def test_reordered_chunks_are_all_cached(self, make_rag): + """Reordering paragraphs changes the document hash but no chunk text. + + Every chunk must be a cache hit and zero extraction must happen — + this is the common 'section moved' docs PR. + """ + v1 = SEP.join([C_ALICE, C_CAROL, C_DAVE]) + v2 = SEP.join([C_DAVE, C_ALICE, C_CAROL]) + + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution()) + before = len(ex.seen_chunk_texts) + + result = await rag.update( + text=v2, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True, + ) + + assert result.no_op is False, "reordering changes the doc hash, so this is a real update" + assert result.metadata["cache_stats"] == {"cached_chunks": 3, "extracted_chunks": 0} + assert ex.seen_chunk_texts[before:] == [], "zero LLM extraction expected" + + snap = await _snapshot(rag) + assert {e[0] for e in snap["entities"]} == { + "Alice", "Acme", "Carol", "Berlin", "Dave", "Munich", + } + assert len(snap["relationships"]) == 3 + await _assert_no_dangling_provenance(rag) + + async def test_duplicate_identical_chunks_each_get_provenance(self, make_rag): + """Two byte-identical chunks map to ONE old chunk id. + + Both new uids must receive their own mention and appear in the + entity's provenance — a naive dict remap would keep only one. + """ + v1 = SEP.join([C_ALICE, C_CAROL]) + v2 = SEP.join([C_ALICE, C_ALICE, C_DAVE]) + + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution()) + + result = await rag.update( + text=v2, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True, + ) + assert result.metadata["cache_stats"] == {"cached_chunks": 2, "extracted_chunks": 1} + + r = await rag._graph_store.query_raw( + "MATCH (:__Entity__ {name:'Alice'})-[:MENTIONED_IN]->(c:Chunk) RETURN count(c)" + ) + assert r.result_set[0][0] == 2, "each duplicate chunk needs its own mention" + + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__ {name:'Alice'}) RETURN e.source_chunk_ids" + ) + srcs = set(r.result_set[0][0] or []) + live = await _live_chunk_ids(rag, "doc") + assert len(srcs) == 2 and srcs <= live + await _assert_no_dangling_provenance(rag) + + async def test_other_document_provenance_untouched(self, make_rag): + """An entity shared with a second document must keep that document's + chunk ids after a cached update of the first.""" + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=SEP.join([C_ALICE, C_CAROL]), document_id="docA", + chunker=SepChunking(), extractor=ex, resolver=ExactMatchResolution()) + await rag.ingest(text=C_EVE, document_id="docB", + chunker=SepChunking(), extractor=ex, resolver=ExactMatchResolution()) + + b_ids = await _live_chunk_ids(rag, "docB") + + await rag.update( + text=SEP.join([C_ALICE, C_DAVE]), document_id="docA", chunker=SepChunking(), + extractor=ex, resolver=ExactMatchResolution(), cache_unchanged_chunks=True, + ) + + r = await rag._graph_store.query_raw( + "MATCH (e:__Entity__ {name:'Acme'}) RETURN e.source_chunk_ids" + ) + srcs = set(r.result_set[0][0] or []) + assert b_ids and b_ids <= srcs, ( + f"docB provenance was dropped by docA's cached update: missing {b_ids - srcs}" + ) + assert srcs & (await _live_chunk_ids(rag, "docA")), "docA provenance missing" + await _assert_no_dangling_provenance(rag) + + async def test_removed_chunk_orphans_cleaned_under_cache(self, make_rag): + """Shrinking a document must still orphan-clean, cache or not.""" + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=SEP.join([C_ALICE, C_CAROL, C_DAVE]), document_id="doc", + chunker=SepChunking(), extractor=ex, resolver=ExactMatchResolution()) + + result = await rag.update( + text=C_ALICE, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True, + ) + assert result.metadata["cache_stats"] == {"cached_chunks": 1, "extracted_chunks": 0} + + snap = await _snapshot(rag) + names = {e[0] for e in snap["entities"]} + assert names == {"Alice", "Acme"}, f"orphans survived: {names}" + assert len(snap["relationships"]) == 1 + await _assert_no_dangling_provenance(rag) + + async def test_repeated_cached_updates_are_stable(self, make_rag): + """Idempotency: the same cached update applied repeatedly must not + accumulate duplicate entities, mentions or provenance entries.""" + v1 = SEP.join([C_ALICE, C_CAROL]) + v2 = SEP.join([C_ALICE, C_DAVE]) + + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution()) + + snaps = [] + for _ in range(3): + await rag.update(text=v2, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True) + await rag.update(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution(), cache_unchanged_chunks=True) + snaps.append(await _snapshot(rag)) + await _assert_no_dangling_provenance(rag) + + assert snaps[0] == snaps[1] == snaps[2], "cached updates are not idempotent" + + async def test_noop_short_circuit_reports_no_cache_stats(self, make_rag): + """Identical content still takes the document-hash fast path; the + cache flag must not disturb it.""" + v1 = SEP.join([C_ALICE, C_CAROL]) + rag = make_rag() + ex = ScriptedExtractor() + await rag.ingest(text=v1, document_id="doc", chunker=SepChunking(), extractor=ex, + resolver=ExactMatchResolution()) + before = len(ex.seen_chunk_texts) + + result = await rag.update(text=v1, document_id="doc", chunker=SepChunking(), + extractor=ex, resolver=ExactMatchResolution(), + cache_unchanged_chunks=True) + + assert result.no_op is True + assert "cache_stats" not in result.metadata + assert ex.seen_chunk_texts[before:] == [] From 9c251fbab286622721025db000438d49b9e80310 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:19:13 +0300 Subject: [PATCH 12/21] fix: never treat a chunk as its own cache entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cache reads chunk texts for `document_id` and matches them against the chunks being written. If those two sets ever overlap, a chunk hashes to itself: the rebuild finds no mentions for a chunk nothing has been extracted from yet, returns empty GraphData, and skips the extractor. Entities are dropped with no exception raised, so the fail-open path can't catch it, and cache stats still report a hit. `update()` is safe today because it writes under a pending document id while the cache reads the resolved id — but that's a property of the caller, not of this class, and `CachedChunkExtraction` is publicly exported. The pipeline persists chunks in step 3, before extraction runs in step 4, so any caller whose read id equals the id being written reads back the chunks it just wrote. Filter the write set out of the lookup. Excluding those ids also keeps them out of `all_old_ids`, so provenance remapping doesn't treat a live chunk as one dying in the cutover. Self-referential chunks now fall through to normal extraction. Verified no-op for the current `update()` path: an end-to-end run against a real FalkorDB produces byte-identical output with and without this change. --- .../cached_chunk_extraction.py | 22 +++++- .../tests/test_cached_chunk_extraction.py | 72 +++++++++++++++++++ 2 files changed, 91 insertions(+), 3 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py index 1589ee85..860ee6cc 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -102,11 +102,25 @@ def __init__( self.cached_chunk_count = 0 self.extracted_chunk_count = 0 - async def _old_chunks_by_hash(self) -> tuple[dict[str, str], set[str]]: - """Return (sha256(chunk text) -> old chunk id, all old chunk ids).""" + async def _old_chunks_by_hash(self, exclude: set[str]) -> tuple[dict[str, str], set[str]]: + """Return (sha256(chunk text) -> old chunk id, all old chunk ids). + + Ids in ``exclude`` — the uids of the chunks currently being written — + are dropped. ``update()`` writes under a pending document id, so the + read and write sets are disjoint there, but that is a property of the + caller, not of this class. If they ever overlap (wired into a path + where the read id is the write id, or the pending indirection + changes), a chunk would hash to *itself*: the rebuild would find no + mentions for a chunk nothing has been extracted from yet, emit empty + ``GraphData``, and skip the extractor — dropping entities silently, + with no exception for the fail-open path to catch and cache stats + still reporting a hit. Excluding them degrades to normal extraction. + """ by_hash: dict[str, str] = {} all_ids: set[str] = set() for cid, text in await self._graph_store.get_document_chunk_texts(self._document_id): + if cid in exclude: + continue all_ids.add(cid) # First occurrence wins; duplicates map to the same content anyway. by_hash.setdefault(_sha256(text), cid) @@ -262,7 +276,9 @@ async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> """Split chunks into cached vs. new, rebuild the former from the graph, extract the latter with the inner strategy, and merge.""" try: - old_by_hash, all_old_ids = await self._old_chunks_by_hash() + old_by_hash, all_old_ids = await self._old_chunks_by_hash( + {c.uid for c in chunks.chunks} + ) except Exception as exc: logger.warning("Chunk cache lookup failed, extracting everything: %s", exc) old_by_hash, all_old_ids = {}, set() diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py index 836c6f2c..22ecdc94 100644 --- a/graphrag_sdk/tests/test_cached_chunk_extraction.py +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -831,3 +831,75 @@ async def test_no_op_short_circuit_unaffected( cache_unchanged_chunks=True, ) assert result.no_op is True + + +class TestSelfMappingGuard: + """A chunk must never be treated as its own cache entry. + + ``update()`` writes new chunks under a pending document id, so its read + set (``resolved_id``) and write set are disjoint and this cannot fire + today. But the pipeline persists chunks in Step 3 *before* extraction + runs in Step 4, and ``CachedChunkExtraction`` is publicly exported — so + any caller whose read id equals the id being written (e.g. wiring it + into ``ingest()``) would read back the chunks just written. + + Each chunk would then hash to itself, the rebuild would find no mentions + for a chunk nothing has been extracted from yet, and the strategy would + return empty ``GraphData`` while reporting a cache hit. No exception is + raised, so the fail-open path cannot catch it: entities vanish silently. + """ + + async def test_self_referential_chunk_is_extracted_not_cached(self, ontology, ctx): + # Store returns the SAME uid the pipeline is writing, with no + # extracted entities behind it yet. + store = FakeGraphStore(chunk_texts=[("new-1", "alpha")], entity_rows=[]) + inner = RecordingExtractor( + result=GraphData( + nodes=[GraphNode(id="alice__person", label="Person", properties={"name": "Alice"})], + mentions=[EntityMention(chunk_id="new-1", entity_id="alice__person")], + ) + ) + strategy = CachedChunkExtraction(inner, store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx + ) + + assert strategy.cached_chunk_count == 0, "chunk was cached against itself" + assert strategy.extracted_chunk_count == 1 + assert len(inner.calls) == 1, "extractor was skipped — entities would be lost" + assert [n.id for n in result.nodes] == ["alice__person"] + assert len(result.mentions) == 1 + + async def test_genuine_old_chunk_still_cached_alongside_self_reference( + self, ontology, ctx + ): + """The guard must drop only the self-reference, not real cache hits.""" + store = FakeGraphStore( + chunk_texts=[("new-1", "alpha"), ("old-2", "beta")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-2", + entity_id="acme__organization", + label="Organization", + name="Acme", + type="Organization", + source_chunk_ids=["old-2"], + ), + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("alpha", 0, "new-1"), _chunk("beta", 1, "new-2")]), + ontology, + ctx, + ) + + assert strategy.cached_chunk_count == 1 + assert strategy.extracted_chunk_count == 1 + assert [c.text for c in inner.calls[0].chunks] == ["alpha"] + assert [n.id for n in result.nodes] == ["acme__organization"] + # Provenance for the cached chunk is remapped onto the new uid. + assert result.nodes[0].properties["source_chunk_ids"] == ["new-2"] From 42a97840428d5abe121222adc4a66ec3ff61eedd Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:10:10 +0300 Subject: [PATCH 13/21] refactor: match embedding models by route segment The config guard compared stored and live embedder names as raw strings, so a graph reached through a different endpoint ("azure/text-embedding-3-large" vs "text-embedding-3-large") was reported as a model mismatch even though the stored vectors are unchanged. Treat one leading route segment as optional on either side, which covers a prefix appearing, disappearing, or changing. Matching is anchored on "/" so only a whole segment is ever ignored -- a substring check would accept "text-embedding-3-large" against "text-embedding-3-large-v2", silently passing the mismatch this guard exists to catch. This drops the provider-prefix lookup and its static fallback list, so the comparison no longer depends on the optional litellm extra or on a list that goes stale as providers are added. --- graphrag_sdk/src/graphrag_sdk/api/main.py | 96 ++++++++--------------- 1 file changed, 31 insertions(+), 65 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/api/main.py b/graphrag_sdk/src/graphrag_sdk/api/main.py index cc8bf47a..f3a4b02d 100644 --- a/graphrag_sdk/src/graphrag_sdk/api/main.py +++ b/graphrag_sdk/src/graphrag_sdk/api/main.py @@ -10,7 +10,6 @@ import logging import os import re -from functools import lru_cache from typing import Any, Literal, overload from uuid import uuid4 @@ -148,71 +147,40 @@ def _neutralize_context_close_tag(text: str) -> str: return _CONTEXT_CLOSE_RE.sub("", text) -# Fallback provider prefixes for :py:func:`_bare_model_name`, used when -# litellm isn't installed (it is an optional extra). Covers the providers -# that actually serve embedding models; the full list is pulled from -# litellm at runtime when available. -_FALLBACK_LITELLM_PROVIDERS = frozenset( - { - "azure", - "azure_ai", - "bedrock", - "cohere", - "databricks", - "deepinfra", - "gemini", - "huggingface", - "mistral", - "nvidia_nim", - "ollama", - "openai", - "openrouter", - "together_ai", - "vertex_ai", - "voyage", - "watsonx", - "xinference", - } -) +def _same_embedding_model(stored: str, current: str) -> bool: + """Whether two model identifiers name the same embedding model. + ``"azure/text-embedding-3-large"`` and ``"text-embedding-3-large"`` are the + *same* model reached through different endpoints — the leading segment is a + routing prefix telling the client where to send the request, it is not part + of the model identity. A plain string comparison would reject a graph whose + provider changed, even though the stored vectors are unchanged. -@lru_cache(maxsize=1) -def _litellm_provider_prefixes() -> frozenset[str]: - """Known litellm routing prefixes, lowercased. + One leading route segment is therefore optional on either side: the names + match if they are equal with or without it, which covers a prefix appearing, + disappearing, or changing (``"azure/…"`` -> ``"openai/…"``). Comparison is + case-insensitive and ignores surrounding whitespace. - Prefers litellm's own ``provider_list`` so the set stays correct as - litellm adds providers; falls back to a static set when litellm isn't - installed. Cached — the list is static for the process lifetime. - """ - try: - import litellm - - providers = {str(getattr(p, "value", p)).lower() for p in litellm.provider_list} - if providers: - return frozenset(providers | _FALLBACK_LITELLM_PROVIDERS) - except Exception: # pragma: no cover - depends on optional extra - logger.debug("litellm provider_list unavailable", exc_info=True) - return _FALLBACK_LITELLM_PROVIDERS - - -def _bare_model_name(name: str) -> str: - """Strip litellm's provider routing prefix from a model identifier. - - ``"azure/text-embedding-3-large"`` and ``"text-embedding-3-large"`` name - the *same* model reached through different endpoints — the prefix tells - litellm where to send the request, it is not part of the model identity. - Comparing the raw strings therefore reports a mismatch for a graph that - merely changed provider, even though the vectors are unchanged. - - Only the first segment is removed, and only when it is a recognised - provider, so multi-segment identifiers such as - ``"openrouter/anthropic/claude-3"`` keep the part that distinguishes the - model. Names without a known prefix are returned unchanged. + Only a whole ``/``-delimited segment can be ignored, never a substring: a + looser check would accept ``"text-embedding-3-large"`` against + ``"text-embedding-3-large-v2"`` — different models with different vectors — + and silently pass the mismatch this guard exists to catch. + + A routing prefix is indistinguishable from a vendor namespace, so + ``"my-org/custom-embedder"`` and ``"custom-embedder"`` also compare equal. """ - head, sep, tail = name.partition("/") - if sep and tail and head.lower() in _litellm_provider_prefixes(): - return tail - return name + + def _variants(name: str) -> set[str]: + name = name.strip().lower() + if not name: + return set() + forms = {name} + _, sep, tail = name.partition("/") + if sep and tail: + forms.add(tail) + return forms + + return bool(_variants(stored) & _variants(current)) def _strip_and_load_json(text: str) -> Any: @@ -2946,9 +2914,7 @@ async def _validate_graph_config(self, *, ctx: Context | None = None) -> None: stored_dim = result.result_set[0][1] current_model = self.embedder.model_name - if stored_model and _bare_model_name(stored_model) != _bare_model_name( - current_model - ): + if stored_model and not _same_embedding_model(stored_model, current_model): raise ConfigError( f"Embedding model mismatch: graph was built with " f"'{stored_model}' but current embedder is " From 19c9aebd53a53880888e84c745a1564c06b21763 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:10:19 +0300 Subject: [PATCH 14/21] test: cover route-segment embedding model matching Replace the provider-prefix normalisation tests with cases for the new comparison: a route appearing, disappearing, or changing between two providers, multi-segment routes, and case/whitespace differences. Pin the negative cases that keep the guard useful -- a trailing word such as "text-embedding-3-large-v2" must not match "text-embedding-3-large", since a looser substring check would accept it and pass a real mismatch. Two config-guard tests changed contract: two differently namespaced names now resolve to the same model, and the test asserting the opposite is replaced. The test stubbing out a missing litellm import is dropped -- the comparison no longer imports it. --- graphrag_sdk/tests/test_facade.py | 116 ++++++++++++++---------------- 1 file changed, 55 insertions(+), 61 deletions(-) diff --git a/graphrag_sdk/tests/test_facade.py b/graphrag_sdk/tests/test_facade.py index 3e640908..0df6d48f 100644 --- a/graphrag_sdk/tests/test_facade.py +++ b/graphrag_sdk/tests/test_facade.py @@ -1174,78 +1174,69 @@ async def test_no_config_node_passes(self, mock_conn, embedder): assert isinstance(result, RetrieverResult) -class TestBareModelName: - """Provider-prefix normalisation used by the config guard.""" +class TestSameEmbeddingModel: + """Provider-route tolerance used by the config guard. + + A stored name may carry a routing prefix the live embedder doesn't (or + vice versa) when a graph moves between providers. Same model, same + vectors — must still compare equal. + """ @pytest.mark.parametrize( - "name,expected", + "stored,current", [ ("azure/text-embedding-3-large", "text-embedding-3-large"), + ("text-embedding-3-large", "azure/text-embedding-3-large"), ("openai/text-embedding-3-large", "text-embedding-3-large"), ("AZURE/text-embedding-3-large", "text-embedding-3-large"), ("vertex_ai/textembedding-gecko", "textembedding-gecko"), + # Identical on both sides, with and without a route. ("text-embedding-3-large", "text-embedding-3-large"), + ("azure/text-embedding-3-large", "azure/text-embedding-3-large"), + # Case and surrounding whitespace are not part of the identity. + ("Text-Embedding-3-Large", "text-embedding-3-large"), + (" azure/text-embedding-3-large ", "text-embedding-3-large"), + # Multi-segment routes collapse to the trailing model id. + ("openrouter/anthropic/claude-3", "anthropic/claude-3"), + # Any vendor namespace, not just the well-known providers. + ("my-org/custom-embedder", "custom-embedder"), + # Route changed on a graph that was already prefixed. + ("azure/text-embedding-3-large", "openai/text-embedding-3-large"), ], ) - def test_known_prefixes_are_stripped(self, name, expected): - from graphrag_sdk.api.main import _bare_model_name + def test_route_prefix_is_ignored(self, stored, current): + from graphrag_sdk.api.main import _same_embedding_model - assert _bare_model_name(name) == expected + assert _same_embedding_model(stored, current) is True @pytest.mark.parametrize( - "name", + "stored,current", [ - # Not a provider — the first segment is part of the model id and - # dropping it would collapse two distinct models onto one name. - "my-org/custom-embedder", - # A bare name that merely contains a slashless provider word. - "azure-lookalike-model", - # Trailing slash leaves nothing to keep. - "azure/", - "", + # Genuinely different models. + ("text-embedding-3-large", "text-embedding-3-small"), + ("azure/text-embedding-3-large", "azure/text-embedding-3-small"), + # Suffix-of-a-word, NOT a route segment. A bare `in` check would + # accept these and silently pass a real mismatch — the vectors + # differ, which is exactly what the guard exists to catch. + ("text-embedding-3-large", "text-embedding-3-large-v2"), + ("text-embedding-3-large-v2", "text-embedding-3-large"), + ("embedding-3-large", "text-embedding-3-large"), + # The route segment itself must be whole. + ("azure/text-embedding-3-large", "3-large"), + # One side empty is not a match. + ("text-embedding-3-large", ""), + ("", "text-embedding-3-large"), ], ) - def test_unknown_or_degenerate_names_unchanged(self, name): - from graphrag_sdk.api.main import _bare_model_name - - assert _bare_model_name(name) == name - - def test_only_first_segment_stripped(self): - """Multi-segment ids keep the part that identifies the model.""" - from graphrag_sdk.api.main import _bare_model_name - - assert ( - _bare_model_name("openrouter/anthropic/claude-3") - == "anthropic/claude-3" - ) - - def test_falls_back_when_litellm_missing(self, monkeypatch): - """litellm is an optional extra; the guard must still work without it.""" - import builtins - - from graphrag_sdk.api.main import ( - _bare_model_name, - _litellm_provider_prefixes, - ) - - _litellm_provider_prefixes.cache_clear() - real_import = builtins.__import__ - - def _no_litellm(name, *args, **kwargs): - if name == "litellm": - raise ImportError("litellm not installed") - return real_import(name, *args, **kwargs) + def test_different_models_do_not_match(self, stored, current): + from graphrag_sdk.api.main import _same_embedding_model - monkeypatch.setattr(builtins, "__import__", _no_litellm) - try: - assert _bare_model_name("azure/text-embedding-3-large") == "text-embedding-3-large" - finally: - _litellm_provider_prefixes.cache_clear() + assert _same_embedding_model(stored, current) is False class TestConfigProviderPrefix: - """The stored model name may carry a litellm routing prefix that the - live embedder doesn't (or vice versa) when a graph is moved between + """The stored model name may carry a routing prefix that the live + embedder doesn't (or vice versa) when a graph is moved between providers. Same model, same vectors — must not raise.""" @staticmethod @@ -1300,15 +1291,18 @@ async def test_genuine_model_change_still_raises(self, mock_conn): "azure/text-embedding-3-large", ) - async def test_unknown_prefix_still_raises(self, mock_conn): - """A non-provider first segment is part of the model identity, so - two such names remain distinct.""" - with pytest.raises(ConfigError, match="Embedding model mismatch"): - await self._retrieve_with( - mock_conn, - self._embedder_named("team-a/mock-embedder"), - "team-b/mock-embedder", - ) + async def test_same_model_under_different_vendors_passes(self, mock_conn): + """A routing prefix is indistinguishable from a vendor namespace, + so ``team-a/mock-embedder`` and ``team-b/mock-embedder`` name the + same model and must not raise. The dimension check below still + guards the case that actually corrupts retrieval. + """ + result = await self._retrieve_with( + mock_conn, + self._embedder_named("team-a/mock-embedder"), + "team-b/mock-embedder", + ) + assert isinstance(result, RetrieverResult) async def test_dimension_mismatch_still_raises_under_prefix(self, mock_conn): """Prefix normalisation must not weaken the dimension check.""" From 582a2d4dcedadc9b304a569699df23ca059a45ce Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:06:34 +0300 Subject: [PATCH 15/21] fix: treat a qualified name on both sides as two models The name comparison stripped a leading segment unconditionally, so any two identifiers sharing a tail matched. That collapsed distinct models onto one name: "BAAI/bge-m3" and "ollama/bge-m3" are a fp32 release and a quantized build, and "sentence-transformers/all-MiniLM-L6-v2" and a finetune under another org are different weights entirely. The dimension check cannot stand in for this -- every realistic collision here keeps the base model's dimensions -- and the resulting failure is soft, since retrieval still returns results, ranked against another model's vectors. Only strip a segment when the other side has none. A bare name carries no route, so a segment present on one side alone is one; two qualified names are two owners and are compared in full. Cost: a route that changes while both sides stay qualified ("azure/x" -> "openai/x") is now read as an owner change and rejected. That is rare next to the collisions it catches, and it fails loudly rather than silently. --- graphrag_sdk/src/graphrag_sdk/api/main.py | 64 ++++++++++++----------- 1 file changed, 34 insertions(+), 30 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/api/main.py b/graphrag_sdk/src/graphrag_sdk/api/main.py index f3a4b02d..bae09970 100644 --- a/graphrag_sdk/src/graphrag_sdk/api/main.py +++ b/graphrag_sdk/src/graphrag_sdk/api/main.py @@ -150,37 +150,41 @@ def _neutralize_context_close_tag(text: str) -> str: def _same_embedding_model(stored: str, current: str) -> bool: """Whether two model identifiers name the same embedding model. - ``"azure/text-embedding-3-large"`` and ``"text-embedding-3-large"`` are the - *same* model reached through different endpoints — the leading segment is a - routing prefix telling the client where to send the request, it is not part - of the model identity. A plain string comparison would reject a graph whose - provider changed, even though the stored vectors are unchanged. - - One leading route segment is therefore optional on either side: the names - match if they are equal with or without it, which covers a prefix appearing, - disappearing, or changing (``"azure/…"`` -> ``"openai/…"``). Comparison is - case-insensitive and ignores surrounding whitespace. - - Only a whole ``/``-delimited segment can be ignored, never a substring: a - looser check would accept ``"text-embedding-3-large"`` against - ``"text-embedding-3-large-v2"`` — different models with different vectors — - and silently pass the mismatch this guard exists to catch. - - A routing prefix is indistinguishable from a vendor namespace, so - ``"my-org/custom-embedder"`` and ``"custom-embedder"`` also compare equal. + A leading segment can mean two different things, and the string alone + doesn't say which: + + - a *route* — ``"azure/text-embedding-3-large"`` is the same model as + ``"text-embedding-3-large"``, just reached through a different endpoint; + - an *owner* — ``"BAAI/bge-m3"`` and ``"ollama/bge-m3"`` are a fp32 release + and a quantized build, same dimensions but different vectors. + + A bare name is the tell: it carries no route, so a segment present on only + one side is one. That case is ignored. When both sides carry a segment they + name owners and are compared in full, which keeps the pairs above distinct. + + This matters because the dimension check can't separate them — every + realistic collision here has identical dimensions — and the resulting + failure is silent: retrieval still returns results, ranked against vectors + from another model. + + Comparison is case-insensitive and ignores surrounding whitespace. Only a + whole ``/`` segment is ever ignored, never a substring, so + ``"text-embedding-3-large"`` and ``"text-embedding-3-large-v2"`` stay + distinct. """ - - def _variants(name: str) -> set[str]: - name = name.strip().lower() - if not name: - return set() - forms = {name} - _, sep, tail = name.partition("/") - if sep and tail: - forms.add(tail) - return forms - - return bool(_variants(stored) & _variants(current)) + a = stored.strip().lower() + b = current.strip().lower() + if not a or not b: + return a == b + if a == b: + return True + _, a_sep, a_tail = a.partition("/") + _, b_sep, b_tail = b.partition("/") + if a_sep and a_tail and not b_sep: + return a_tail == b + if b_sep and b_tail and not a_sep: + return b_tail == a + return False def _strip_and_load_json(text: str) -> Any: From f6ca63e0dadcf101308d3dd38f2f9bc71951fe60 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:06:42 +0300 Subject: [PATCH 16/21] test: pin owner-qualified names as distinct models Cover the pairs the dimension check cannot see -- a finetune under another org, a quantized build, and the same model id served by two vendors -- so a future loosening of the name rule has to break a test to land. Restore test_unknown_prefix_still_raises, which had been inverted to assert that two owner-qualified names pass. They do not: identical dimensions with different weights is the one case neither check would catch. Move the both-sides-qualified route change to the negative cases and state why it is the accepted cost of the rule. --- graphrag_sdk/tests/test_facade.py | 75 ++++++++++++++++++------------- 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/graphrag_sdk/tests/test_facade.py b/graphrag_sdk/tests/test_facade.py index 0df6d48f..fc0060eb 100644 --- a/graphrag_sdk/tests/test_facade.py +++ b/graphrag_sdk/tests/test_facade.py @@ -1175,33 +1175,30 @@ async def test_no_config_node_passes(self, mock_conn, embedder): class TestSameEmbeddingModel: - """Provider-route tolerance used by the config guard. + """Model-name comparison used by the config guard. - A stored name may carry a routing prefix the live embedder doesn't (or - vice versa) when a graph moves between providers. Same model, same - vectors — must still compare equal. + A leading segment is a route when the other side has none, and an owner + when both sides have one. Routes are ignorable; owners are identity. """ @pytest.mark.parametrize( "stored,current", [ + # A segment on one side only: the bare side names no route, so + # the segment is one. This is the case that reaches production — + # a caller records the bare name and later builds a routed one. ("azure/text-embedding-3-large", "text-embedding-3-large"), ("text-embedding-3-large", "azure/text-embedding-3-large"), ("openai/text-embedding-3-large", "text-embedding-3-large"), ("AZURE/text-embedding-3-large", "text-embedding-3-large"), ("vertex_ai/textembedding-gecko", "textembedding-gecko"), + ("my-org/custom-embedder", "custom-embedder"), # Identical on both sides, with and without a route. ("text-embedding-3-large", "text-embedding-3-large"), ("azure/text-embedding-3-large", "azure/text-embedding-3-large"), # Case and surrounding whitespace are not part of the identity. ("Text-Embedding-3-Large", "text-embedding-3-large"), (" azure/text-embedding-3-large ", "text-embedding-3-large"), - # Multi-segment routes collapse to the trailing model id. - ("openrouter/anthropic/claude-3", "anthropic/claude-3"), - # Any vendor namespace, not just the well-known providers. - ("my-org/custom-embedder", "custom-embedder"), - # Route changed on a graph that was already prefixed. - ("azure/text-embedding-3-large", "openai/text-embedding-3-large"), ], ) def test_route_prefix_is_ignored(self, stored, current): @@ -1215,9 +1212,20 @@ def test_route_prefix_is_ignored(self, stored, current): # Genuinely different models. ("text-embedding-3-large", "text-embedding-3-small"), ("azure/text-embedding-3-large", "azure/text-embedding-3-small"), - # Suffix-of-a-word, NOT a route segment. A bare `in` check would - # accept these and silently pass a real mismatch — the vectors - # differ, which is exactly what the guard exists to catch. + # Two owners. These are the pairs the dimension check cannot see: + # a finetune under another org, and a quantized build, both keep + # the base model's dimensions while producing different vectors. + ("sentence-transformers/all-MiniLM-L6-v2", "myorg/all-MiniLM-L6-v2"), + ("BAAI/bge-m3", "ollama/bge-m3"), + ("intfloat/e5-large-v2", "rando/e5-large-v2"), + ("openai/text-embedding-3-large", "mistralai/text-embedding-3-large"), + # Cost of the above: a route that changes while both sides stay + # qualified is read as an owner change and rejected. Rare next to + # the collisions it buys, and it fails loudly rather than silently. + ("azure/text-embedding-3-large", "openai/text-embedding-3-large"), + ("openrouter/anthropic/claude-3", "anthropic/claude-3"), + # Suffix-of-a-word, NOT a route segment. A bare substring check + # would accept these and silently pass a real mismatch. ("text-embedding-3-large", "text-embedding-3-large-v2"), ("text-embedding-3-large-v2", "text-embedding-3-large"), ("embedding-3-large", "text-embedding-3-large"), @@ -1273,13 +1281,15 @@ async def test_stored_bare_current_prefixed_passes(self, mock_conn): ) assert isinstance(result, RetrieverResult) - async def test_differing_providers_same_model_passes(self, mock_conn): - result = await self._retrieve_with( - mock_conn, - self._embedder_named("openai/mock-embedder"), - "azure/mock-embedder", - ) - assert isinstance(result, RetrieverResult) + async def test_differing_providers_both_qualified_raises(self, mock_conn): + """With a segment on both sides there is nothing to mark either as a + route, so they are read as two owners and rejected.""" + with pytest.raises(ConfigError, match="Embedding model mismatch"): + await self._retrieve_with( + mock_conn, + self._embedder_named("openai/mock-embedder"), + "azure/mock-embedder", + ) async def test_genuine_model_change_still_raises(self, mock_conn): """The guard must keep catching a real swap — a prefix must not be @@ -1291,18 +1301,21 @@ async def test_genuine_model_change_still_raises(self, mock_conn): "azure/text-embedding-3-large", ) - async def test_same_model_under_different_vendors_passes(self, mock_conn): - """A routing prefix is indistinguishable from a vendor namespace, - so ``team-a/mock-embedder`` and ``team-b/mock-embedder`` name the - same model and must not raise. The dimension check below still - guards the case that actually corrupts retrieval. + async def test_unknown_prefix_still_raises(self, mock_conn): + """Two owner-qualified names are two different models. + + The dimension check cannot stand in for this one: a finetune or a + quantized build keeps the base model's dimensions, so it passes while + producing incompatible vectors. Retrieval then fails soft — results + come back, ranked against another model's embeddings — which can sit + in a graph indefinitely. """ - result = await self._retrieve_with( - mock_conn, - self._embedder_named("team-a/mock-embedder"), - "team-b/mock-embedder", - ) - assert isinstance(result, RetrieverResult) + with pytest.raises(ConfigError, match="Embedding model mismatch"): + await self._retrieve_with( + mock_conn, + self._embedder_named("team-a/mock-embedder"), + "team-b/mock-embedder", + ) async def test_dimension_mismatch_still_raises_under_prefix(self, mock_conn): """Prefix normalisation must not weaken the dimension check.""" From ddd11117eaca7bed39577827653c12de268bf158 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:22:11 +0300 Subject: [PATCH 17/21] fix: fail fast when the database is unreachable Driver failures escaped as raw redis exceptions, so the chunk cache's broad fallback swallowed them and re-ran extraction on every chunk -- burning LLM calls before the write failed anyway. FalkorDB probes the server while the client is constructed, so the error surfaced from _ensure_client() before the retry loop ever ran. All four exit paths now raise DatabaseError, and the cache re-raises DatabaseError and LatencyBudgetExceededError instead of falling back. The wrapped message keeps the driver's original text -- vector_store and ontology_store both match on it to detect existing indexes. --- .../src/graphrag_sdk/core/connection.py | 45 +++++++++++++++---- .../cached_chunk_extraction.py | 18 +++++++- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/core/connection.py b/graphrag_sdk/src/graphrag_sdk/core/connection.py index c47127f8..a367b5ef 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/connection.py +++ b/graphrag_sdk/src/graphrag_sdk/core/connection.py @@ -11,6 +11,8 @@ from typing import Any from urllib.parse import urlparse +from graphrag_sdk.core.exceptions import DatabaseError + logger = logging.getLogger(__name__) @@ -123,9 +125,21 @@ def _ensure_client(self) -> None: pool_kwargs["ssl_keyfile"] = self.config.ssl_keyfile pool_kwargs["ssl_check_hostname"] = self.config.ssl_check_hostname - self._pool = BlockingConnectionPool(**pool_kwargs) - self._driver = FalkorDB(connection_pool=self._pool) - self._graph = self._driver.select_graph(self.config.graph_name) + self._pool = self._pool or BlockingConnectionPool(**pool_kwargs) + try: + self._driver = FalkorDB(connection_pool=self._pool) + self._graph = self._driver.select_graph(self.config.graph_name) + except Exception as exc: + # The driver probes the server during construction, so an + # unreachable database surfaces here rather than at query time. + # The pool is kept: it connects lazily and is reusable, and + # only close() can release it. Dropping it here would strand + # a pool per failed attempt with no way to aclose() it. + self._driver = None + self._graph = None + raise DatabaseError( + f"Could not connect to FalkorDB at {self.config.host}:{self.config.port}: {exc}" + ) from exc logger.info( "Connected to FalkorDB (async) at %s:%s (tls=%s)", @@ -158,12 +172,18 @@ async def query( Returns: ``QueryResult`` from the async FalkorDB driver. + + Raises: + DatabaseError: The query could not be completed — the connection is + unhealthy, retries were exhausted, or the failure is permanent. + Driver-level exceptions are wrapped so callers can distinguish + infrastructure failures from application errors. """ self._ensure_client() assert self._graph is not None # for type-checkers if not await self._breaker.allow_request(): - raise ConnectionError( + raise DatabaseError( "Circuit breaker is open — FalkorDB connection is unhealthy. " "Requests will resume after recovery timeout." ) @@ -186,7 +206,7 @@ async def query( exc, ) logger.debug("Non-transient FalkorDB query failure details", exc_info=True) - raise + raise DatabaseError(f"FalkorDB query failed: {exc}") from exc await self._breaker.record_failure() logger.warning( "Query attempt %d/%d failed: %s", @@ -210,8 +230,10 @@ async def query( "FalkorDB query failure details", exc_info=(type(last_exc), last_exc, last_exc.__traceback__), ) - raise last_exc - raise RuntimeError("FalkorDB query failed without an exception") + raise DatabaseError( + f"FalkorDB query failed after {self.config.retry_count} attempts: {last_exc}" + ) from last_exc + raise DatabaseError("FalkorDB query failed without an exception") # Substrings that indicate a non-transient (permanent) error — # retrying will never succeed. @@ -229,9 +251,14 @@ def _is_non_transient(cls, exc: Exception) -> bool: # ── Health & Admin ──────────────────────────────────────────── async def ping(self) -> bool: - """Send a Redis PING to verify the connection is alive.""" - self._ensure_client() + """Send a Redis PING to verify the connection is alive. + + Returns ``False`` rather than raising when the server cannot be + reached — reporting "not alive" is the point of the call. + """ try: + self._ensure_client() + from redis.asyncio import Redis redis: Redis = Redis(connection_pool=self._pool) diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py index 860ee6cc..55d0ad2e 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -27,8 +27,11 @@ then strips the old chunk ids. - Mentions are re-emitted against the new chunk uid, so the ``update()`` orphan cleanup keeps every entity that unchanged chunks still support. -- Every cache failure falls back to real extraction (fail-open): the worst - case is paying for LLM calls that could have been skipped, never data loss. +- A cache failure falls back to real extraction (fail-open): the worst case is + paying for LLM calls that could have been skipped, never data loss. The + exception is an unreachable graph (``DatabaseError``) or an exhausted + latency budget, which propagate — extraction only needs the LLM, so it + would bill in full and then fail at the write phase regardless. Semantics to be aware of (documented on ``GraphRAG.update()``): @@ -48,6 +51,7 @@ from typing import TYPE_CHECKING, Any from graphrag_sdk.core.context import Context +from graphrag_sdk.core.exceptions import DatabaseError, LatencyBudgetExceededError from graphrag_sdk.core.models import ( EntityMention, GraphData, @@ -279,6 +283,12 @@ async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> old_by_hash, all_old_ids = await self._old_chunks_by_hash( {c.uid for c in chunks.chunks} ) + except (DatabaseError, LatencyBudgetExceededError): + # Fail-open covers a cache that can't answer, not a graph that + # isn't there. Extracting every chunk needs only the LLM, so it + # would succeed and bill in full before the write phase hit the + # same dead connection and failed anyway. + raise except Exception as exc: logger.warning("Chunk cache lookup failed, extracting everything: %s", exc) old_by_hash, all_old_ids = {}, set() @@ -308,6 +318,10 @@ async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> [(old_id, chunk.uid) for chunk, old_id in cached], all_old_ids ) ) + except (DatabaseError, LatencyBudgetExceededError): + # Same reasoning as the lookup above: an unreachable graph + # can't be recovered from by doing more LLM work. + raise except Exception as exc: # Cache miss must never lose data — fall back to real extraction. logger.warning( From 3667777f54a14aa6d8511ac1bc1dcd671e1e7a7b Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:22:31 +0300 Subject: [PATCH 18/21] test: pin database failures as DatabaseError, not LLM calls Covers every driver exit path, that the message keeps the driver's text, that a failed connect reuses its pool instead of stranding one, and that the cache re-raises DatabaseError rather than falling back. --- .../tests/test_cached_chunk_extraction.py | 70 ++++++++++ graphrag_sdk/tests/test_connection.py | 124 ++++++++++++++++++ 2 files changed, 194 insertions(+) diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py index 22ecdc94..c82cddd3 100644 --- a/graphrag_sdk/tests/test_cached_chunk_extraction.py +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -14,6 +14,7 @@ import pytest from graphrag_sdk.core.context import Context +from graphrag_sdk.core.exceptions import DatabaseError, LatencyBudgetExceededError from graphrag_sdk.core.models import ( ChunkEntityRow, ChunkRelationshipRow, @@ -903,3 +904,72 @@ async def test_genuine_old_chunk_still_cached_alongside_self_reference( assert [n.id for n in result.nodes] == ["acme__organization"] # Provenance for the cached chunk is remapped onto the new uid. assert result.nodes[0].properties["source_chunk_ids"] == ["new-2"] + + +class TestInfrastructureErrorsPropagate: + """Fail-open covers a cache that can't answer, not a graph that isn't there. + + Extraction only needs the LLM, so on an unreachable graph it succeeds and + bills in full, then the write phase hits the same dead connection and + fails anyway. The error is better raised before that spend. + """ + + async def test_lookup_database_error_propagates(self, ontology, ctx): + store = FakeGraphStore() + + async def _boom(document_id): + raise DatabaseError("connection refused") + + store.get_document_chunk_texts = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + with pytest.raises(DatabaseError): + await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) + assert inner.calls == [], "extraction must not run once the graph is unreachable" + + async def test_lookup_budget_error_propagates(self, ontology, ctx): + store = FakeGraphStore() + + async def _boom(document_id): + raise LatencyBudgetExceededError("out of time") + + store.get_document_chunk_texts = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + with pytest.raises(LatencyBudgetExceededError): + await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) + assert inner.calls == [] + + async def test_rebuild_database_error_propagates(self, ontology, ctx): + """Same rule at the second fallback, which runs after a cache hit.""" + store = FakeGraphStore(chunk_texts=[("old-1", "a")]) + + async def _boom(chunk_ids): + raise DatabaseError("connection reset") + + store.get_entities_mentioned_in_chunks = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + with pytest.raises(DatabaseError): + await strategy.extract( + TextChunks(chunks=[_chunk("a", uid="new-1")]), ontology, ctx + ) + assert inner.calls == [] + + async def test_other_errors_still_fail_open(self, ontology, ctx): + """A cache that merely misbehaves must not block the update.""" + store = FakeGraphStore() + + async def _boom(document_id): + raise ValueError("malformed row") + + store.get_document_chunk_texts = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) + assert len(inner.calls) == 1, "fail-open path must still extract" + assert strategy.cached_chunk_count == 0 diff --git a/graphrag_sdk/tests/test_connection.py b/graphrag_sdk/tests/test_connection.py index b981e660..b113def0 100644 --- a/graphrag_sdk/tests/test_connection.py +++ b/graphrag_sdk/tests/test_connection.py @@ -6,6 +6,7 @@ import pytest from graphrag_sdk.core.connection import ConnectionConfig, FalkorDBConnection +from graphrag_sdk.core.exceptions import DatabaseError class TestConnectionConfig: @@ -198,3 +199,126 @@ def test_pool_passes_ssl_kwargs_when_enabled(self): assert kwargs["ssl_certfile"] == "/etc/client.pem" assert kwargs["ssl_keyfile"] == "/etc/client.key" assert kwargs["ssl_check_hostname"] is True + + +class TestDriverErrorsBecomeDatabaseError: + """Driver-level failures reach callers as ``DatabaseError``. + + Callers decide what to do about an unreachable database — the chunk + extraction cache, for one, stops instead of paying for an LLM run whose + result it could not store. That decision needs a database failure to be + recognisable, so it cannot be left as a ``redis`` exception that only + reads as a generic ``Exception``. + """ + + def test_unreachable_server_on_connect(self): + """The driver contacts the server while being constructed, so an + unreachable database surfaces here rather than at query time.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + conn = FalkorDBConnection(ConnectionConfig(host="h", port=1)) + with patch("redis.asyncio.BlockingConnectionPool"), \ + patch("falkordb.asyncio.FalkorDB") as mock_falkor: + mock_falkor.side_effect = RedisConnectionError("Connection refused") + with pytest.raises(DatabaseError, match="Could not connect to FalkorDB"): + conn._ensure_client() + + def test_failed_connect_leaves_no_half_built_client(self): + """A later attempt must reconnect rather than reuse the wreckage.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + conn = FalkorDBConnection(ConnectionConfig(host="h", port=1)) + with patch("redis.asyncio.BlockingConnectionPool"), \ + patch("falkordb.asyncio.FalkorDB") as mock_falkor: + mock_falkor.side_effect = RedisConnectionError("Connection refused") + with pytest.raises(DatabaseError): + conn._ensure_client() + + assert conn._driver is None + assert conn._graph is None + + def test_failed_connect_keeps_the_pool_closeable(self): + """Only close() can release a pool, so a failed attempt must not + strand one — and a retry must not overwrite it with a second.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + conn = FalkorDBConnection(ConnectionConfig(host="h", port=1)) + with patch("redis.asyncio.BlockingConnectionPool") as mock_pool, \ + patch("falkordb.asyncio.FalkorDB") as mock_falkor: + mock_falkor.side_effect = RedisConnectionError("Connection refused") + for _ in range(3): + with pytest.raises(DatabaseError): + conn._ensure_client() + + assert mock_pool.call_count == 1, "each retry allocated another pool" + assert conn._pool is mock_pool.return_value + + def test_missing_falkordb_package_still_raises_import_error(self): + """A packaging problem is not a database problem.""" + conn = FalkorDBConnection() + with patch.dict("sys.modules", {"falkordb.asyncio": None, "falkordb": None}): + with pytest.raises(ImportError): + conn._ensure_client() + + async def test_query_wraps_error_once_retries_are_spent(self): + from redis.exceptions import ConnectionError as RedisConnectionError + + conn = FalkorDBConnection(ConnectionConfig(retry_count=2, retry_delay=0.001)) + mock_graph = MagicMock() + mock_graph.query = AsyncMock(side_effect=RedisConnectionError("Connection reset")) + conn._driver = MagicMock() + conn._graph = mock_graph + + with pytest.raises(DatabaseError) as excinfo: + await conn.query("MATCH (n) RETURN n") + + assert mock_graph.query.await_count == 2 + assert isinstance(excinfo.value.__cause__, RedisConnectionError) + + async def test_query_wraps_permanent_error_without_retrying(self): + from redis.exceptions import ResponseError + + conn = FalkorDBConnection(ConnectionConfig(retry_count=3, retry_delay=0.001)) + mock_graph = MagicMock() + mock_graph.query = AsyncMock(side_effect=ResponseError("Attribute already indexed")) + conn._driver = MagicMock() + conn._graph = mock_graph + + with pytest.raises(DatabaseError): + await conn.query("CREATE INDEX ...") + + assert mock_graph.query.await_count == 1 + + async def test_wrapped_message_keeps_the_driver_text(self): + """Index creation tolerates 'already indexed' by reading the message, + so wrapping must not hide what the driver said.""" + from redis.exceptions import ResponseError + + conn = FalkorDBConnection(ConnectionConfig(retry_count=1, retry_delay=0.001)) + mock_graph = MagicMock() + mock_graph.query = AsyncMock(side_effect=ResponseError("Attribute already indexed")) + conn._driver = MagicMock() + conn._graph = mock_graph + + with pytest.raises(DatabaseError, match="already indexed"): + await conn.query("CREATE INDEX ...") + + async def test_open_circuit_breaker(self): + conn = FalkorDBConnection() + conn._driver = MagicMock() + conn._graph = MagicMock() + conn._breaker.allow_request = AsyncMock(return_value=False) + + with pytest.raises(DatabaseError, match="Circuit breaker is open"): + await conn.query("MATCH (n) RETURN n") + + async def test_ping_reports_unreachable_instead_of_raising(self): + """Liveness checks read the bool; an unreachable server is the + case ping() exists to report, not one it should raise on.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + conn = FalkorDBConnection(ConnectionConfig(host="h", port=1)) + with patch("redis.asyncio.BlockingConnectionPool"), \ + patch("falkordb.asyncio.FalkorDB") as mock_falkor: + mock_falkor.side_effect = RedisConnectionError("Connection refused") + assert await conn.ping() is False From 7b1088ff2611de23b8b0045bb31f3113bb2ce33f Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:23:34 +0300 Subject: [PATCH 19/21] docs: correct the fail-open claim and note the breaking change The cache no longer falls back on *any* failure, and wrapping driver errors in DatabaseError changes what callers must catch. --- CHANGELOG.md | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bc14b87..f74dac01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,9 +26,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 unchanged. - **`CachedChunkExtraction(inner, graph_store, document_id)`** — the underlying decorator `ExtractionStrategy`, exported at top level for - advanced pipelines. Fail-open by construction: any cache lookup or - rebuild failure falls back to full extraction (worst case is paying - for skippable LLM calls, never data loss). Caveats: the graph is the + advanced pipelines. Fail-open: a cache lookup or rebuild failure + falls back to full extraction (worst case is paying for skippable + LLM calls, never data loss). An unreachable graph (`DatabaseError`) + and an exhausted latency budget are the exceptions and propagate — + extraction only needs the LLM, so it would bill in full and then + fail at the write phase regardless. Caveats: the graph is the cache, so manually deleted entities are not resurrected from unchanged chunks, and ontology/prompt/model changes do not re-extract unchanged chunks — pass `cache_unchanged_chunks=False` @@ -39,6 +42,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 accessors backing the cache (with `ChunkEntityRow` / `ChunkRelationshipRow` typed rows in `core.models`). +### Changed + +- **`FalkorDBConnection.query()` now raises `DatabaseError` for every + driver-level failure** — unreachable server, open circuit breaker, + exhausted retries, or a permanent error. Previously the raw + `redis.exceptions.*` exception escaped, which callers could only + catch as a generic `Exception`; telling an infrastructure failure + apart from an application error is what lets the chunk cache stop + instead of paying for an LLM run it could not store. The originating + exception is preserved as `__cause__` and its text is kept in the + message, so message-based handling (such as tolerating "index + already exists") is unaffected. Catch `DatabaseError` instead of + `redis.exceptions.ConnectionError` if you were relying on the + driver's own types. +- **`FalkorDBConnection.ping()` returns `False` when the server is + unreachable** instead of raising. Reporting "not alive" is the point + of the call. + ## [1.3.0] - 2026-06-04 Ontology discovery (#271): bootstrap an ontology straight from a From 60f04af12a885cc64b71e62dfdd64048b48bb152 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Sun, 9 Aug 2026 13:40:13 +0300 Subject: [PATCH 20/21] fix: fall back to extraction when the cache cannot rebuild a chunk Two cache-rebuild paths could silently delete live data. A label-less entity was skipped with a warning, on the assumption the entity survived. It does, but IngestionPipeline._filter_quality drops every relationship incident to a node absent from the batch, so the rebuilt provenance never landed and the post-cutover sweep then deleted edges the unchanged chunk still supported. Re-ingesting an unmodified document lost facts. A stored relationship whose endpoint no cached chunk mentions fails the same way. The pipeline only ever persists an edge alongside its endpoints, so this needs a graph written by something else, but the outcome is identical. Neither case can reproduce the graph, so refuse to rebuild and let the existing fallback hand the document to real extraction, which is always correct. Both now also surface in cache_stats as extracted chunks rather than disappearing into a log line. --- .../cached_chunk_extraction.py | 42 ++++++++-- .../tests/test_cached_chunk_extraction.py | 84 ++++++++++++++++++- 2 files changed, 115 insertions(+), 11 deletions(-) diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py index 55d0ad2e..3f1d7476 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -28,8 +28,10 @@ - Mentions are re-emitted against the new chunk uid, so the ``update()`` orphan cleanup keeps every entity that unchanged chunks still support. - A cache failure falls back to real extraction (fail-open): the worst case is - paying for LLM calls that could have been skipped, never data loss. The - exception is an unreachable graph (``DatabaseError``) or an exhausted + paying for LLM calls that could have been skipped, never data loss. This + includes a stored entity that carries no concrete label, which cannot be + re-emitted without either minting a duplicate node or losing its edges. + The exception is an unreachable graph (``DatabaseError``) or an exhausted latency budget, which propagate — extraction only needs the LLM, so it would bill in full and then fail at the write phase regardless. @@ -69,6 +71,14 @@ logger = logging.getLogger(__name__) +class _UnrebuildableCacheEntry(Exception): + """Raised when cached data cannot be faithfully re-emitted. + + Internal to this module: ``extract()`` treats it like any other cache + failure and falls back to real extraction, which is always correct. + """ + + def _sha256(text: str) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() @@ -192,23 +202,41 @@ async def _graph_data_from_cache( for eid, props in ent_props.items(): label = ent_labels[eid] if not label: - # No usable label: upsert_nodes would MERGE on a fallback + # No usable label, so upsert_nodes would MERGE on a fallback # label and mint a duplicate node instead of matching this - # one. Skip the node write — the mention above still keeps - # the entity alive through orphan cleanup. - logger.warning("Skipping cache node write for label-less entity %s", eid) - continue + # one. Skipping just this node is not safe either: the + # pipeline's quality filter drops every relationship + # incident to a node it never saw, so the rebuilt provenance + # never lands and the post-cutover sweep then deletes edges + # the unchanged chunk still supports. Neither branch can + # reproduce the graph, so refuse to rebuild and let the + # caller fall back to real extraction. + raise _UnrebuildableCacheEntry(f"entity {eid!r} carries no concrete label") props["source_chunk_ids"] = ent_sources[eid] nodes.append(GraphNode(id=eid, label=label, properties=props)) relationships: list[GraphRelationship] = [] if ent_props: + node_ids = {n.id for n in nodes} rel_by_pair: dict[tuple[str, str], dict[str, Any]] = {} for rel in await self._graph_store.get_relationships_for_chunks(old_ids): new_uids = id_map.get(rel.chunk_id) if not new_uids: continue key = (rel.start_entity_id, rel.end_entity_id) + if key[0] not in node_ids or key[1] not in node_ids: + # The pipeline only persists an edge alongside the nodes + # it was extracted with (see IngestionPipeline._filter_ + # quality), so a stored edge whose endpoints this chunk + # does not mention means the graph was written by + # something else. Re-emitting it would be dropped by that + # same filter, losing the provenance and letting the + # post-cutover sweep delete a live fact — so hand the + # document to real extraction instead. + raise _UnrebuildableCacheEntry( + f"relationship {key[0]!r}->{key[1]!r} has an endpoint that " + "no cached chunk mentions" + ) props = rel_by_pair.get(key) if props is None: props = {"source_chunk_ids": []} diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py index c82cddd3..c8bd7f35 100644 --- a/graphrag_sdk/tests/test_cached_chunk_extraction.py +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -353,7 +353,12 @@ async def test_relationship_between_same_pair_deduped(self, ontology, ctx): class TestLabelLessEntity: - async def test_label_less_entity_skips_node_but_keeps_mention(self, ontology, ctx): + async def test_label_less_entity_falls_back_to_extraction(self, ontology, ctx): + """A stored entity with no concrete label cannot be re-emitted: + writing it would mint a duplicate node, and skipping it makes the + pipeline's quality filter drop its relationships, which the + post-cutover sweep then deletes. Rebuild must give way to real + extraction rather than lose edges.""" store = FakeGraphStore( chunk_texts=[("old-1", "alpha")], entity_rows=[ @@ -366,14 +371,85 @@ async def test_label_less_entity_skips_node_but_keeps_mention(self, ontology, ct ) ], ) - strategy = CachedChunkExtraction(RecordingExtractor(), store, "doc-1") + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx) + + assert len(inner.calls) == 1 + assert [c.uid for c in inner.calls[0].chunks] == ["new-1"] + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 1 + + async def test_relationship_endpoint_not_mentioned_falls_back(self, ontology, ctx): + """A stored edge whose endpoint no cached chunk mentions cannot be + re-emitted: IngestionPipeline._filter_quality drops relationships + whose endpoints are absent from the same batch, so the rebuilt + provenance would never land and the post-cutover sweep would then + delete an edge the unchanged chunk still supports.""" + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="known", + label="Person", + name="Known", + source_chunk_ids=["old-1"], + ) + ], + rel_rows=[ + ChunkRelationshipRow( + chunk_id="old-1", + start_entity_id="known", + end_entity_id="stranger", # never mentioned by old-1 + rel_type="knows", + ) + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx) + + assert len(inner.calls) == 1 + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 1 + + async def test_relationship_with_both_endpoints_mentioned_is_cached(self, ontology, ctx): + """The guard must not fire on the normal case.""" + store = FakeGraphStore( + chunk_texts=[("old-1", "alpha")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id=eid, + label="Person", + name=eid, + source_chunk_ids=["old-1"], + ) + for eid in ("a", "b") + ], + rel_rows=[ + ChunkRelationshipRow( + chunk_id="old-1", + start_entity_id="a", + end_entity_id="b", + rel_type="knows", + ) + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") result = await strategy.extract( TextChunks(chunks=[_chunk("alpha", 0, "new-1")]), ontology, ctx ) - assert result.nodes == [] # no node write — would mint a duplicate - assert result.mentions == [EntityMention(chunk_id="new-1", entity_id="ghost")] + assert inner.calls == [] + assert strategy.cached_chunk_count == 1 + assert len(result.relationships) == 1 + assert result.relationships[0].properties["source_chunk_ids"] == ["new-1"] class TestFailOpen: From 8b057b7a3619b7c7d044ca989381694da90fdb00 Mon Sep 17 00:00:00 2001 From: Naseem Ali <34807727+Naseem77@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:05:54 +0300 Subject: [PATCH 21/21] fix: union cached and freshly extracted relationships on merge _merge concatenated relationship lists, so a chunk that was rebuilt from the cache and a chunk that was re-extracted could both emit the same (start, type, end) edge. ExactMatchResolution keeps the first occurrence and rebuilds the edge from it, dropping the rest, and cached parts are appended before extracted ones - so a stale fact won over its correction and the new chunk id never reached source_chunk_ids. A later update then emptied the list and the sweep deleted the edge. Full re-extraction does not have this problem because _aggregate_relations collapses duplicates upstream with provenance unioned. Relationships are now keyed on (start_node_id, type, end_node_id) the same way nodes already were: later parts win on properties, and source_chunk_ids is unioned across all occurrences. Also: - Split DatabaseUnavailableError out of DatabaseError. Wrapping every driver failure as DatabaseError made the cache re-raise on a query the server had rejected, turning a case that falls open cleanly into a hard failure. Only a server that never answered - unreachable, open circuit breaker, exhausted retries - propagates now. The new type subclasses DatabaseError, so existing handlers are unchanged. query()'s permanent-error check also recognizes syntax errors, invalid input, unknown functions, type mismatches and missing procedures, so a rejected query fails fast instead of spending the whole retry budget first. - Scope cache fallback per chunk. An entity whose label is missing, or an edge whose endpoint is not mentioned in the citing chunk, voided the entire document's cache. Only the affected chunks are re-extracted now, computed as a fixpoint - dropping a chunk removes entities only it mentioned, which can strand edges other chunks cite. - ping() re-raises ImportError instead of reporting a missing falkordb package as a server that is down. - Correct a stale docstring on get_entities_mentioned_in_chunks. --- CHANGELOG.md | 28 +- .../src/graphrag_sdk/core/connection.py | 25 +- .../src/graphrag_sdk/core/exceptions.py | 16 + .../cached_chunk_extraction.py | 208 +++++++++---- .../src/graphrag_sdk/storage/graph_store.py | 8 +- .../tests/test_cached_chunk_extraction.py | 283 ++++++++++++++---- 6 files changed, 435 insertions(+), 133 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f74dac01..03696d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,11 +27,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **`CachedChunkExtraction(inner, graph_store, document_id)`** — the underlying decorator `ExtractionStrategy`, exported at top level for advanced pipelines. Fail-open: a cache lookup or rebuild failure - falls back to full extraction (worst case is paying for skippable - LLM calls, never data loss). An unreachable graph (`DatabaseError`) - and an exhausted latency budget are the exceptions and propagate — - extraction only needs the LLM, so it would bill in full and then - fail at the write phase regardless. Caveats: the graph is the + falls back to extraction (worst case is paying for skippable + LLM calls, never data loss). Fallback is scoped per chunk — a single + chunk the cache cannot rebuild is re-extracted on its own, and any + chunk transitively depending on it, rather than voiding the whole + document's cache. An unreachable graph + (`DatabaseUnavailableError`) and an exhausted latency budget are the + exceptions and propagate — extraction only needs the LLM, so it would + bill in full and then fail at the write phase regardless. A query the + server *rejects* (`DatabaseError`) still falls open, because the + server answered and the write phase will work. Caveats: the graph is the cache, so manually deleted entities are not resurrected from unchanged chunks, and ontology/prompt/model changes do not re-extract unchanged chunks — pass `cache_unchanged_chunks=False` @@ -58,7 +63,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 driver's own types. - **`FalkorDBConnection.ping()` returns `False` when the server is unreachable** instead of raising. Reporting "not alive" is the point - of the call. + of the call. A missing `falkordb` package still raises `ImportError` + — a broken install is a packaging bug, not a down server, and + reporting it as one sends operators to the wrong place. +- **New `DatabaseUnavailableError(DatabaseError)`** distinguishes "the + server never answered" (unreachable, open circuit breaker, exhausted + retries) from "the server answered and rejected the query" (plain + `DatabaseError`). Existing `except DatabaseError` handlers are + unaffected. The chunk cache uses the split to decide whether falling + open is safe. `query()`'s permanent-error short-circuit now also + recognizes syntax errors, invalid input, unknown functions, type + mismatches, and missing procedures, so a rejected query fails fast + instead of burning the full retry budget first. ## [1.3.0] - 2026-06-04 diff --git a/graphrag_sdk/src/graphrag_sdk/core/connection.py b/graphrag_sdk/src/graphrag_sdk/core/connection.py index a367b5ef..7d20de49 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/connection.py +++ b/graphrag_sdk/src/graphrag_sdk/core/connection.py @@ -11,7 +11,7 @@ from typing import Any from urllib.parse import urlparse -from graphrag_sdk.core.exceptions import DatabaseError +from graphrag_sdk.core.exceptions import DatabaseError, DatabaseUnavailableError logger = logging.getLogger(__name__) @@ -137,7 +137,7 @@ def _ensure_client(self) -> None: # a pool per failed attempt with no way to aclose() it. self._driver = None self._graph = None - raise DatabaseError( + raise DatabaseUnavailableError( f"Could not connect to FalkorDB at {self.config.host}:{self.config.port}: {exc}" ) from exc @@ -183,7 +183,7 @@ async def query( assert self._graph is not None # for type-checkers if not await self._breaker.allow_request(): - raise DatabaseError( + raise DatabaseUnavailableError( "Circuit breaker is open — FalkorDB connection is unhealthy. " "Requests will resume after recovery timeout." ) @@ -230,7 +230,7 @@ async def query( "FalkorDB query failure details", exc_info=(type(last_exc), last_exc, last_exc.__traceback__), ) - raise DatabaseError( + raise DatabaseUnavailableError( f"FalkorDB query failed after {self.config.retry_count} attempts: {last_exc}" ) from last_exc raise DatabaseError("FalkorDB query failed without an exception") @@ -241,6 +241,14 @@ async def query( "already indexed", "already exists", "unknown index", + # Malformed or unsupported query: the server answered and rejected + # it. Retrying re-sends the same bytes, so it can only fail again — + # and exhausting the budget would misreport a rejection as an outage. + "syntax error", + "invalid input", + "unknown function", + "type mismatch", + "procedure not found", ) @classmethod @@ -254,7 +262,9 @@ async def ping(self) -> bool: """Send a Redis PING to verify the connection is alive. Returns ``False`` rather than raising when the server cannot be - reached — reporting "not alive" is the point of the call. + reached — reporting "not alive" is the point of the call. An + ``ImportError`` from a missing driver still propagates: that is a + broken install, not a dead server. """ try: self._ensure_client() @@ -263,6 +273,11 @@ async def ping(self) -> bool: redis: Redis = Redis(connection_pool=self._pool) return await redis.ping() + except ImportError: + # The driver isn't installed. Reporting that as "server down" + # sends operators to inspect a healthy database instead of the + # broken install, so let it surface. + raise except Exception: logger.debug("Ping failed", exc_info=True) return False diff --git a/graphrag_sdk/src/graphrag_sdk/core/exceptions.py b/graphrag_sdk/src/graphrag_sdk/core/exceptions.py index 99c3c379..677d646c 100644 --- a/graphrag_sdk/src/graphrag_sdk/core/exceptions.py +++ b/graphrag_sdk/src/graphrag_sdk/core/exceptions.py @@ -94,6 +94,22 @@ class DatabaseError(GraphRAGError): pass +class DatabaseUnavailableError(DatabaseError): + """Raised when FalkorDB cannot be reached at all. + + Distinct from a plain :class:`DatabaseError`, which means the server + answered and *rejected* the query. Callers that fall back to a slower + path on cache/read failures need that distinction: a rejected query + leaves the rest of the pipeline perfectly usable, while an unreachable + server will fail the write phase no matter how much work is done first. + + Subclasses ``DatabaseError``, so existing ``except DatabaseError`` + handlers keep working unchanged. + """ + + pass + + class IndexError_(GraphRAGError): """Raised when index creation/management fails. diff --git a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py index 3f1d7476..8589d7b3 100644 --- a/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -28,12 +28,16 @@ - Mentions are re-emitted against the new chunk uid, so the ``update()`` orphan cleanup keeps every entity that unchanged chunks still support. - A cache failure falls back to real extraction (fail-open): the worst case is - paying for LLM calls that could have been skipped, never data loss. This - includes a stored entity that carries no concrete label, which cannot be - re-emitted without either minting a duplicate node or losing its edges. - The exception is an unreachable graph (``DatabaseError``) or an exhausted - latency budget, which propagate — extraction only needs the LLM, so it - would bill in full and then fail at the write phase regardless. + paying for LLM calls that could have been skipped, never data loss. The + fallback is per chunk where possible — a chunk that cannot be rebuilt + faithfully (a stored entity carrying no concrete label, or an edge whose + endpoint no cached chunk mentions) is extracted on its own, leaving the + rest of the document cached. + The exception is an unreachable graph (``DatabaseUnavailableError``) or an + exhausted latency budget, which propagate — extraction only needs the LLM, + so it would bill in full and then fail at the write phase regardless. A + plain ``DatabaseError`` (the server answered and rejected the query) does + fall open: the rest of the pipeline is unaffected. Semantics to be aware of (documented on ``GraphRAG.update()``): @@ -53,7 +57,7 @@ from typing import TYPE_CHECKING, Any from graphrag_sdk.core.context import Context -from graphrag_sdk.core.exceptions import DatabaseError, LatencyBudgetExceededError +from graphrag_sdk.core.exceptions import DatabaseUnavailableError, LatencyBudgetExceededError from graphrag_sdk.core.models import ( EntityMention, GraphData, @@ -71,14 +75,6 @@ logger = logging.getLogger(__name__) -class _UnrebuildableCacheEntry(Exception): - """Raised when cached data cannot be faithfully re-emitted. - - Internal to this module: ``extract()`` treats it like any other cache - failure and falls back to real extraction, which is always correct. - """ - - def _sha256(text: str) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() @@ -140,9 +136,53 @@ async def _old_chunks_by_hash(self, exclude: set[str]) -> tuple[dict[str, str], by_hash.setdefault(_sha256(text), cid) return by_hash, all_ids + @staticmethod + def _usable_chunks( + old_ids: list[str], + ent_rows: list[Any], + rel_rows: list[Any], + ) -> set[str]: + """Largest subset of cached chunks that can be rebuilt faithfully. + + A chunk is unusable when it mentions an entity carrying no concrete + label (``upsert_nodes`` would MERGE on a fallback label and mint a + duplicate node), or when it cites a relationship whose endpoint no + remaining cached chunk mentions. In both cases + ``IngestionPipeline._filter_quality`` would drop the affected edges, + the rebuilt provenance would never land, and the post-cutover sweep + would then delete facts the unchanged chunk still supports. + + Dropping a chunk removes the entities only it mentioned, which can + strand edges cited by other chunks, so this iterates to a fixpoint. + Only the affected chunks are excluded — the rest of the document + still gets its cache. + """ + ents_by_chunk: dict[str, set[str]] = {} + label_by_ent: dict[str, str | None] = {} + for row in ent_rows: + ents_by_chunk.setdefault(row.chunk_id, set()).add(row.entity_id) + if not label_by_ent.get(row.entity_id): + label_by_ent[row.entity_id] = row.label + + good = set(old_ids) + while good: + node_ids = {e for c in good for e in ents_by_chunk.get(c, ())} + bad: set[str] = set() + for c in good: + if any(not label_by_ent.get(e) for e in ents_by_chunk.get(c, ())): + bad.add(c) + for rel in rel_rows: + if rel.chunk_id in good and rel.chunk_id not in bad: + if rel.start_entity_id not in node_ids or rel.end_entity_id not in node_ids: + bad.add(rel.chunk_id) + if not bad: + break + good -= bad + return good + async def _graph_data_from_cache( self, pairs: list[tuple[str, str]], all_old_ids: set[str] - ) -> GraphData: + ) -> tuple[GraphData, set[str]]: """Rebuild GraphData for all unchanged chunks from the live graph. Args: @@ -152,6 +192,11 @@ async def _graph_data_from_cache( mentions and provenance. all_old_ids: every pre-update chunk id of this document — used to drop provenance pointing at chunks that die in the cutover. + + Returns: + ``(graph_data, unusable_old_ids)``. Chunks in ``unusable_old_ids`` + could not be rebuilt faithfully and are excluded from the result; + the caller sends their new chunks to real extraction instead. """ # old chunk id -> ALL new chunk uids that reuse its extraction. id_map: dict[str, list[str]] = {} @@ -159,12 +204,29 @@ async def _graph_data_from_cache( id_map.setdefault(old_id, []).append(new_uid) old_ids = list(id_map) + ent_rows = await self._graph_store.get_entities_mentioned_in_chunks(old_ids) + rel_rows = await self._graph_store.get_relationships_for_chunks(old_ids) + + good = self._usable_chunks(old_ids, ent_rows, rel_rows) + unusable = set(old_ids) - good + if unusable: + logger.warning( + "Cache cannot faithfully rebuild %d of %d unchanged chunk(s); " + "extracting those instead", + len(unusable), + len(old_ids), + ) + if not good: + return GraphData(nodes=[], relationships=[], mentions=[]), unusable + + id_map = {old: uids for old, uids in id_map.items() if old in good} + mentions: list[EntityMention] = [] ent_props: dict[str, dict[str, Any]] = {} ent_labels: dict[str, str | None] = {} ent_sources: dict[str, list[str]] = {} - for row in await self._graph_store.get_entities_mentioned_in_chunks(old_ids): + for row in ent_rows: new_uids = id_map.get(row.chunk_id) if not new_uids: continue @@ -202,16 +264,11 @@ async def _graph_data_from_cache( for eid, props in ent_props.items(): label = ent_labels[eid] if not label: - # No usable label, so upsert_nodes would MERGE on a fallback - # label and mint a duplicate node instead of matching this - # one. Skipping just this node is not safe either: the - # pipeline's quality filter drops every relationship - # incident to a node it never saw, so the rebuilt provenance - # never lands and the post-cutover sweep then deletes edges - # the unchanged chunk still supports. Neither branch can - # reproduce the graph, so refuse to rebuild and let the - # caller fall back to real extraction. - raise _UnrebuildableCacheEntry(f"entity {eid!r} carries no concrete label") + # _usable_chunks already excluded every chunk mentioning a + # label-less entity, so reaching here means the row set + # changed underneath us. Skip defensively rather than mint a + # duplicate node on a fallback label. + continue props["source_chunk_ids"] = ent_sources[eid] nodes.append(GraphNode(id=eid, label=label, properties=props)) @@ -219,24 +276,17 @@ async def _graph_data_from_cache( if ent_props: node_ids = {n.id for n in nodes} rel_by_pair: dict[tuple[str, str], dict[str, Any]] = {} - for rel in await self._graph_store.get_relationships_for_chunks(old_ids): + for rel in rel_rows: new_uids = id_map.get(rel.chunk_id) if not new_uids: continue key = (rel.start_entity_id, rel.end_entity_id) if key[0] not in node_ids or key[1] not in node_ids: - # The pipeline only persists an edge alongside the nodes - # it was extracted with (see IngestionPipeline._filter_ - # quality), so a stored edge whose endpoints this chunk - # does not mention means the graph was written by - # something else. Re-emitting it would be dropped by that - # same filter, losing the provenance and letting the - # post-cutover sweep delete a live fact — so hand the - # document to real extraction instead. - raise _UnrebuildableCacheEntry( - f"relationship {key[0]!r}->{key[1]!r} has an endpoint that " - "no cached chunk mentions" - ) + # Same as above: the fixpoint guarantees both endpoints + # are present, so this is unreachable for a consistent + # read. Dropping the edge here is safe because the chunk + # that cites it would already have been excluded. + continue props = rel_by_pair.get(key) if props is None: props = {"source_chunk_ids": []} @@ -259,46 +309,70 @@ async def _graph_data_from_cache( for (s, e), p in rel_by_pair.items() ] - return GraphData(nodes=nodes, relationships=relationships, mentions=mentions) + return GraphData(nodes=nodes, relationships=relationships, mentions=mentions), unusable @staticmethod def _merge(parts: list[GraphData]) -> GraphData: """Merge cached and freshly extracted GraphData. - Later parts win on conflicting node properties (fresh extraction - follows cached parts, so updated descriptions take precedence) — - EXCEPT ``source_chunk_ids``, which is a union: an entity present in + Later parts win on conflicting properties (fresh extraction follows + cached parts, so updated descriptions take precedence) — EXCEPT + ``source_chunk_ids``, which is a union: an entity or fact present in both a cached and an extracted chunk must keep both provenances. + + Relationships are keyed on ``(start, type, end)`` for the same reason + nodes are keyed on id. Concatenating them instead would leave two + rows for one fact, and ``ExactMatchResolution`` keeps only the first + and rebuilds from it — discarding the fresh part's corrected + properties *and* its provenance, which the post-cutover sweep would + later collect as an unsupported edge. Full re-extraction never hits + that because ``_aggregate_relations`` collapses duplicates upstream; + merging here keeps the cached path equivalent. """ nodes_by_id: dict[str, GraphNode] = {} - relationships: list[GraphRelationship] = [] + rels_by_key: dict[tuple[str, str, str], GraphRelationship] = {} mentions: list[EntityMention] = [] extracted_entities = [] extracted_relations = [] + + def _union_sources(old_props: dict[str, Any], new_props: dict[str, Any]) -> dict[str, Any]: + merged = {**old_props, **new_props} + old_src = old_props.get("source_chunk_ids") or [] + new_src = new_props.get("source_chunk_ids") or [] + if old_src or new_src: + merged["source_chunk_ids"] = list(old_src) + [ + c for c in new_src if c not in old_src + ] + return merged + for part in parts: for node in part.nodes: existing = nodes_by_id.get(node.id) if existing is None: nodes_by_id[node.id] = node else: - old_props = existing.properties or {} - new_props = node.properties or {} - merged = {**old_props, **new_props} - old_src = old_props.get("source_chunk_ids") or [] - new_src = new_props.get("source_chunk_ids") or [] - if old_src or new_src: - merged["source_chunk_ids"] = old_src + [ - c for c in new_src if c not in old_src - ] + merged = _union_sources(existing.properties or {}, node.properties or {}) label = node.label or existing.label nodes_by_id[node.id] = GraphNode(id=node.id, label=label, properties=merged) - relationships.extend(part.relationships) + for rel in part.relationships: + key = (rel.start_node_id, rel.type, rel.end_node_id) + current = rels_by_key.get(key) + if current is None: + rels_by_key[key] = rel + else: + merged = _union_sources(current.properties or {}, rel.properties or {}) + rels_by_key[key] = GraphRelationship( + start_node_id=rel.start_node_id, + end_node_id=rel.end_node_id, + type=rel.type, + properties=merged, + ) mentions.extend(part.mentions) extracted_entities.extend(part.extracted_entities) extracted_relations.extend(part.extracted_relations) return GraphData( nodes=list(nodes_by_id.values()), - relationships=relationships, + relationships=list(rels_by_key.values()), mentions=mentions, extracted_entities=extracted_entities, extracted_relations=extracted_relations, @@ -311,11 +385,13 @@ async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> old_by_hash, all_old_ids = await self._old_chunks_by_hash( {c.uid for c in chunks.chunks} ) - except (DatabaseError, LatencyBudgetExceededError): + except (DatabaseUnavailableError, LatencyBudgetExceededError): # Fail-open covers a cache that can't answer, not a graph that # isn't there. Extracting every chunk needs only the LLM, so it # would succeed and bill in full before the write phase hit the - # same dead connection and failed anyway. + # same dead connection and failed anyway. A plain DatabaseError + # means the server answered and rejected the query — the rest of + # the pipeline still works, so that one falls open below. raise except Exception as exc: logger.warning("Chunk cache lookup failed, extracting everything: %s", exc) @@ -341,12 +417,18 @@ async def extract(self, chunks: TextChunks, ontology: Ontology, ctx: Context) -> parts: list[GraphData] = [] if cached: try: - parts.append( - await self._graph_data_from_cache( - [(old_id, chunk.uid) for chunk, old_id in cached], all_old_ids - ) + data, unusable = await self._graph_data_from_cache( + [(old_id, chunk.uid) for chunk, old_id in cached], all_old_ids ) - except (DatabaseError, LatencyBudgetExceededError): + if unusable: + # Only the chunks that could not be rebuilt go to the + # LLM; the rest of the document keeps its cache. + to_extract.extend(chunk for chunk, old_id in cached if old_id in unusable) + moved = sum(1 for _, old_id in cached if old_id in unusable) + self.cached_chunk_count -= moved + self.extracted_chunk_count += moved + parts.append(data) + except (DatabaseUnavailableError, LatencyBudgetExceededError): # Same reasoning as the lookup above: an unreachable graph # can't be recovered from by doing more LLM work. raise diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index d33916d0..c2cef098 100644 --- a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py +++ b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py @@ -720,10 +720,10 @@ async def get_entities_mentioned_in_chunks(self, chunk_ids: list[str]) -> list[C label (matching how ``upsert_nodes`` writes concrete label + ``__Entity__`` marker). ``None`` when no concrete label exists — deliberately NO fallback to ``e.type``: a MERGE on a label the - node doesn't carry would mint a duplicate node with the same id, - so the consumer must skip the node write instead - (``CachedChunkExtraction`` keeps the mention, which is enough to - survive orphan cleanup). + node doesn't carry would mint a duplicate node with the same id. + ``CachedChunkExtraction`` treats every chunk mentioning such an + entity as unrebuildable and sends those chunks to real extraction + instead. Batched by ``_BATCH_SIZE`` to keep parameter sizes bounded. """ diff --git a/graphrag_sdk/tests/test_cached_chunk_extraction.py b/graphrag_sdk/tests/test_cached_chunk_extraction.py index c8bd7f35..8eb20d1a 100644 --- a/graphrag_sdk/tests/test_cached_chunk_extraction.py +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -14,13 +14,18 @@ import pytest from graphrag_sdk.core.context import Context -from graphrag_sdk.core.exceptions import DatabaseError, LatencyBudgetExceededError +from graphrag_sdk.core.exceptions import ( + DatabaseError, + DatabaseUnavailableError, + LatencyBudgetExceededError, +) from graphrag_sdk.core.models import ( ChunkEntityRow, ChunkRelationshipRow, EntityMention, GraphData, GraphNode, + GraphRelationship, Ontology, TextChunk, TextChunks, @@ -52,9 +57,7 @@ def __init__( async def get_document_chunk_texts(self, document_id: str) -> list[tuple[str, str]]: return self.chunk_texts - async def get_entities_mentioned_in_chunks( - self, chunk_ids: list[str] - ) -> list[ChunkEntityRow]: + async def get_entities_mentioned_in_chunks(self, chunk_ids: list[str]) -> list[ChunkEntityRow]: wanted = set(chunk_ids) return [r for r in self.entity_rows if r.chunk_id in wanted] @@ -131,9 +134,7 @@ async def test_all_chunks_cached_inner_never_called(self, ontology, ctx): inner = RecordingExtractor() strategy = CachedChunkExtraction(inner, store, "doc-1") - new_chunks = TextChunks( - chunks=[_chunk("alpha", 0, "new-1"), _chunk("beta", 1, "new-2")] - ) + new_chunks = TextChunks(chunks=[_chunk("alpha", 0, "new-1"), _chunk("beta", 1, "new-2")]) result = await strategy.extract(new_chunks, ontology, ctx) assert inner.calls == [] @@ -496,9 +497,7 @@ async def test_stats_reset_between_calls(self, ontology, ctx): class TestMerge: - async def test_entity_in_cached_and_extracted_chunk_unions_provenance( - self, ontology, ctx - ): + async def test_entity_in_cached_and_extracted_chunk_unions_provenance(self, ontology, ctx): """Same entity in an unchanged chunk (cached) and a changed chunk (extracted): fresh properties win, provenance is the union.""" store = FakeGraphStore( @@ -548,12 +547,8 @@ async def test_extracted_entities_and_relations_pass_through(self, ontology, ctx fresh = GraphData( extracted_entities=[ExtractedEntity(name="Alice", type="Person")], ) - strategy = CachedChunkExtraction( - RecordingExtractor(fresh), FakeGraphStore(), "doc-1" - ) - result = await strategy.extract( - TextChunks(chunks=[_chunk("anything")]), ontology, ctx - ) + strategy = CachedChunkExtraction(RecordingExtractor(fresh), FakeGraphStore(), "doc-1") + result = await strategy.extract(TextChunks(chunks=[_chunk("anything")]), ontology, ctx) assert len(result.extracted_entities) == 1 @@ -564,9 +559,7 @@ class TestGraphStoreCacheAccessors: def graph_store(self, mock_connection): return GraphStore(mock_connection) - async def test_get_document_chunk_texts_filters_bad_rows( - self, graph_store, mock_connection - ): + async def test_get_document_chunk_texts_filters_bad_rows(self, graph_store, mock_connection): from unittest.mock import MagicMock mock_connection.query = AsyncMock( @@ -584,9 +577,7 @@ async def test_get_document_chunk_texts_filters_bad_rows( cypher = mock_connection.query.call_args[0][0] assert "PART_OF" in cypher and "c.text" in cypher - async def test_get_entities_mentioned_in_chunks_maps_rows( - self, graph_store, mock_connection - ): + async def test_get_entities_mentioned_in_chunks_maps_rows(self, graph_store, mock_connection): from unittest.mock import MagicMock mock_connection.query = AsyncMock( @@ -628,9 +619,7 @@ async def test_get_entities_label_none_when_no_concrete_label( rows = await graph_store.get_entities_mentioned_in_chunks(["c1"]) assert rows[0].label is None - async def test_get_relationships_for_chunks_maps_rows( - self, graph_store, mock_connection - ): + async def test_get_relationships_for_chunks_maps_rows(self, graph_store, mock_connection): from unittest.mock import MagicMock mock_connection.query = AsyncMock( @@ -653,9 +642,7 @@ async def test_get_relationships_for_chunks_maps_rows( # Single edge scan per batch — no per-chunk UNWIND re-scan. assert "UNWIND" not in cypher - async def test_get_relationships_one_row_per_matching_chunk( - self, graph_store, mock_connection - ): + async def test_get_relationships_one_row_per_matching_chunk(self, graph_store, mock_connection): from unittest.mock import MagicMock mock_connection.query = AsyncMock( @@ -671,9 +658,7 @@ async def test_get_relationships_one_row_per_matching_chunk( ("c2", "a", "b"), } - async def test_scalar_where_list_expected_treated_as_empty( - self, graph_store, mock_connection - ): + async def test_scalar_where_list_expected_treated_as_empty(self, graph_store, mock_connection): """Tampered graphs may hold a scalar where a list is expected — it must not iterate as characters into bogus labels/provenance.""" from unittest.mock import MagicMock @@ -688,9 +673,7 @@ async def test_scalar_where_list_expected_treated_as_empty( assert rows[0].label is None # scalar labels → no concrete label assert rows[0].source_chunk_ids == [] # "abc" must not become ['a','b','c'] - async def test_relationship_scalar_provenance_skipped( - self, graph_store, mock_connection - ): + async def test_relationship_scalar_provenance_skipped(self, graph_store, mock_connection): from unittest.mock import MagicMock mock_connection.query = AsyncMock( @@ -735,9 +718,7 @@ def _texts(self): # across both versions; only chunk 2 changes. FixedSizeChunking cuts # pure character windows without stripping, so padding is stable. part1 = _first_n_chars_stable( - "Alice works at Acme Corporation with her colleague Bob today.".ljust( - self.CHUNK, "|" - ), + "Alice works at Acme Corporation with her colleague Bob today.".ljust(self.CHUNK, "|"), self.CHUNK, ) v1 = part1 + "Carol manages the Berlin office of Acme Corporation." @@ -814,8 +795,7 @@ async def count(name): # Alice's MENTIONED_IN edge must target a live chunk. r = await rag._graph_store.query_raw( - "MATCH (e:__Entity__ {name: 'Alice'})-[:MENTIONED_IN]->(c:Chunk) " - "RETURN collect(c.id)" + "MATCH (e:__Entity__ {name: 'Alice'})-[:MENTIONED_IN]->(c:Chunk) RETURN collect(c.id)" ) assert set(r.result_set[0][0] or []) <= live_chunk_ids @@ -859,9 +839,7 @@ async def test_manually_deleted_entity_not_resurrected( ) assert r.result_set[0][0] == 1 - async def test_default_off_extracts_all_chunks( - self, real_falkordb_rag_factory, scripted_llm - ): + async def test_default_off_extracts_all_chunks(self, real_falkordb_rag_factory, scripted_llm): """cache_unchanged_chunks defaults to False → both chunks hit the LLM on update (4 scripted responses, all consumed).""" from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( @@ -886,9 +864,7 @@ async def test_default_off_extracts_all_chunks( ) assert r.result_set[0][0] == 1 - async def test_no_op_short_circuit_unaffected( - self, real_falkordb_rag_factory, scripted_llm - ): + async def test_no_op_short_circuit_unaffected(self, real_falkordb_rag_factory, scripted_llm): from graphrag_sdk.ingestion.resolution_strategies.exact_match import ( ExactMatchResolution, ) @@ -948,9 +924,7 @@ async def test_self_referential_chunk_is_extracted_not_cached(self, ontology, ct assert [n.id for n in result.nodes] == ["alice__person"] assert len(result.mentions) == 1 - async def test_genuine_old_chunk_still_cached_alongside_self_reference( - self, ontology, ctx - ): + async def test_genuine_old_chunk_still_cached_alongside_self_reference(self, ontology, ctx): """The guard must drop only the self-reference, not real cache hits.""" store = FakeGraphStore( chunk_texts=[("new-1", "alpha"), ("old-2", "beta")], @@ -994,16 +968,34 @@ async def test_lookup_database_error_propagates(self, ontology, ctx): store = FakeGraphStore() async def _boom(document_id): - raise DatabaseError("connection refused") + raise DatabaseUnavailableError("connection refused") store.get_document_chunk_texts = _boom inner = RecordingExtractor() strategy = CachedChunkExtraction(inner, store, "doc-1") - with pytest.raises(DatabaseError): + with pytest.raises(DatabaseUnavailableError): await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) assert inner.calls == [], "extraction must not run once the graph is unreachable" + async def test_rejected_query_falls_open(self, ontology, ctx): + """A *rejected* query is not an outage. The server answered, so the + write phase will work and full extraction is the correct fallback — + otherwise a cache-only query the server dislikes would break an + update that ``cache_unchanged_chunks=False`` handles fine.""" + store = FakeGraphStore() + + async def _boom(document_id): + raise DatabaseError("syntax error near 'any'") + + store.get_document_chunk_texts = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) + assert len(inner.calls) == 1 + assert strategy.cached_chunk_count == 0 + async def test_lookup_budget_error_propagates(self, ontology, ctx): store = FakeGraphStore() @@ -1023,18 +1015,32 @@ async def test_rebuild_database_error_propagates(self, ontology, ctx): store = FakeGraphStore(chunk_texts=[("old-1", "a")]) async def _boom(chunk_ids): - raise DatabaseError("connection reset") + raise DatabaseUnavailableError("connection reset") store.get_entities_mentioned_in_chunks = _boom inner = RecordingExtractor() strategy = CachedChunkExtraction(inner, store, "doc-1") - with pytest.raises(DatabaseError): - await strategy.extract( - TextChunks(chunks=[_chunk("a", uid="new-1")]), ontology, ctx - ) + with pytest.raises(DatabaseUnavailableError): + await strategy.extract(TextChunks(chunks=[_chunk("a", uid="new-1")]), ontology, ctx) assert inner.calls == [] + async def test_rebuild_rejected_query_falls_open(self, ontology, ctx): + """Same distinction at the rebuild step.""" + store = FakeGraphStore(chunk_texts=[("old-1", "a")]) + + async def _boom(chunk_ids): + raise DatabaseError("unknown function 'any'") + + store.get_entities_mentioned_in_chunks = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract(TextChunks(chunks=[_chunk("a", uid="new-1")]), ontology, ctx) + assert len(inner.calls) == 1 + assert strategy.cached_chunk_count == 0 + assert strategy.extracted_chunk_count == 1 + async def test_other_errors_still_fail_open(self, ontology, ctx): """A cache that merely misbehaves must not block the update.""" store = FakeGraphStore() @@ -1049,3 +1055,170 @@ async def _boom(document_id): await strategy.extract(TextChunks(chunks=[_chunk("a")]), ontology, ctx) assert len(inner.calls) == 1, "fail-open path must still extract" assert strategy.cached_chunk_count == 0 + + +class TestRelationshipMergeUnion: + """A fact carried by both a cached and a freshly extracted chunk must + survive as ONE edge with both provenances. + + ExactMatchResolution keys relationships on (start, type, end), keeps the + first occurrence and rebuilds from it — it does not merge properties the + way the node path does. So two rows for one fact means the second row is + discarded outright: the fresh chunk's corrected description is lost, and + its chunk id never reaches source_chunk_ids, leaving the edge to be + swept as unsupported on a later update. Full re-extraction never hits + this because _aggregate_relations collapses duplicates upstream. + """ + + def _rel(self, fact, srcs): + return GraphRelationship( + start_node_id="alice", + end_node_id="acme", + type="RELATES", + properties={"fact": fact, "source_chunk_ids": list(srcs)}, + ) + + def test_same_edge_from_cached_and_extracted_is_unioned(self): + cached = GraphData(nodes=[], relationships=[self._rel("old fact", ["newA"])]) + fresh = GraphData(nodes=[], relationships=[self._rel("corrected fact", ["newB"])]) + + merged = CachedChunkExtraction._merge([cached, fresh]) + + assert len(merged.relationships) == 1 + props = merged.relationships[0].properties + assert props["fact"] == "corrected fact", "fresh extraction must win on properties" + assert props["source_chunk_ids"] == ["newA", "newB"], "both provenances must survive" + + def test_different_edge_types_stay_separate(self): + a = GraphData( + nodes=[], + relationships=[ + GraphRelationship( + start_node_id="alice", + end_node_id="acme", + type="RELATES", + properties={"source_chunk_ids": ["c1"]}, + ) + ], + ) + b = GraphData( + nodes=[], + relationships=[ + GraphRelationship( + start_node_id="alice", + end_node_id="acme", + type="MENTIONS", + properties={"source_chunk_ids": ["c2"]}, + ) + ], + ) + merged = CachedChunkExtraction._merge([a, b]) + assert len(merged.relationships) == 2 + + def test_direction_is_not_collapsed(self): + fwd = GraphData( + nodes=[], + relationships=[ + GraphRelationship( + start_node_id="a", + end_node_id="b", + type="RELATES", + properties={"source_chunk_ids": ["c1"]}, + ) + ], + ) + rev = GraphData( + nodes=[], + relationships=[ + GraphRelationship( + start_node_id="b", + end_node_id="a", + type="RELATES", + properties={"source_chunk_ids": ["c2"]}, + ) + ], + ) + merged = CachedChunkExtraction._merge([fwd, rev]) + assert len(merged.relationships) == 2 + + +class TestPerChunkFallbackScope: + """An unrebuildable chunk must not cost the whole document its cache.""" + + async def test_only_the_bad_chunk_is_extracted(self, ontology, ctx): + store = FakeGraphStore( + chunk_texts=[("old-1", "good"), ("old-2", "bad")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="ok", + label="Person", + name="Ok", + source_chunk_ids=["old-1"], + ), + ChunkEntityRow( + chunk_id="old-2", + entity_id="broken", + label="", # no concrete label + name="Broken", + source_chunk_ids=["old-2"], + ), + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + result = await strategy.extract( + TextChunks(chunks=[_chunk("good", 0, "new-1"), _chunk("bad", 1, "new-2")]), + ontology, + ctx, + ) + + assert strategy.cached_chunk_count == 1 + assert strategy.extracted_chunk_count == 1 + assert [c.uid for c in inner.calls[0].chunks] == ["new-2"] + assert any(n.id == "ok" for n in result.nodes), "the good chunk keeps its cache" + + async def test_dropping_a_chunk_cascades_to_edges_that_needed_it(self, ontology, ctx): + """Removing a chunk removes the entities only it mentioned, which can + strand an edge cited by another chunk. The fixpoint must catch that.""" + store = FakeGraphStore( + chunk_texts=[("old-1", "a"), ("old-2", "b")], + entity_rows=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="known", + label="Person", + name="Known", + source_chunk_ids=["old-1"], + ), + # only old-2 mentions 'shared', and old-2 is unrebuildable + ChunkEntityRow( + chunk_id="old-2", + entity_id="shared", + label="", + name="Shared", + source_chunk_ids=["old-2"], + ), + ], + rel_rows=[ + ChunkRelationshipRow( + chunk_id="old-1", + start_entity_id="known", + end_entity_id="shared", + rel_type="knows", + ) + ], + ) + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + await strategy.extract( + TextChunks(chunks=[_chunk("a", 0, "new-1"), _chunk("b", 1, "new-2")]), + ontology, + ctx, + ) + + assert strategy.cached_chunk_count == 0, "both chunks must fall back" + assert strategy.extracted_chunk_count == 2 + assert sorted(c.uid for c in inner.calls[0].chunks) == ["new-1", "new-2"]