fix(core): build KG from ainsert_custom_chunks extraction (#3352) by ysys143 · Pull Request #3353 · HKUDS/LightRAG

ysys143 · 2026-07-04T07:08:06Z

Summary

ainsert_custom_chunks fires entity extraction inside an asyncio.gather and discards its return value, so merge_nodes_and_edges never runs. The extracted entities/relationships are dropped: the knowledge graph and entities_vdb / relationships_vdb stay empty while the extraction LLM calls (and their cost/latency) are still spent. KG-dependent query modes (local / global / hybrid / mix) then return no context.

The normal file pipeline (lightrag/pipeline.py) captures the extraction result and merges it via merge_nodes_and_edges. ainsert_custom_chunks omitted that step.

There is also a secondary error-masking defect: extraction is called without pipeline_status / pipeline_status_lock (both default to None), so if extraction raises, the except block runs async with pipeline_status_lock: on a None lock and raises TypeError: 'NoneType' object does not support the asynchronous context manager protocol, masking the real error.

Changes

lightrag/lightrag.py:

Persist chunks/docs first (so text_chunks exists before extraction reads it), then capture the _process_extract_entities result and pass it to merge_nodes_and_edges — mirroring the file pipeline. Import merge_nodes_and_edges from lightrag.operate.
Pass pipeline_status / pipeline_status_lock down to _process_extract_entities, and guard its except block against a None lock so the real extraction error surfaces.

tests/pipeline/test_custom_chunks_kg_merge.py (new) — regression tests:

extracted chunk_results are handed to merge_nodes_and_edges (with the right graph/vdb/doc_id);
no entities extracted → merge is not called;
with pipeline_status_lock=None, a failing extraction surfaces the real error, not a TypeError.

tests/pipeline/test_document_file_path_normalization.py — minimal update of the existing custom-chunks test for the new extraction call signature (status/lock are now passed through).

Behavior (before → after), same instance / same text

ainsert_custom_chunks of two chunks about Marie/Pierre Curie:

Before — extraction runs, graph stays empty:

Chunk 1 of 2 extracted 4 Ent + 3 Rel
Chunk 2 of 2 extracted 3 Ent + 3 Rel
Writing graph with 0 nodes, 0 edges        # no merge ran
get_all_labels() -> 0
hybrid query -> [no-context]

After — extracted entities are merged, matching what ainsert produces for the same text:

Chunk N of M extracted X Ent + Y Rel
Merging stage ... / Completed merging: <N> entities, <M> relations
Writing graph with <N> nodes, <M> edges
get_all_labels() -> <N>
hybrid query -> grounded answer

Test verification

New tests pass against the fix; all three fail against the pre-fix code (confirming they pin the defect).
ruff check clean on the changed files.
Reproduced live end-to-end with real Gemini (gemini-flash-lite-latest + gemini-embedding-001) — details and full logs in [Bug]: ainsert_custom_chunks runs entity extraction but never merges it into the KG #3352.

danielaskdd · 2026-07-06T10:30:16Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4798dc5e8a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

ainsert_custom_chunks fired entity extraction inside asyncio.gather and discarded its return value, so merge_nodes_and_edges never ran: the knowledge graph and entity/relationship vector stores stayed empty while the extraction LLM calls were still spent. KG-dependent query modes then returned no context. Capture the extraction result and merge it into the KG (mirroring the file pipeline in pipeline.py), persisting chunks/docs first so text_chunks exists before extraction reads it. Also pass pipeline_status / pipeline_status_lock down to extraction and guard _process_extract_entities' except block against a None lock, so an extraction failure surfaces the real error instead of a TypeError from `async with None`. Add regression tests covering both defects.

The file-path normalization tests carried no `offline` marker, so CI's `pytest -m offline` leg silently deselected the whole file. The custom-chunks test updated in HKUDS#3352 also called `initialize_pipeline_status` on a LightRAG built via `__new__` (which bypasses `initialize_storages`), without first initializing the shared-data dictionaries — raising `ValueError: Shared dictionaries not initialized`, a failure CI never exercised. Mark the module `offline` so the file runs in CI, and call `initialize_share_data()` before `initialize_pipeline_status()`.

If merge_nodes_and_edges (or extraction) raises, ainsert_custom_chunks previously still flushed the full_docs/text_chunks/chunks_vdb base writes via the finally block. A retry with the same doc_id then short-circuits at the duplicate checks ("already in the storage" / "All chunks are already in the storage"), leaving the document stored with an empty/incomplete KG and no way to repair it — the exact state this method exists to prevent. Wrap the write path in try/except/else: commit via _insert_done_with_cleanup only on success; on failure delete this call's just-written doc/chunk records so a retry rebuilds the KG. Flush the rollback with plain _insert_done, not _insert_done_with_cleanup — the cleanup helper discards pending ops on IndexFlushError, which for buffered-delete backends (e.g. OpenSearch KV) would drop the rollback deletes and re-block the retry, and leave orphan chunk vectors in buffered vector stores (its docstring warns against deletion-path use). Bind doc_key / inserting_chunks before the try so the rollback can always reference them. Adds a regression test that a merge failure leaves the insert retryable. Raised in the Codex review of HKUDS#3353.

danielaskdd · 2026-07-06T13:29:24Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ff68f855bd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

A KG merge/extraction failure previously left the insert in a broken state: the base full_docs/text_chunks/chunks_vdb writes were flushed by the finally block, so a retry with the same doc_id short-circuited at the duplicate checks with an empty/incomplete KG and no way to repair it. Adopt the file pipeline's doc_status model instead of hand-rolling a rollback. (An earlier delete-then-flush attempt leaked partial KG: merge_nodes_and_edges writes graph/vdb records it does not roll back, and flushing on failure persisted them with source_ids pointing at the just-deleted chunks.) - Gate reprocessing on doc_status (PROCESSED -> no-op; absent/FAILED/crashed PROCESSING -> rebuild), not on full_docs/text_chunks presence. Chunks are never deleted, so a partial node's source_id always resolves to a live chunk. - Mark PROCESSING up front (immediate-write), PROCESSED only after a clean flush, FAILED + discard (never flush) on failure. On a flush error the doc stays PROCESSING and is rebuilt on retry, mirroring the pipeline's flush-failure recovery. - Extraction is no longer inside the base-write asyncio.gather, so a sibling upsert racing after the failure handler can no longer strand state (the remaining gather writes are idempotent). Regression tests cover: retry rebuilds after a merge failure; the failed doc is FAILED with its chunks kept (no dangling source_id); a partial-merge failure is not orphaned; a success-path flush failure leaves the doc rebuildable rather than PROCESSED-but-empty; an empty extraction still completes. Addresses the Codex review findings on HKUDS#3353.

A KG merge/extraction failure previously left the insert in a broken state: the base full_docs/text_chunks/chunks_vdb writes were flushed by the finally block, so a retry with the same doc_id short-circuited at the duplicate checks with an empty/incomplete KG and no way to repair it. Adopt the file pipeline's doc_status model instead of hand-rolling a rollback. (An earlier delete-then-flush attempt leaked partial KG: merge_nodes_and_edges writes graph/vdb records it does not roll back, and flushing on failure persisted them with source_ids pointing at the just-deleted chunks.) - Gate reprocessing on doc_status (PROCESSED -> no-op; absent/FAILED/crashed PROCESSING -> rebuild), not on full_docs/text_chunks presence. Chunks are never deleted, so a partial node's source_id always resolves to a live chunk. - Mark PROCESSING up front (immediate-write), PROCESSED only after a clean flush, FAILED + discard (never flush) on failure. On a flush error the doc stays PROCESSING and is rebuilt on retry. FAILED-mark and discard are guarded independently so a failure in one still runs the other. - Before rebuilding, purge the prior attempt's chunks + KG via _purge_doc_chunks_and_kg (mirroring the pipeline's resume purge) so a retry of a completed-but-not-PROCESSED merge rebuilds cleanly instead of re-appending descriptions or orphaning chunks. - Extraction is no longer inside the base-write asyncio.gather, so a sibling upsert racing after the failure handler can no longer strand state. Known limitation (shared with the file pipeline, tracked separately): a merge that writes partial graph state then fails before its full_entities index row is written leaves nothing for purge to find, so a retry re-merges and entity descriptions accumulate (bounded; collapsed by later re-summarization). The proper fix is cross-dedup of stored vs new descriptions inside merge_nodes_and_edges. Regression tests: retry rebuilds after a merge failure; the failed doc is FAILED with its chunks kept; a partial-merge failure is not orphaned; retry purges the prior chunks_list; a first insert does not purge; a success-path flush failure leaves the doc rebuildable; an empty extraction still completes. Addresses the Codex review findings on HKUDS#3353.

A KG merge/extraction failure previously left the insert in a broken state: the base full_docs/text_chunks/chunks_vdb writes were flushed by the finally block, so a retry with the same doc_id short-circuited at the duplicate checks with an empty/incomplete KG and no way to repair it. Adopt the file pipeline's doc_status model instead of hand-rolling a rollback. (An earlier delete-then-flush attempt leaked partial KG: merge_nodes_and_edges writes graph/vdb records it does not roll back, and flushing on failure persisted them with source_ids pointing at the just-deleted chunks.) - Gate reprocessing on doc_status (PROCESSED -> no-op; absent/FAILED/crashed PROCESSING -> rebuild), not on full_docs/text_chunks presence. Chunks are never deleted, so a partial node's source_id always resolves to a live chunk. - Mark PROCESSING (immediate-write), PROCESSED only after a clean flush, FAILED + discard (never flush) on failure. On a flush error the doc stays PROCESSING and is rebuilt on retry. FAILED-mark and discard are guarded independently. - Before rebuilding, purge the prior attempt's chunks + KG via _purge_doc_chunks_and_kg (mirroring the pipeline resume purge), and do it BEFORE overwriting doc_status with the new chunks_list — so a purge failure preserves the prior chunks_list as the recovery anchor instead of stranding the old chunks/KG under new ids. - Extraction is no longer inside the base-write asyncio.gather, so a sibling upsert racing after the failure handler can no longer strand state. Known limitations (documented in code, tracked separately): - A merge that writes partial graph state then fails before its full_entities index row is written leaves nothing for purge to find, so a retry re-merges and descriptions accumulate (bounded; a merge_nodes_and_edges property shared with the file pipeline's resume purge). - No per-doc serialization: concurrent inserts with the same doc_id can interleave. Single-writer use only (the file pipeline serializes via enqueue locks; this path never had one). Regression tests: retry rebuilds after a merge failure; failed doc is FAILED with chunks kept; partial-merge failure is not orphaned; retry purges the prior chunks_list; a purge failure preserves the prior chunks_list anchor; a first insert does not purge; a success-path flush failure leaves the doc rebuildable; an empty extraction still completes. Addresses the Codex review findings on HKUDS#3353.

ysys143 · 2026-07-07T04:35:01Z

Pausing this PR.

The original #3352 fix (extraction result was discarded, so the KG was never built) stands. But making the retry/reprocess path fully robust runs into shared merge/purge core limitations rather than anything specific to this method:

entity/relation descriptions accumulate on any reprocess/resume (no cross-dedup of stored vs new) — filed as [Bug]: merge_nodes_and_edges accumulates entity/relation descriptions on any reprocess/resume (file pipeline + ainsert, not just custom chunks) #3367;
when full_entities isn't persisted (a partial merge, or a flush that drops only the full_entities row), _purge_doc_chunks_and_kg deletes the chunks by id but can't find the graph nodes (it discovers them via full_entities), leaving nodes whose source_id points at deleted chunks.

Both affect the file pipeline and ainsert too, since they share merge_nodes_and_edges / _purge_doc_chunks_and_kg. They belong in the core, not in this deprecated direct path. Leaving this for a maintainer to decide whether it's worth taking further.

danielaskdd · 2026-07-15T11:08:08Z

Hi @ysys143,

The design below is intended to be the comprehensive solution to the shared merge/purge partial-failure and idempotent-recovery limitations that led to pausing this PR:

Proposed design and implementation plan: recoverable document mutations

Since you identified and analyzed the blockers in #3353, could you please help evaluate whether the proposed invariants, write-ahead recovery indexes, custom-chunk patch journal, SDK-driven idempotent roll-forward, and scan rollback are reasonable, and whether any failure modes are still missing? Your feedback would be valuable before implementation starts.

Thanks!

danielaskdd · 2026-07-17T08:14:04Z

Closing this PR as #3416 has been merged.

chatgpt-codex-connector Bot reviewed Jul 6, 2026

View reviewed changes

Comment thread lightrag/lightrag.py

ysys143 added 2 commits July 6, 2026 20:30

ysys143 force-pushed the fix/ainsert-custom-chunks-kg-merge branch from 4798dc5 to ff68f85 Compare July 6, 2026 11:33

chatgpt-codex-connector Bot reviewed Jul 6, 2026

View reviewed changes

Comment thread lightrag/lightrag.py Outdated

Comment thread lightrag/lightrag.py

ysys143 force-pushed the fix/ainsert-custom-chunks-kg-merge branch from ff68f85 to eba27cc Compare July 7, 2026 02:14

ysys143 force-pushed the fix/ainsert-custom-chunks-kg-merge branch from eba27cc to 0bd64c6 Compare July 7, 2026 02:52

ysys143 force-pushed the fix/ainsert-custom-chunks-kg-merge branch from 0bd64c6 to 784ff10 Compare July 7, 2026 03:17

ysys143 mentioned this pull request Jul 7, 2026

[Bug]: merge_nodes_and_edges accumulates entity/relation descriptions on any reprocess/resume (file pipeline + ainsert, not just custom chunks) #3367

Closed

2 tasks

danielaskdd marked this pull request as draft July 15, 2026 02:46

danielaskdd mentioned this pull request Jul 15, 2026

Partial failures in KG merge/purge can leave cross-store inconsistencies and orphan graph data #3400

Closed

danielaskdd mentioned this pull request Jul 17, 2026

feat(core): recoverable document mutations — write-ahead anchors, journaled custom chunks, scan rollback (#3400, all 5 phases) #3416

Merged

5 tasks

danielaskdd closed this Jul 17, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(core): build KG from ainsert_custom_chunks extraction (#3352)#3353

fix(core): build KG from ainsert_custom_chunks extraction (#3352)#3353
ysys143 wants to merge 3 commits into
HKUDS:mainfrom
ysys143:fix/ainsert-custom-chunks-kg-merge

ysys143 commented Jul 4, 2026

Uh oh!

danielaskdd commented Jul 6, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

danielaskdd commented Jul 6, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

ysys143 commented Jul 7, 2026

Uh oh!

danielaskdd commented Jul 15, 2026

Uh oh!

danielaskdd commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ysys143 commented Jul 4, 2026

Summary

Changes

Behavior (before → after), same instance / same text

Test verification

Uh oh!

danielaskdd commented Jul 6, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

danielaskdd commented Jul 6, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

ysys143 commented Jul 7, 2026

Uh oh!

danielaskdd commented Jul 15, 2026

Uh oh!

danielaskdd commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants