diff --git a/CHANGELOG.md b/CHANGELOG.md index ff0a9a56..03696d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,75 @@ 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: a cache lookup or rebuild failure + 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` + (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`). + +### 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. 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 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/api/main.py b/graphrag_sdk/src/graphrag_sdk/api/main.py index c774e96b..bae09970 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, @@ -144,6 +147,46 @@ def _neutralize_context_close_tag(text: str) -> str: return _CONTEXT_CLOSE_RE.sub("", text) +def _same_embedding_model(stored: str, current: str) -> bool: + """Whether two model identifiers name the same embedding model. + + 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. + """ + 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: """Parse LLM-emitted JSON, tolerating optional markdown fences. @@ -1923,6 +1966,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 +2028,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 +2170,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 +2265,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 +2406,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 +2460,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 +2557,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(), ) @@ -2829,7 +2918,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 stored_model != 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 " @@ -3138,6 +3227,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 +3251,7 @@ def update_sync( chunker=chunker, extractor=extractor, resolver=resolver, + cache_unchanged_chunks=cache_unchanged_chunks, if_missing=if_missing, ctx=ctx, ) @@ -3193,6 +3284,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 +3307,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, diff --git a/graphrag_sdk/src/graphrag_sdk/core/connection.py b/graphrag_sdk/src/graphrag_sdk/core/connection.py index c47127f8..7d20de49 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, DatabaseUnavailableError + 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 DatabaseUnavailableError( + 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 DatabaseUnavailableError( "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 DatabaseUnavailableError( + 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. @@ -219,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 @@ -229,13 +259,25 @@ 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. An + ``ImportError`` from a missing driver still propagates: that is a + broken install, not a dead server. + """ try: + self._ensure_client() + from redis.asyncio import Redis 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/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 ───────────────────────────────────────────────── 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", 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..8589d7b3 --- /dev/null +++ b/graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py @@ -0,0 +1,449 @@ +# 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. +- 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 + 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()``): + +- 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.exceptions import DatabaseUnavailableError, LatencyBudgetExceededError +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, 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) + 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] + ) -> tuple[GraphData, set[str]]: + """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. + + 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]] = {} + for old_id, new_uid in pairs: + 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 ent_rows: + 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: + # _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)) + + 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 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: + # 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": []} + 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), unusable + + @staticmethod + def _merge(parts: list[GraphData]) -> GraphData: + """Merge cached and freshly extracted GraphData. + + 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] = {} + 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: + 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) + 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=list(rels_by_key.values()), + 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( + {c.uid for c in chunks.chunks} + ) + 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. 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) + 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: + data, unusable = await self._graph_data_from_cache( + [(old_id, chunk.uid) for chunk, old_id in cached], all_old_ids + ) + 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 + 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) diff --git a/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py b/graphrag_sdk/src/graphrag_sdk/storage/graph_store.py index 457ce2b8..c2cef098 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,135 @@ 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. + ``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. + """ + 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 + # 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, + 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 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 + # 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( + 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, 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..8eb20d1a --- /dev/null +++ b/graphrag_sdk/tests/test_cached_chunk_extraction.py @@ -0,0 +1,1224 @@ +"""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.exceptions import ( + DatabaseError, + DatabaseUnavailableError, + LatencyBudgetExceededError, +) +from graphrag_sdk.core.models import ( + ChunkEntityRow, + ChunkRelationshipRow, + EntityMention, + GraphData, + GraphNode, + GraphRelationship, + 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_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=[ + ChunkEntityRow( + chunk_id="old-1", + entity_id="ghost", + label=None, + name="Ghost", + source_chunk_ids=["old-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 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: + 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_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([]) == [] + 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 + + +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"] + + +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 DatabaseUnavailableError("connection refused") + + store.get_document_chunk_texts = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + 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() + + 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 DatabaseUnavailableError("connection reset") + + store.get_entities_mentioned_in_chunks = _boom + inner = RecordingExtractor() + strategy = CachedChunkExtraction(inner, store, "doc-1") + + 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() + + 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 + + +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"] 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:] == [] 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 diff --git a/graphrag_sdk/tests/test_facade.py b/graphrag_sdk/tests/test_facade.py index 36108731..fc0060eb 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,166 @@ async def test_no_config_node_passes(self, mock_conn, embedder): assert isinstance(result, RetrieverResult) +class TestSameEmbeddingModel: + """Model-name comparison used by the config guard. + + 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"), + ], + ) + def test_route_prefix_is_ignored(self, stored, current): + from graphrag_sdk.api.main import _same_embedding_model + + assert _same_embedding_model(stored, current) is True + + @pytest.mark.parametrize( + "stored,current", + [ + # Genuinely different models. + ("text-embedding-3-large", "text-embedding-3-small"), + ("azure/text-embedding-3-large", "azure/text-embedding-3-small"), + # 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"), + # 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_different_models_do_not_match(self, stored, current): + from graphrag_sdk.api.main import _same_embedding_model + + assert _same_embedding_model(stored, current) is False + + +class TestConfigProviderPrefix: + """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 + 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_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 + 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): + """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. + """ + 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."""