feat: chunk-level extraction cache for update() (cache_unchanged_chunks) - #288
Conversation
Typed rows for graph-backed cache reads: one entity per (chunk, entity) mention pair and one RELATES edge per (chunk, edge) provenance pair. Consumed by the upcoming chunk-level extraction cache. No behavior change.
Three schema-owning read methods backing the chunk-level extraction cache: - get_document_chunk_texts(document_id) — (chunk_id, text) snapshot of a document's live chunks, for hashing before the update cutover. - get_entities_mentioned_in_chunks(chunk_ids) — previously extracted entities per chunk. No e.type label fallback: a MERGE on a label the node doesn't carry would mint a duplicate node, so label-less rows return None and the consumer skips the node write. - get_relationships_for_chunks(chunk_ids) — RELATES edges whose source_chunk_ids provenance includes the chunks. Single edge scan per batch (any(c IN r.source_chunk_ids WHERE c IN $cids)) rather than a per-chunk re-scan; chunk intersection is client-side. All batched by _BATCH_SIZE with parameterized Cypher.
Decorator ExtractionStrategy for document updates: chunks whose text is byte-identical (SHA-256) to a chunk already stored for the same document skip LLM extraction — their entities, relationships, and mentions are rebuilt from the live graph and remapped onto the new chunk uids. Only genuinely new/changed chunks reach the inner extractor. The graph itself is the cache; nothing new is stored. Safe by construction: - deterministic entity ids MERGE onto existing nodes (SET n += preserves embeddings; only non-empty props are emitted) - entity source_chunk_ids remapped: this doc's old ids -> new uid(s) or dropped at cutover; other documents' ids pass through untouched - duplicate identical chunks each get their own mentions and provenance - label-less entities skip the node write (MERGE would mint a duplicate) but keep the mention, surviving orphan cleanup - fail-open: any cache lookup/rebuild error falls back to full extraction — worst case is skippable LLM spend, never data loss cached_chunk_count / extracted_chunk_count expose the split for reporting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
New opt-in flag (default False — existing behavior unchanged) on update(), update_sync(), apply_changes(), and apply_changes_sync(). When set, update() wraps the effective extractor in CachedChunkExtraction during Phase 3 — old chunks still exist until the Phase 5 cutover, so the cache read is safe by construction and the crash-safe pending/commit/rollforward state machine is untouched. Cache effectiveness is surfaced in UpdateResult.metadata["cache_stats"] (cached_chunks / extracted_chunks). The if_missing="ingest" fresh-ingest fallthrough is not wrapped (a new document has no chunks to reuse). Caveats documented on the docstring: the graph is the cache, so manual entity deletions are not resurrected from unchanged chunks, and ontology/prompt/model changes do not re-extract unchanged chunks — pass False to force a full rebuild.
22 unit tests (fake graph store + recording inner extractor) covering the full matrix: all-cached / all-new / mixed splits, provenance remap (own old ids remapped, foreign ids pass through, dropped ids die), duplicate identical chunks, label-less entity guard, relationship rebuild + same-pair dedup, fail-open on lookup and rebuild failures, cached+extracted merge (fresh props win, provenance unions), stats reset, and the three GraphStore accessors (row mapping, bad-row filtering, single-scan query shape). 4 integration tests (RUN_INTEGRATION=1, real FalkorDB) with a strict scripted LLM — an LLM call for an unchanged chunk raises, so passing tests are the proof that caching skips extraction: - unchanged chunk skips LLM; cache_stats reported; provenance and mentions point only at live chunks; replaced-chunk entity cleaned - manually deleted entity is not resurrected by a cached update - default off: both chunks re-extracted, no cache_stats - whole-document no_op short-circuit unaffected
Export CachedChunkExtraction at top level, from graphrag_sdk.ingestion, and from graphrag_sdk.ingestion.extraction_strategies, mirroring GraphExtraction. Document the feature, its cache_stats reporting, and its caveats under Unreleased in the CHANGELOG.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughGraphRAG adds an opt-in chunk-level extraction cache. The cache reuses graph data for unchanged chunks, remaps provenance to current chunk IDs, reports cache statistics, and supports asynchronous, synchronous, and batch update APIs. ChangesChunk-Level Extraction Cache
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant GraphRAG
participant CachedChunkExtraction
participant GraphStore
participant ExtractionStrategy
Client->>GraphRAG: update(cache_unchanged_chunks=True)
GraphRAG->>CachedChunkExtraction: extract(document chunks)
CachedChunkExtraction->>GraphStore: load prior chunk text and graph provenance
CachedChunkExtraction->>ExtractionStrategy: extract changed chunks
CachedChunkExtraction-->>GraphRAG: return merged GraphData and cache statistics
GraphRAG-->>Client: return UpdateResult
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Adds an opt-in, chunk-level extraction cache for GraphRAG.update() that reuses previously extracted entities/relationships for byte-identical chunks, avoiding unnecessary LLM calls while preserving existing update safety semantics (pending/commit/cutover).
Changes:
- Introduces
CachedChunkExtractiondecorator strategy and wires it intoupdate(),update_sync(),apply_changes(), andapply_changes_sync()behindcache_unchanged_chunks=True, emittingUpdateResult.metadata["cache_stats"]. - Extends
GraphStorewith read accessors to snapshot chunk texts and read back entities/relationships needed to rebuildGraphDatafrom the live graph. - Adds comprehensive unit + integration tests covering split/remap/merge/fail-open behavior and the new storage accessors; documents behavior in
CHANGELOG.md.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| graphrag_sdk/tests/test_cached_chunk_extraction.py | New unit + integration coverage for cached chunk extraction and GraphStore accessors. |
| graphrag_sdk/src/graphrag_sdk/storage/graph_store.py | Adds three cache read accessors used to rebuild extraction results from the graph. |
| graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py | Implements the decorator extraction strategy that splits cached vs extracted chunks and merges results. |
| graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/init.py | Exports CachedChunkExtraction from extraction strategies package. |
| graphrag_sdk/src/graphrag_sdk/ingestion/init.py | Exports CachedChunkExtraction from ingestion package. |
| graphrag_sdk/src/graphrag_sdk/core/models.py | Adds typed row models ChunkEntityRow / ChunkRelationshipRow for cache rebuild reads. |
| graphrag_sdk/src/graphrag_sdk/api/main.py | Adds cache_unchanged_chunks flag to update/apply_changes APIs and attaches cache stats to results. |
| graphrag_sdk/src/graphrag_sdk/init.py | Re-exports CachedChunkExtraction at top level. |
| CHANGELOG.md | Documents the new chunk-level cache feature, caveats, and API surface changes. |
Comments suppressed due to low confidence (1)
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:799
srcs(ther.source_chunk_idsprovenance) is iterated without validating its type. If the property is missing or corrupted into a scalar (e.g., a string), this will iterate characters and emit incorrectChunkRelationshipRows. For consistency with the other cache snapshot accessors (which skip bad rows), guard withisinstance(srcs, list)and only consider string chunk ids.
for cid in srcs or []:
if cid not in batch_set:
continue
rows.append(
ChunkRelationshipRow(
chunk_id=cid,
start_entity_id=start_id,
end_entity_id=end_id,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
graphrag_sdk/tests/test_cached_chunk_extraction.py (2)
613-614: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the redundant
@pytest.mark.asyncio.Every other async test in this file omits it, and auto mode already collects coroutines.
♻️ Proposed change
-@pytest.mark.asyncio `@pytest.mark.integration` class TestCachedUpdateIntegration:Based on learnings: pytest-asyncio in
graphrag_sdk/pyproject.tomlusesasyncio_mode = "auto", sopytest.mark.asynciois not needed on async tests in this repo.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_cached_chunk_extraction.py` around lines 613 - 614, Remove the redundant `@pytest.mark.asyncio` decorator from the affected async test, leaving the `@pytest.mark.integration` decorator and test behavior unchanged; the repository’s asyncio auto mode will handle coroutine collection.Source: Learnings
751-776: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the scripted responses were consumed.
MockLLM(strict=True)only fails on over-consumption;._call_index >= len(self._responses)is false when fewer LLM calls occur. Add an assertion such asassert llm._call_index == 4so an under-extracting regression doesn’t silently pass.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_cached_chunk_extraction.py` around lines 751 - 776, Add an assertion in test_default_off_extracts_all_chunks verifying llm._call_index equals 4 after the update, ensuring all scripted LLM responses were consumed and detecting under-extraction.graphrag_sdk/src/graphrag_sdk/storage/graph_store.py (1)
763-811: 🚀 Performance & Scalability | 🔵 TrivialFull RELATES scan per batch — acceptable now, but worth flagging.
The docstring already documents that
source_chunk_idslist membership is a plain list property, so no index can serve the membership test — but the edge scan runs ONCE per batch, not once per chunk id. For documents whose chunk count exceeds_BATCH_SIZE(500), this full(a:__Entity__)-[r:RELATES]->(b:__Entity__)scan repeats once per batch — cost scales with total RELATES edge count × number of batches, which could offset the cache's savings on very large graphs with long documents.No action needed now since the tradeoff is intentional and documented, but worth monitoring if
RELATESedge counts grow large in production deployments using this cache option.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/storage/graph_store.py` around lines 763 - 811, No code changes are needed: retain the intentional batched RELATES scan in get_relationships_for_chunks and its existing documentation, and monitor performance as edge counts or documents spanning multiple _BATCH_SIZE batches grow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@graphrag_sdk/src/graphrag_sdk/storage/graph_store.py`:
- Around line 763-811: No code changes are needed: retain the intentional
batched RELATES scan in get_relationships_for_chunks and its existing
documentation, and monitor performance as edge counts or documents spanning
multiple _BATCH_SIZE batches grow.
In `@graphrag_sdk/tests/test_cached_chunk_extraction.py`:
- Around line 613-614: Remove the redundant `@pytest.mark.asyncio` decorator from
the affected async test, leaving the `@pytest.mark.integration` decorator and test
behavior unchanged; the repository’s asyncio auto mode will handle coroutine
collection.
- Around line 751-776: Add an assertion in test_default_off_extracts_all_chunks
verifying llm._call_index equals 4 after the update, ensuring all scripted LLM
responses were consumed and detecting under-extraction.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 85ef1aad-8337-4aed-8fea-e09d593364c3
📒 Files selected for processing (9)
CHANGELOG.mdgraphrag_sdk/src/graphrag_sdk/__init__.pygraphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/src/graphrag_sdk/core/models.pygraphrag_sdk/src/graphrag_sdk/ingestion/__init__.pygraphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/__init__.pygraphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.pygraphrag_sdk/src/graphrag_sdk/storage/graph_store.pygraphrag_sdk/tests/test_cached_chunk_extraction.py
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:756
get_entities_mentioned_in_chunks()iteratessource_chunk_idswith(source_chunk_ids or []). If the stored value is accidentally a string (tampered graph), this iterates characters and silently emits invalid provenance. Also,cid/eidshould be enforced as non-empty strings to avoid propagating corrupted ids into the cache rebuild path.
for row in result.result_set or []:
cid, eid, labels, name, etype, description, source_chunk_ids = row
if not cid or not eid:
continue
label = next((lb for lb in (labels or []) if lb != "__Entity__"), None)
rows.append(
ChunkEntityRow(
chunk_id=cid,
entity_id=eid,
label=label,
name=name if isinstance(name, str) else None,
type=etype if isinstance(etype, str) else None,
description=description if isinstance(description, str) else None,
source_chunk_ids=[
s for s in (source_chunk_ids or []) if isinstance(s, str)
],
)
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:711
get_document_chunk_texts()claims to skip rows with a missing id, but it currently accepts any truthycidvalue (including non-strings). Since callers treat these ids as chunk-id strings, letting non-strings through can cause cache misses or incorrect provenance remapping; it’s safer to enforcecidis a non-emptystrin addition totextbeing astr.
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
graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py:113
_old_chunks_by_hash()collapses multiple old chunks with identical text down to a single old chunk id (setdefault). If a document previously contained duplicate chunk texts, and those duplicates ended up with different extracted mentions/edges (LLM nondeterminism, manual edits, or partial writes), the cache rebuild will only consult one of the old chunk ids and can drop entities/relationships that were only attached to the other identical-text chunk(s), potentially causing unintended orphan cleanup. Consider tracking all old chunk ids per text hash and rebuilding from the union (with mention de-duplication for repeated (entity_id, new_uid) pairs).
async def _old_chunks_by_hash(self) -> tuple[dict[str, str], set[str]]:
"""Return (sha256(chunk text) -> old chunk id, all old chunk ids)."""
by_hash: dict[str, str] = {}
all_ids: set[str] = set()
for cid, text in await self._graph_store.get_document_chunk_texts(self._document_id):
all_ids.add(cid)
# First occurrence wins; duplicates map to the same content anyway.
by_hash.setdefault(_sha256(text), cid)
return by_hash, all_ids
graphrag_sdk/tests/test_cached_chunk_extraction.py:36
- The
_sha()helper is defined but never used in this test module, which adds dead code and can confuse readers about intended assertions.
def _sha(text: str) -> str:
return hashlib.sha256(text.encode("utf-8")).hexdigest()
Tampered/partial graphs may hold a scalar (e.g. a string) where a list is expected in labels(e) or source_chunk_ids. Iterating it would yield characters as bogus labels/provenance instead of treating the row as a cache miss. Treat non-list values as empty (entities) or skip the row (relationships), per the documented skip-bad-rows contract. Addresses Copilot review feedback on PR #288.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:710
get_document_chunk_texts()is documented as filtering out tampered/partial rows, but it currently accepts non-stringc.idvalues (it only checks truthiness). That can leak invalid IDs into the cache key map and later provenance remapping, potentially causing mismatches (e.g.,intids won’t compare equal to expectedstrids). Filter toisinstance(cid, str)as well (and keep the existing text-type guard).
for row in result.result_set or []:
cid, text = row[0], row[1]
if cid and isinstance(text, str):
out.append((cid, text))
`EmbedderConfig.to_embedder()` prepends the provider to the model name when `provider` is set (e.g. `azure/text-embedding-3-large`), but the name persisted on the graph's config node is unprefixed. `_validate_graph_config` compared the two raw strings, so simply routing an existing graph through a provider — same model, same dimensions — raised ConfigError and made every subsequent update fail. Normalize both sides before comparing: strip a leading `<provider>/` segment when the prefix is a known litellm provider. The provider list is read from litellm at first use and cached, with a static fallback for when that import is unavailable. The check itself is deliberately kept: a genuine model swap silently corrupts vector search, since old and new embeddings stay mutually incomparable without ever raising.
Twelve tests over the fix in the previous commit. `TestBareModelName` pins the helper's boundaries: known prefixes are stripped, unknown ones are not, and names that merely contain a slash (bare model names with a `/` in them, e.g. HuggingFace repo ids) are left intact. Case is deliberately significant. `TestConfigProviderPrefix` drives the validator end to end, asserting the prod scenario now passes (`azure/X` vs stored `X`) while a real model change still raises. Both mutants were checked: reverting to the raw comparison fails four tests, and stripping any leading path segment fails two.
Eight tests against a real FalkorDB, exercising cases the existing suite does not reach. The most valuable is `test_cached_and_uncached_updates_converge`: it runs the same edit twice, once with the cache and once without, and asserts the resulting graphs are identical. That makes the cache's correctness property explicit rather than checking hand-written expectations, and it catches faults in the rebuild Cypher itself — a mutation swapping two returned columns is invisible to fixture-backed tests but fails here. The rest cover reordered chunks, duplicate identical chunks, provenance isolation across documents, orphan cleanup when a chunk disappears, stability across repeated cached updates, and the no-op short circuit reporting no cache stats. Uses a local `ScriptedExtractor` that emits relationships; the shared scripted fixture always returns an empty relationship list, so the rebuild path was previously never executed.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
graphrag_sdk/tests/test_cached_chunk_extraction_edges.py (1)
229-231: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the redundant
pytest.mark.asynciomarker.
graphrag_sdk/pyproject.tomlsetsasyncio_mode = "auto", so pytest-asyncio already collectsasync deftests in this repository. Keeppytest.mark.integration, which carries real meaning for the env-gated run.♻️ Proposed change
-@pytest.mark.asyncio `@pytest.mark.integration` class TestChunkCacheEdges:Based on learnings: because pytest-asyncio in
graphrag_sdk/pyproject.tomlusesasyncio_mode = "auto", async test functions should be auto-detected, andpytest.mark.asynciois generally not needed in this repository.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_cached_chunk_extraction_edges.py` around lines 229 - 231, Remove the redundant pytest.mark.asyncio decorator from TestChunkCacheEdges while preserving pytest.mark.integration; async tests should continue relying on the repository’s asyncio_mode = "auto" configuration for collection.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@graphrag_sdk/tests/test_cached_chunk_extraction_edges.py`:
- Line 81: Update the backfill callsite of compute_entity_id in the entity-cache
path to pass ent_name as the first argument and label as the second, matching
its (name, entity_type) signature. Preserve the resulting IDs so cached entity
rows MERGE onto the corresponding GraphExtraction IDs.
---
Nitpick comments:
In `@graphrag_sdk/tests/test_cached_chunk_extraction_edges.py`:
- Around line 229-231: Remove the redundant pytest.mark.asyncio decorator from
TestChunkCacheEdges while preserving pytest.mark.integration; async tests should
continue relying on the repository’s asyncio_mode = "auto" configuration for
collection.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: df4804c7-0261-4f87-be01-abc371cc2bff
📒 Files selected for processing (3)
graphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/tests/test_cached_chunk_extraction_edges.pygraphrag_sdk/tests/test_facade.py
| if line.startswith("E|"): | ||
| _, name, etype, desc = line.split("|", 3) | ||
| types[name.strip().lower()] = etype | ||
| eid = compute_entity_id(name, etype) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Resolve the definition and the production call sites of compute_entity_id.
set -euo pipefail
fd -t f 'entity_extractors.py' --exec ast-grep outline {} --items all
fd -t f 'entity_extractors.py' --exec ast-grep run \
--pattern 'def compute_entity_id($$$):
$$$' --lang python {}
echo '--- call sites ---'
rg -nP -C2 '\bcompute_entity_id\s*\(' --type=pyRepository: FalkorDB/GraphRAG-SDK
Length of output: 2261
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- compute_entity_id helper and production calls ---'
rg -n -C 4 '\bcompute_entity_id\s*\(' --type=py graphrag_sdk/src graphrag_sdk/tests
echo '--- relevant production code slices ---'
sed -n '108,132p' graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py
sed -n '1100,1128p' graphrag_sdk/src/graphrag_sdk/api/main.py
echo '--- test file slices ---'
sed -n '68,118p' graphrag_sdk/tests/test_cached_chunk_extraction_edges.pyRepository: FalkorDB/GraphRAG-SDK
Length of output: 15192
🏁
Length of output: 25
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- main.py backfill surrounding context ---'
sed -n '1070,1132p' graphrag_sdk/src/graphrag_sdk/api/main.py | nl -ba
echo '--- imports in main.py top ---'
sed -n '1,80p' graphrag_sdk/src/graphrag_sdk/api/main.py | nl -ba
echo '--- deterministic argument-order probe for compute_entity_id ---'
python3 - <<'PY'
def compute_entity_id(name: str, entity_type: str = "") -> str:
base = name.strip().lower().replace(" ", "_")
if entity_type:
return f"{base}__{entity_type.strip().lower()}"
return base
ent_name = "Alice Corp"
label = "Person"
print("name, label:", compute_entity_id(ent_name, label))
print("label, name:", compute_entity_id(label, ent_name))
PYRepository: FalkorDB/GraphRAG-SDK
Length of output: 244
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- main.py backfill surrounding context ---'
sed -n '1070,1132p' graphrag_sdk/src/graphrag_sdk/api/main.py
echo '--- imports in main.py top ---'
sed -n '1,80p' graphrag_sdk/src/graphrag_sdk/api/main.py
echo '--- deterministic argument-order probe for compute_entity_id ---'
python3 - <<'PY'
def compute_entity_id(name: str, entity_type: str = "") -> str:
base = name.strip().lower().replace(" ", "_")
if entity_type:
return f"{base}__{entity_type.strip().lower()}"
return base
ent_name = "Alice Corp"
label = "Person"
print("name, label:", compute_entity_id(ent_name, label))
print("label, name:", compute_entity_id(label, ent_name))
PYRepository: FalkorDB/GraphRAG-SDK
Length of output: 5653
Fix the compute_entity_id callsite in the backfill path.
compute_entity_id has the signature compute_entity_id(name, entity_type), so graphrag_sdk/src/graphrag_sdk/api/main.py:1117 should pass ent_name first and label second. The current call generates type-qualified ids where the entity name becomes the label prefix, so cached entity rows will not MERGE onto GraphExtraction ids.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@graphrag_sdk/tests/test_cached_chunk_extraction_edges.py` at line 81, Update
the backfill callsite of compute_entity_id in the entity-cache path to pass
ent_name as the first argument and label as the second, matching its (name,
entity_type) signature. Preserve the resulting IDs so cached entity rows MERGE
onto the corresponding GraphExtraction IDs.
712d7b6 to
160fd8f
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (2)
graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py:304
- On cache rebuild failure, cached chunks are appended to
to_extract, which can reorder chunks relative to the originalchunks.chunksinput (e.g., cached chunk(s) that originally appeared before changed ones get moved to the end). If the inner extractor is order-sensitive (chunk indices, cross-chunk context), the fail-open path can diverge from a normal full extraction. Prefer falling back to extracting all original chunks in their original order and update the stats to match.
to_extract.extend(chunk for chunk, _ in cached)
self.extracted_chunk_count += self.cached_chunk_count
self.cached_chunk_count = 0
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:710
get_document_chunk_texts()aims to be robust to tampered/partial graphs, but it currently accepts any truthycidvalue (including non-strings). A non-string chunk id can later break cache remapping (e.g.,all_old_idsmembership checks) and lead to stale provenance being treated as “foreign” and preserved. Filter out non-string chunk ids the same way non-string texts are filtered.
cid, text = row[0], row[1]
if cid and isinstance(text, str):
out.append((cid, text))
The config guard compared stored and live embedder names as raw strings, so
a graph reached through a different endpoint ("azure/text-embedding-3-large"
vs "text-embedding-3-large") was reported as a model mismatch even though the
stored vectors are unchanged.
Treat one leading route segment as optional on either side, which covers a
prefix appearing, disappearing, or changing. Matching is anchored on "/" so
only a whole segment is ever ignored -- a substring check would accept
"text-embedding-3-large" against "text-embedding-3-large-v2", silently passing
the mismatch this guard exists to catch.
This drops the provider-prefix lookup and its static fallback list, so the
comparison no longer depends on the optional litellm extra or on a list that
goes stale as providers are added.
Replace the provider-prefix normalisation tests with cases for the new comparison: a route appearing, disappearing, or changing between two providers, multi-segment routes, and case/whitespace differences. Pin the negative cases that keep the guard useful -- a trailing word such as "text-embedding-3-large-v2" must not match "text-embedding-3-large", since a looser substring check would accept it and pass a real mismatch. Two config-guard tests changed contract: two differently namespaced names now resolve to the same model, and the test asserting the opposite is replaced. The test stubbing out a missing litellm import is dropped -- the comparison no longer imports it.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (2)
graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/cached_chunk_extraction.py:320
- In the fail-open path, cached chunks are appended to
to_extractafter the already-uncached chunks. This changes the original chunk order, so a cache-rebuild failure can produce a different extraction batch ordering than the no-cache behavior (and some extractors may be order/context sensitive). To preserve semantics, fall back to extracting all chunks in their original order when cache rebuild fails.
# 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
graphrag_sdk/src/graphrag_sdk/storage/graph_store.py:792
get_relationships_for_chunks()usesany(c IN r.source_chunk_ids ...)without guarding againstr.source_chunk_idsbeing null. Elsewhere (e.g.,delete_stale_relationships) the code explicitly checksr.source_chunk_ids IS NOT NULLbefore applyingany(...). Adding the same guard here avoids null-handling differences/edge-case errors on partially populated graphs and keeps query semantics consistent.
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},
galshubeli
left a comment
galshubeli
left a comment
There was a problem hiding this comment.
Re-reviewed after the embedding-match rework. Splitting this into the two independent changes.
Cache (the bulk of the PR) — solid
I specifically tried to break several things and couldn't:
rel_by_pairkeyed on(start, end)ignoringrel_typeis correct, not lossy —upsert_relationshipsdoesMERGE (a)-[r:RELATES]->(b), so storage already permits one edge per pair. The rebuild matches the schema.- Provenance remap (own old ids to new uids, other documents' ids passed through) is right, and the
RELATESasymmetry is safe becauseupsert_relationshipsUNIONs. - Duplicate-text chunks via
id_map: old_id -> list[new_uid]are handled properly. _mergeordering[cached, extracted]with union-on-source_chunk_idsand last-wins elsewhere gives fresh extraction precedence. Correct.- The scalar-vs-list guards in the new accessors are genuinely good defensive work.
Comments below are the real ones: the edge scan in get_relationships_for_chunks, fail-open on database errors, and the public export of CachedChunkExtraction. None are blocking-by-design, all are worth fixing before merge.
Embedding model matching — going the wrong way
Dropping the hardcoded provider list was right. But the replacement is strictly more permissive than what it replaced: the old rule only stripped known providers, the new one strips any first segment. Everything the old code rejected still passes, plus more.
I ran the new function against realistic names:
| stored | current | result |
|---|---|---|
sentence-transformers/all-MiniLM-L6-v2 |
myorg/all-MiniLM-L6-v2 |
match |
BAAI/bge-m3 |
ollama/bge-m3 |
match |
intfloat/e5-large-v2 |
rando/e5-large-v2 |
match |
All same-dimension, all different vectors, all silently accepted. These are reachable through the custom-embedder path the SDK documents in examples/04_custom_provider.py, which points at SentenceTransformer — whose model names are exactly org/model repo ids.
test_unknown_prefix_still_raises was inverted to assert these now pass, on the grounds that "the dimension check below still guards the case that actually corrupts retrieval." That claim isn't right — the dimension check compares stored_dim != embedding_dimension, and every case above has identical dimensions. Retrieval then fails soft: results come back, ranked by cosine against incomparable vectors, with no error.
The docstring's own line — "a routing prefix is indistinguishable from a vendor namespace" — is exactly the problem, and it's not fixable in string space, because the information isn't in the string.
The useful part: you already compute the real answer and discard it. _validate_graph_config embeds a fixed probe string and uses only len(probe). Store that vector on the config node and compare cosine — free (the call already happens), passes the Azure/OpenAI route change this PR is about, and catches all three rows above plus dimensions= changes and silent vendor model refreshes. Details on _same_embedding_model.
Ranked: fix EmbedderConfig upstream (both sides are already symmetric inside the SDK — the asymmetry comes from the server) > probe fingerprint > explicit embedding_model_policy override > current string widening.
Either way this should be a separate PR. It's unrelated to the cache, it relaxes a global invariant, and in a 2,100-line change where the rest is strong it'll get carried along on that strength.
| 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): |
There was a problem hiding this comment.
Worth stepping back on this comparison generally.
Inside the SDK both sides are already symmetric — _write_graph_config writes self.embedder.model_name and this reads it back to compare against self.embedder.model_name. A graph written by this SDK and read by this SDK cannot hit the prefix mismatch. The asymmetry is entirely the server's EmbedderConfig.to_embedder() storing a bare name and later constructing a prefixed one.
So a global invariant is being loosened for every SDK user to compensate for a caller bug in a different repo. And because _write_graph_config uses MERGE ... SET, affected graphs rewrite the stored name on their next successful ingest — only the bootstrap needs a one-shot fix.
The original incident (a ConfigError that then broke every subsequent update) is a legitimate complaint, but it's a complaint about an un-overridable, inscrutable error, not about the error existing. That's fixable without widening the check:
if raw names match -> pass
if probe vectors are similar -> pass, rewrite stored name
if basenames match and dims match -> ConfigError naming the override flag
otherwise -> ConfigError
The third branch is where this PR's users land, and the message carries its own fix: "same base model via a different provider — pass embedding_model_policy='allow_provider_change' if these vectors are comparable." One deliberate line during a rare, supervised migration, and it self-heals on the next write. That's a much better shape than permanently relaxing the check for everyone.
Preference order: fix EmbedderConfig upstream > probe-vector fingerprint > explicit override flag > the current string widening.
| 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) " |
There was a problem hiding this comment.
This is a full graph edge scan sitting in the hot path of the feature.
source_chunk_ids is a plain list property, so nothing can index this predicate — every RELATES edge in the database is visited, once per 500-chunk batch. The docstring frames "the edge scan runs ONCE per batch, not once per chunk id" as the optimization, but the baseline it optimizes from is O(all edges in the graph).
So cost scales with total graph size, not with the document being updated. Editing one paragraph of a small document in a graph holding a few hundred other documents pays for all of them — which is the exact cost profile cache_unchanged_chunks exists to remove. The gain gets thinner as the graph grows, which is backwards.
Chunk ids are indexed, so reaching the edges through the chunks avoids the scan:
MATCH (c:Chunk) WHERE c.id IN $cids
MATCH (e:__Entity__)-[:MENTIONED_IN]->(c)
MATCH (e)-[r:RELATES]-(o:__Entity__)
RETURN ...The provenance intersection is already computed client-side from the returned source_chunk_ids, so the filtering below doesn't need to change.
| old_by_hash, all_old_ids = await self._old_chunks_by_hash( | ||
| {c.uid for c in chunks.chunks} | ||
| ) | ||
| except Exception as exc: |
There was a problem hiding this comment.
Fail-open is right for a cache miss but wrong for an infrastructure error.
The realistic trigger for this except is a DatabaseError — dropped connection, query timeout — not a logic fault. When the graph is unreachable, this decides to run full LLM extraction over every chunk in the document. That extraction succeeds (it only needs the LLM), and then the write phase hits the same dead connection and fails anyway.
So an error that could surface in milliseconds instead costs a full document of LLM spend and lands in the same place. The PR frames the worst case as "LLM spend, never data loss" — this is that worst case being reached by the most likely failure mode rather than an unlikely one.
Same at the rebuild fallback below (line 311).
Letting DatabaseError and LatencyBudgetExceededError propagate and keeping fail-open for everything else would fix it. _validate_graph_config already uses exactly this pattern — except ConfigError: raise / except LatencyBudgetExceededError: raise / except Exception: log.
| 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 — |
There was a problem hiding this comment.
This docstring describes a silent-data-loss path that the class can't structurally prevent, and the class is public API (exported from both graphrag_sdk.__init__ and ingestion.__init__).
The failure is spelled out here in full: a chunk hashes to itself, the rebuild finds no mentions, emits empty GraphData, skips the extractor, drops the entities — with no exception for the fail-open path to catch — and cached_chunk_count still reports a hit. Every safety net in the design misses this one.
The exclude set only helps when the new chunks carry different uids from the old. That's a property of how update() happens to wire this up, not an invariant of the class. Someone constructing CachedChunkExtraction(inner=..., graph_store=..., document_id=<live doc>) and running it over that document's own chunks gets entities deleted and a success result.
The most recent commit before the embedding work is fix: never treat a chunk as its own cache entry, which is decent evidence this is reachable in practice.
Either keep it internal for now (drop it from the public exports) and revisit once the contract has settled, or make the requirement enforceable — e.g. take the read document id and the write document id as separate arguments and raise when they're equal, so the caller has to state intent rather than get it right by accident.
| # label and mint a duplicate node instead of matching this | ||
| # one. Skip the node write — the mention above still keeps | ||
| # the entity alive through orphan cleanup. | ||
| logger.warning("Skipping cache node write for label-less entity %s", eid) |
There was a problem hiding this comment.
This drop isn't visible in cache_stats, which makes the feature's only operational signal unreliable.
A label-less entity is skipped with a log line. Nothing raises, so fail-open never fires, and the chunk is still counted in cached_chunk_count as fully reused. Same for rows discarded by the not cid or not eid guard in get_entities_mentioned_in_chunks.
So the graph can come out thinner than a full re-extraction would produce while UpdateResult.metadata["cache_stats"] reports a clean 100% hit rate. For a fail-open cache that number is how an operator knows things are healthy — it shouldn't be the one thing blind to partial rebuilds. A skipped_entities counter alongside cached_chunks / extracted_chunks would make it observable without changing behavior.
| ) | ||
| for row in result.result_set or []: | ||
| cid, eid, labels, name, etype, description, source_chunk_ids = row | ||
| if not cid or not eid: |
There was a problem hiding this comment.
not cid is falsy, not absent — a chunk id of "0" is discarded here as though the row were malformed.
Almost certainly unreachable with current uid generation, but this is a data-integrity guard, and conflating falsy with missing inside one stops being theoretical the day id generation changes. if cid is None or eid is None costs nothing. Same at line 795 for start_id / end_id.
The name comparison stripped a leading segment unconditionally, so any two
identifiers sharing a tail matched. That collapsed distinct models onto one
name: "BAAI/bge-m3" and "ollama/bge-m3" are a fp32 release and a quantized
build, and "sentence-transformers/all-MiniLM-L6-v2" and a finetune under
another org are different weights entirely.
The dimension check cannot stand in for this -- every realistic collision here
keeps the base model's dimensions -- and the resulting failure is soft, since
retrieval still returns results, ranked against another model's vectors.
Only strip a segment when the other side has none. A bare name carries no
route, so a segment present on one side alone is one; two qualified names are
two owners and are compared in full.
Cost: a route that changes while both sides stay qualified ("azure/x" ->
"openai/x") is now read as an owner change and rejected. That is rare next to
the collisions it catches, and it fails loudly rather than silently.
Cover the pairs the dimension check cannot see -- a finetune under another org, a quantized build, and the same model id served by two vendors -- so a future loosening of the name rule has to break a test to land. Restore test_unknown_prefix_still_raises, which had been inverted to assert that two owner-qualified names pass. They do not: identical dimensions with different weights is the one case neither check would catch. Move the both-sides-qualified route change to the negative cases and state why it is the accepted cost of the rule.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (2)
graphrag_sdk/tests/test_facade.py:1195
- This test asserts that an arbitrary unknown prefix (
my-org/) should be treated as an ignorable route when compared to a bare model name. That makes the config guard too permissive for owner-qualified model names (org finetunes, namespaces). If prefix-stripping is intended only for known provider routes, this case should not be in the "route prefix is ignored" set.
("my-org/custom-embedder", "custom-embedder"),
graphrag_sdk/src/graphrag_sdk/api/main.py:187
_same_embedding_modelcurrently treats any leading segment as a routable prefix when the other side is bare (e.g. it would acceptBAAI/bge-m3vsbge-m3). That weakens the config guard and can silently allow incompatible embeddings when the prefix is actually an owner/namespace rather than a provider route. Consider only stripping prefixes that are known embedding providers (e.g.azure,openai,vertex_ai) and treat unknown prefixes as identity (i.e. mismatch when compared to a bare name).
_, 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:
Driver failures escaped as raw redis exceptions, so the chunk cache's broad fallback swallowed them and re-ran extraction on every chunk -- burning LLM calls before the write failed anyway. FalkorDB probes the server while the client is constructed, so the error surfaced from _ensure_client() before the retry loop ever ran. All four exit paths now raise DatabaseError, and the cache re-raises DatabaseError and LatencyBudgetExceededError instead of falling back. The wrapped message keeps the driver's original text -- vector_store and ontology_store both match on it to detect existing indexes.
Covers every driver exit path, that the message keeps the driver's text, that a failed connect reuses its pool instead of stranding one, and that the cache re-raises DatabaseError rather than falling back.
The cache no longer falls back on *any* failure, and wrapping driver errors in DatabaseError changes what callers must catch.
| ("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. |
| 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 |
Two cache-rebuild paths could silently delete live data. A label-less entity was skipped with a warning, on the assumption the entity survived. It does, but IngestionPipeline._filter_quality drops every relationship incident to a node absent from the batch, so the rebuilt provenance never landed and the post-cutover sweep then deleted edges the unchanged chunk still supported. Re-ingesting an unmodified document lost facts. A stored relationship whose endpoint no cached chunk mentions fails the same way. The pipeline only ever persists an edge alongside its endpoints, so this needs a graph written by something else, but the outcome is identical. Neither case can reproduce the graph, so refuse to rebuild and let the existing fallback hand the document to real extraction, which is always correct. Both now also surface in cache_stats as extracted chunks rather than disappearing into a log line.
|
Blocker (1)
Should fix before merge (2)
Nits (3)
Only #1 blocks merge. The embedding-match rework, |
_merge concatenated relationship lists, so a chunk that was rebuilt from the cache and a chunk that was re-extracted could both emit the same (start, type, end) edge. ExactMatchResolution keeps the first occurrence and rebuilds the edge from it, dropping the rest, and cached parts are appended before extracted ones - so a stale fact won over its correction and the new chunk id never reached source_chunk_ids. A later update then emptied the list and the sweep deleted the edge. Full re-extraction does not have this problem because _aggregate_relations collapses duplicates upstream with provenance unioned. Relationships are now keyed on (start_node_id, type, end_node_id) the same way nodes already were: later parts win on properties, and source_chunk_ids is unioned across all occurrences. Also: - Split DatabaseUnavailableError out of DatabaseError. Wrapping every driver failure as DatabaseError made the cache re-raise on a query the server had rejected, turning a case that falls open cleanly into a hard failure. Only a server that never answered - unreachable, open circuit breaker, exhausted retries - propagates now. The new type subclasses DatabaseError, so existing handlers are unchanged. query()'s permanent-error check also recognizes syntax errors, invalid input, unknown functions, type mismatches and missing procedures, so a rejected query fails fast instead of spending the whole retry budget first. - Scope cache fallback per chunk. An entity whose label is missing, or an edge whose endpoint is not mentioned in the citing chunk, voided the entire document's cache. Only the affected chunks are re-extracted now, computed as a fixpoint - dropping a chunk removes entities only it mentioned, which can strand edges other chunks cite. - ping() re-raises ImportError instead of reporting a missing falkordb package as a server that is down. - Correct a stale docstring on get_entities_mentioned_in_chunks.
Bump version to 1.4.0 and cut the [Unreleased] section into [1.4.0]: the chunk-level extraction cache for update() (FalkorDB#288) and the DatabaseError / DatabaseUnavailableError contract that backs it. Also bump graphrag_sdk.__version__, which was left at 1.2.0 when 1.3.0 was cut — it is what the API's /version endpoint reports. Minor bump — the cache is opt-in (cache_unchanged_chunks=False by default) and existing ingest paths are unaffected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
What
Two changes, both aimed at making
update()cheap and reliable.1. Chunk-level extraction cache (opt-in:
cache_unchanged_chunks=True)A chunk whose text is byte-identical (SHA-256) to one already stored for the same document skips LLM extraction entirely. Its entities, relationships and mentions are rebuilt from the live graph and remapped onto the new chunk uids.
Editing one paragraph of a 50-chunk document now costs ~1 extraction instead of 50.
This extends the existing document-level no-op short-circuit down to chunk granularity, reusing the
extractor=hookupdate()already exposes. No newupdate()function — the existing one gains a flag.2. Fix: config validation rejected provider-prefixed model names
EmbedderConfig.to_embedder()buildsazure/text-embedding-3-largewhen a provider is set, but the name stored on the graph is unprefixed._validate_graph_configcompared the two raw strings — so pointing an existing graph at a provider (same model, same dimensions) raisedConfigErrorand broke every update from then on.Both sides are now normalized before comparing: a leading
<provider>/is stripped when the prefix is a known litellm provider. The check itself stays — a genuine model swap silently corrupts vector search, since old and new embeddings are never comparable and nothing ever raises.How
CachedChunkExtraction(new, exported) — a decoratorExtractionStrategy. The graph is the cache; nothing extra is stored.GraphStoreread accessors +ChunkEntityRow/ChunkRelationshipRow— schema coupling stays in the storage layer.cache_unchanged_chunks: bool = Falseonupdate(),update_sync(),apply_changes(),apply_changes_sync(). Stats land inUpdateResult.metadata["cache_stats"]._bare_model_name()— provider list read from litellm once and cached, with a static fallback.Safety
SET n +=preserves embeddings).source_chunk_idsremap: own old ids → new uid(s) or dropped; other documents' ids pass through. RELATES provenance is UNIONed.Caveats (also on the docstring + CHANGELOG)
False; existing behavior is byte-identical until you opt in.Testing
1056 passed, 0 failed; ruff clean.
Verified end to end against a live server + FalkorDB over the real HTTP path: 18/19 chunks reused (95%) on a one-paragraph edit, with the changed chunk correctly re-extracted.
Summary by CodeRabbit
CachedChunkExtractionstrategy.