Skip to content

[Bug]: ainsert_custom_chunks runs entity extraction but never merges it into the KG #3352

Description

@ysys143

Do you need to file an issue?

  • I have searched the existing issues and this bug is not already filed.
  • I believe this is a legitimate bug, not just a question or feature request.

Describe the bug

LightRAG.ainsert_custom_chunks(full_text, text_chunks, doc_id) is expected to ingest caller-supplied chunks and build the knowledge graph from them (the way ainsert does for its own chunks). In practice it embeds the chunks but silently discards the extracted entities and relationships, so no graph is built. KG-dependent query modes then return nothing, while the extraction LLM calls (and their cost/latency) were still spent.

Root cause — in lightrag/lightrag.py, ainsert_custom_chunks fires extraction inside an asyncio.gather and throws away its return value; there is no merge_nodes_and_edges call:

# lightrag/lightrag.py  (~L1559)
tasks = [
    self.chunks_vdb.upsert(inserting_chunks),
    self._process_extract_entities(inserting_chunks),   # <-- returns chunk_results, DISCARDED
    self.full_docs.upsert(new_docs),
    self.text_chunks.upsert(inserting_chunks),
]
await asyncio.gather(*tasks)
# ... finally: await self._insert_done_with_cleanup()   # flush only; nothing was written to the graph

_process_extract_entities itself is correct — it returns the extracted chunk_results. The normal file pipeline (lightrag/pipeline.py) captures that return and merges it via merge_nodes_and_edges, which is the step that actually writes the graph + entity/relationship vector stores. ainsert_custom_chunks omits that step, so the graph and entities_vdb / relationships_vdb are never populated; _insert_done_with_cleanup only flushes storages and there is nothing to flush for the graph.

Secondary defect (error masking)ainsert_custom_chunks calls self._process_extract_entities(inserting_chunks) without pipeline_status / pipeline_status_lock (both default to None). If extraction raises, the except block does async with pipeline_status_lock: on a None lock, raising TypeError: 'NoneType' object does not support the asynchronous context manager protocol, which masks the original extraction error.

Steps to reproduce

Verified live on lightrag-hku==1.5.4 with real Gemini (LLM + embedding). The script below ingests the same caller-supplied chunks two ways on two fresh storesainsert_custom_chunks (A) vs ainsert (B) — so the only variable is whether the KG merge runs. A full runnable version (with the exact LightRAG(...) config) is in LightRAG Config Used below.

import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.kg.shared_storage import initialize_pipeline_status

async def main():
    rag = LightRAG(
        working_dir="./repro_store",
        llm_model_func=...,          # any working chat LLM
        embedding_func=...,          # any working embedding
    )
    await rag.initialize_storages()
    await initialize_pipeline_status()

    chunks = [
        "Marie Curie discovered polonium and radium. She won two Nobel Prizes.",
        "Pierre Curie collaborated with Marie Curie on radioactivity research.",
    ]
    await rag.ainsert_custom_chunks(full_text="\n\n".join(chunks), text_chunks=chunks, doc_id="doc-1")
    # Logs show: "Chunk N of M extracted X Ent + Y Rel"  then  "Writing graph with 0 nodes, 0 edges"

    labels = await rag.chunk_entity_relation_graph.get_all_labels()
    print("entities in graph:", len(labels))   # -> 0  (BUG: extraction was discarded)

    ans = await rag.aquery("Who did Marie Curie work with?", param=QueryParam(mode="hybrid"))
    print(ans)                                  # -> "[no-context]" / cannot answer
    await rag.finalize_storages()

asyncio.run(main())

Empirical evidence — ingesting 85 caller-supplied chunks of one PDF via ainsert_custom_chunks:

  • Ingest logs show extraction ran (e.g. Chunk 85 of 85 extracted 19 Ent + 17 Rel) and chunks flush: embedding 85 vectors.
  • On-disk stores afterwards: graph_chunk_entity_relation.graphml0 nodes, 0 edges; vdb_entities.json → empty; vdb_chunks.json → 85 chunk vectors.
  • Query mode="hybrid"Raw search results: 0 entities, 0 relations, 0 vector chunks[no-context].
  • Query mode="naive" → grounded answer (only the chunk path survived).
  • The same document ingested via ainsert(... split_by_character=...) on the same instance produced ~500 entities / ~511 edges and a working hybrid query.

Expected Behavior

ainsert_custom_chunks should build the KG from the extracted entities/relationships (like the file pipeline does): after ingest, get_all_labels() returns the extracted entities and local/global/hybrid/mix queries answer from the KG.

LightRAG Config Used

The defect is pure control flow (the KG merge call is simply absent), so it reproduces on any LLM + embedding backend with all-default storages. The exact config used for the reproduction below — a runnable, backend-agnostic setup with real Gemini:

import asyncio, os
from lightrag import LightRAG, QueryParam
from lightrag.utils import EmbeddingFunc
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.llm.gemini import gemini_model_complete, gemini_embed

os.environ.setdefault("GEMINI_API_KEY", "<your-gemini-api-key>")  # or export it before running

# gemini_embed is pre-wrapped @wrap_embedding_func_with_attrs(embedding_dim=1536) but does NOT
# send output_dimensionality to the API, so the API returns 3072 and EmbeddingFunc's count check
# (total_elements // declared_dim) mismatches. Pin embedding_dim=1536 so declared == actual.
EMB_DIM = 1536
async def embed(texts, **_):
    return await gemini_embed(texts, model="gemini-embedding-001", embedding_dim=EMB_DIM)

def make_rag(wd):
    return LightRAG(
        working_dir=wd,
        llm_model_func=gemini_model_complete,
        llm_model_name="gemini-flash-lite-latest",   # read via hashing_kv.global_config
        embedding_func=EmbeddingFunc(embedding_dim=EMB_DIM, max_token_size=2048, func=embed),
        # storages: all library defaults — JsonKVStorage / NanoVectorDBStorage /
        # NetworkXStorage / JsonDocStatusStorage (in-process, no external services)
    )

CHUNKS = [
    "Marie Curie discovered polonium and radium. She won two Nobel Prizes.",
    "Pierre Curie collaborated with Marie Curie on radioactivity research.",
]

async def run(tag, wd, ingest):
    rag = make_rag(wd)
    await rag.initialize_storages()
    await initialize_pipeline_status()
    await ingest(rag)
    labels = await rag.chunk_entity_relation_graph.get_all_labels()
    print(f"[{tag}] get_all_labels() -> {len(labels)} entities: {labels}")
    for mode in ("hybrid", "naive"):
        ans = await rag.aquery("Who did Marie Curie work with?", param=QueryParam(mode=mode))
        print(f"[{tag}] mode={mode:6} -> {ans[:120]!r}")
    await rag.finalize_storages()

async def main():
    await run("A ainsert_custom_chunks", "./repro_custom",
              lambda r: r.ainsert_custom_chunks("\n\n".join(CHUNKS), CHUNKS, doc_id="doc-1"))
    await run("B ainsert (control)", "./repro_normal",
              lambda r: r.ainsert("\n\n".join(CHUNKS)))

asyncio.run(main())

Logs and screenshots

Captured from the two-store run above (LightRAG 1.5.4, real Gemini). Same instance, same embedding, same text — the only difference is the missing merge step.

A — ainsert_custom_chunks (buggy): extraction runs, graph stays empty

INFO: Inserting 1 docs
INFO: Chunk 1 of 2 extracted 4 Ent + 3 Rel chunk-8359c50a...
INFO: Chunk 2 of 2 extracted 3 Ent + 3 Rel chunk-90d52c1a...
INFO: [] chunks flush: embedding 2 vectors in 1 batch(es)
INFO: [] Writing graph with 0 nodes, 0 edges          # <-- no merge ran
#   (no "Merging stage" / "Phase 1/2/3" / "Completed merging" lines appear at all)

[A] get_all_labels() -> 0 entities: []
    graph_chunk_entity_relation.graphml -> 0 nodes, 0 edges
    vdb_chunks.json -> 2 vectors ; vdb_entities.json -> 0 ; vdb_relationships.json -> 0

INFO: Raw search results: 0 entities, 0 relations, 0 vector chunks
INFO: [kg_query] No query context could be built; returning no-result.
INFO: Naive query: 2 chunks (chunk_top_k:20)
[A] mode=hybrid -> "Sorry, I'm not able to provide an answer to that question.[no-context]"
[A] mode=naive  -> "Marie Curie collaborated with Pierre Curie on her research regarding radioactivity. ..."

B — ainsert (control, same text/instance): merge runs, graph populated

INFO: Chunk 1 of 1 extracted 6 Ent + 6 Rel doc-...-chunk-000
INFO: Merging stage 1/1: unknown_source               # <-- present in B, absent in A
INFO: Phase 1: Processing 6 entities ...
INFO: Phase 2: Processing 6 relations ...
INFO: Completed merging: 6 entities, 0 extra entities, 6 relations
INFO: [] Writing graph with 6 nodes, 6 edges

[B] get_all_labels() -> 6 entities: ['Marie Curie','Nobel Prize','Pierre Curie','Polonium','Radioactivity','Radium']
    graph_chunk_entity_relation.graphml -> 6 nodes, 6 edges
    vdb_entities.json -> 6 ; vdb_relationships.json -> 6

INFO: Raw search results: 6 entities, 6 relations, 0 vector chunks
[B] mode=hybrid -> "Marie Curie conducted her research in the field of radioactivity in collaboration with Pierre Curie. ..."
[B] mode=naive  -> "Marie Curie collaborated with Pierre Curie ..."

The absence of the Merging stage / Completed merging log lines in A is the falsifiable signature that merge_nodes_and_edges is never called from ainsert_custom_chunks.

Proposed fix

Capture the extraction result and merge it, mirroring pipeline.py (persist chunks first so text_chunks exists before extraction reads it, then extract, then merge):

await asyncio.gather(
    self.chunks_vdb.upsert(inserting_chunks),
    self.full_docs.upsert(new_docs),
    self.text_chunks.upsert(inserting_chunks),
)
chunk_results = await self._process_extract_entities(inserting_chunks)   # capture, don't discard
if chunk_results:
    pipeline_status = await get_namespace_data("pipeline_status", workspace=self.workspace)
    pipeline_status_lock = get_namespace_lock("pipeline_status", workspace=self.workspace)
    await merge_nodes_and_edges(
        chunk_results=chunk_results,
        knowledge_graph_inst=self.chunk_entity_relation_graph,
        entity_vdb=self.entities_vdb,
        relationships_vdb=self.relationships_vdb,
        global_config=self._build_global_config(),
        full_entities_storage=self.full_entities,
        full_relations_storage=self.full_relations,
        doc_id=doc_key,
        pipeline_status=pipeline_status,
        pipeline_status_lock=pipeline_status_lock,
        llm_response_cache=self.llm_response_cache,
        entity_chunks_storage=self.entity_chunks,
        relation_chunks_storage=self.relation_chunks,
        file_path=file_path,
    )

Also guard the _process_extract_entities except block against a None pipeline_status_lock so an extraction failure surfaces the real error instead of a NoneType async-context-manager TypeError.

Additional Information

Environment

  • LightRAG Version : 1.5.4 — latest release (PyPI lightrag-hku==1.5.4, git tag v1.5.4 @ 9a45b64). Identical ainsert_custom_chunks body on current main @ 2a729d3 (_version.py = 1.5.5), as of 2026-07-03.

  • Storages : JsonKVStorage / NanoVectorDBStorage / NetworkXStorage / JsonDocStatusStorage (defaults)

  • LLM / Embedding : Google Gemini (gemini-flash-lite-latest + gemini-embedding-001, dim 1536)

  • Python : 3.12

  • OS : macOS

  • Related Issues: [Bug]: ainsert_custom_chunks doc status #1199 (different symptom — doc_status not updated by ainsert_custom_chunks), [Question]: ainsert_custom_chunks deprecation #1760 (ainsert_custom_chunks deprecation question)

  • Affected file: lightrag/lightrag.pyainsert_custom_chunks, _process_extract_entities

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions