fix(graphrag): give a failed chunk one more attempt before giving up

A chunk whose extraction failed was marked failed and the build moved on. The
checkpoint treats failed chunks as pending, but nothing ever reran the build,
so one transient error -- a provider hiccup, a single response that did not
parse -- left a permanent hole in the graph until someone rebuilt the whole
source.

Failed chunks now get one more attempt after the rest of the build, so a burst
of rate limiting has time to pass. Extraction errors, unparseable responses and
failed writes are all retried; a chunk is marked failed only when its retry
fails too, so it costs at most two calls. The chunks given up on are logged by
id, and progress still ends at the total.

docsgpt/graphrag/extraction.py

@@ -170,13 +170,13 @@ def _extract_chunk(
         )
     except Exception as exc:
         logger.warning(
-            "Graph extraction call failed for chunk %s, skipping: %s", chunk_id, exc
+            "Graph extraction call failed for chunk %s: %s", chunk_id, exc
         )
         return None
     parsed = _parse_extraction(response)
     if parsed is None:
         logger.warning(
-            "Graph extraction returned unparseable output for chunk %s; marking it failed.",
+            "Graph extraction returned unparseable output for chunk %s.",
             chunk_id,
         )
     return parsed
@@ -203,8 +203,10 @@ def extract_graph_for_source(
     Resumable and idempotent: chunks already marked ``done`` are skipped via the
     ``graph_ingest_progress`` checkpoint, so a retry never re-extracts (and never
     re-bills). Processes at most the resolved chunk cap; excess chunks are
-    reported under ``skipped_over_cap``. A malformed response or an LLM error on
-    a single chunk marks it ``failed`` and continues — the pipeline never crashes.
+    reported under ``skipped_over_cap``. A malformed response, an LLM error or a
+    failed write on a single chunk is retried once after the rest of the build;
+    a chunk that fails again is marked ``failed`` and the run continues — the
+    pipeline never crashes.
 
     Each chunk is written in a single transaction with one batched embedding
     call (entity + relationship-endpoint names together).
@@ -273,13 +275,9 @@ def extract_graph_for_source(
         """One chunk's LLM extraction — the only step run concurrently.
 
         A chunk spends almost all of its time waiting on the model, so that is
-        what runs in the pool. Everything else stays on the calling thread:
-        graph writes, so transactions and the progress checkpoint are exactly
-        what they were serially, and embedding. Inside a Celery worker the
-        embeddings client decides to embed locally from the task on the
-        *current thread's* stack; a pool thread has none, so it would instead
-        dispatch an embed task to the worker and wait on it, which Celery
-        refuses inside a task — failing every chunk of the build.
+        what runs in the pool. Graph writes and embedding stay on the calling
+        thread, so transactions and the progress checkpoint are exactly what
+        they were serially and the pool never touches the embeddings client.
         """
         chunk, chunk_id = item
         text = _chunk_text(chunk)
@@ -294,56 +292,85 @@ def extract_graph_for_source(
             relationships = _build_relationships(extracted["relationships"])
         except Exception as exc:
             logger.warning(
-                "Graph extraction failed for chunk %s, skipping: %s", chunk_id, exc
+                "Graph extraction failed for chunk %s: %s", chunk_id, exc
             )
             return chunk_id, "failed", None
         return chunk_id, "ok", (entities, relationships)
 
+    def _write(chunk_id, status, payload) -> bool:
+        """Apply one prepared chunk to the graph; False when it did not land."""
+        nonlocal node_upserts, edges, chunks_processed
+        if status == "empty":
+            store.mark_chunk(source_id, chunk_id, "done")
+            chunks_processed += 1
+            return True
+        if status == "failed":
+            return False
+
+        entities, relationships = payload
+        try:
+            name_embeddings = _embed_names(embedding, entities, relationships)
+            _embed_facts(embedding, relationships)
+            chunk_nodes, chunk_edges = store.apply_chunk(
+                source_id, chunk_id, entities, relationships, name_embeddings
+            )
+        except Exception as exc:
+            logger.warning(
+                "Graph extraction embed/write failed for chunk %s: %s", chunk_id, exc
+            )
+            return False
+        # ``apply_chunk`` marks the chunk done inside the transaction that
+        # writes its rows, so the checkpoint cannot disagree with the graph
+        # and a replayed write cannot apply the chunk twice.
+        node_upserts += chunk_nodes
+        edges += chunk_edges
+        chunks_processed += 1
+        return True
+
     workers = max(1, int(getattr(settings, "GRAPHRAG_EXTRACTION_WORKERS", 1) or 1))
     pool = None
     if workers > 1 and len(to_process) > 1:
         pool = ThreadPoolExecutor(max_workers=workers)
-        # ``map`` yields in submission order, so chunks are still applied in the
-        # order they were given and a run stays reproducible.
-        prepared = pool.map(_prepare, to_process)
-    else:
-        prepared = (_prepare(item) for item in to_process)
+
+    def _pass(items):
+        """Extract and write ``items``; return the ones that did not land."""
+        if pool is not None:
+            # ``map`` yields in submission order, so chunks are still applied in
+            # the order they were given and a run stays reproducible.
+            prepared = pool.map(_prepare, items)
+        else:
+            prepared = (_prepare(item) for item in items)
+        missed = []
+        for item, (chunk_id, status, payload) in zip(items, prepared):
+            if not _write(chunk_id, status, payload):
+                missed.append(item)
+            _report()
+        return missed
 
     try:
-        for chunk_id, status, payload in prepared:
-            if status == "empty":
-                store.mark_chunk(source_id, chunk_id, "done")
-                chunks_processed += 1
-                _report()
-                continue
-            if status == "failed":
-                store.mark_chunk(source_id, chunk_id, "failed")
-                failed_chunks += 1
-                _report()
-                continue
-
-            entities, relationships = payload
-            try:
-                # On this thread, not in the pool — see ``_prepare``.
-                name_embeddings = _embed_names(embedding, entities, relationships)
-                _embed_facts(embedding, relationships)
-                chunk_nodes, chunk_edges = store.apply_chunk(
-                    source_id, chunk_id, entities, relationships, name_embeddings
-                )
-                node_upserts += chunk_nodes
-                edges += chunk_edges
-                # ``apply_chunk`` marks the chunk done inside the transaction that
-                # writes its rows, so the checkpoint cannot disagree with the graph
-                # and a replayed write cannot apply the chunk twice.
-                chunks_processed += 1
-            except Exception as exc:
-                logger.warning(
-                    "Graph extraction embed/write failed for chunk %s, skipping: %s",
-                    chunk_id,
-                    exc,
-                )
-                store.mark_chunk(source_id, chunk_id, "failed")
-                failed_chunks += 1
+        missed = _pass(to_process)
+        if missed:
+            # A failure is usually transient — a provider error, one response
+            # that did not parse — and the checkpoint only picks it up on a
+            # rerun nothing schedules. One more attempt, after the rest of the
+            # build so a burst of rate limiting has passed, and no more: a
+            # chunk that cannot be extracted costs at most two calls.
+            logger.info(
+                "Graph extraction retrying %d failed chunk(s) for source %s",
+                len(missed),
+                source_id,
+            )
+            missed = _pass(missed)
+        for _, chunk_id in missed:
+            store.mark_chunk(source_id, chunk_id, "failed")
+        failed_chunks = len(missed)
+        if missed:
+            logger.warning(
+                "Graph extraction gave up on %d chunk(s) for source %s after a retry: %s",
+                failed_chunks,
+                source_id,
+                ", ".join(str(chunk_id) for _, chunk_id in missed),
+            )
             _report()
     finally:
         if pool is not None:

tests/graphrag/test_extraction.py

@@ -88,6 +88,37 @@ class _StubLLM:
         return response
 
 
+class _ScriptedLLM:
+    """Stub LLM answering per chunk, so results do not depend on call order.
+
+    The extraction pool runs calls concurrently, so a stub that hands out
+    responses in call order gives each chunk whichever response its thread
+    happened to grab first. ``script`` maps a chunk's text to the responses
+    for that chunk, consumed one per call.
+    """
+
+    def __init__(self, script):
+        self._script = {text: list(responses) for text, responses in script.items()}
+        self.model_id = "stub-model"
+        self.calls = []
+        self._token_usage_source = None
+        self._request_id = None
+
+    def gen(self, model=None, messages=None, **kwargs):
+        text = messages[-1]["content"].removeprefix("<chunk>\n").removesuffix("\n</chunk>")
+        self.calls.append(text)
+        responses = self._script.get(text)
+        if not responses:
+            raise AssertionError(f"unexpected extraction call for {text!r}")
+        response = responses.pop(0)
+        if isinstance(response, Exception):
+            raise response
+        return response
+
+    def calls_for(self, text):
+        return self.calls.count(text)
+
+
 class _StubEmbedding:
     """Stub embeddings model producing deterministic fixed-dim vectors."""
 
@@ -365,11 +396,12 @@ class TestExtractionLive:
     ):
         """Only the LLM call may run in the extraction pool, never embedding.
 
-        Inside a Celery worker the embeddings client decides to embed locally
-        from the task on the *current thread's* stack. A pool thread has none,
-        so from there it dispatches an embed task to the worker and waits on
-        it — which Celery refuses inside a task, so every chunk of a graph
-        build failed.
+        Embedding from the pool is what broke every graph build inside a
+        worker: the embeddings client used to decide "embed locally" from the
+        task on the *current thread's* stack, so a pool thread dispatched to
+        the worker instead and Celery refused the wait. ``in_worker`` is
+        process-wide now, but the pool still has no reason to touch the
+        embeddings client — it exists to overlap model latency.
         """
         import threading
 
@@ -505,11 +537,13 @@ class TestExtractionLive:
                 entities=[{"name": "Ada", "type": "person", "description": "d"}],
                 relationships=[],
             )
-            llm = _StubLLM([
-                "not json at all",
-                RuntimeError("model exploded"),
-                good,
-            ])
+            # Each failing chunk fails its retry too; one that recovers on
+            # retry is covered in ``TestFailedChunksAreRetried``.
+            llm = _ScriptedLLM({
+                "garbage": ["not json at all", "still not json"],
+                "boom": [RuntimeError("model exploded"), RuntimeError("model exploded again")],
+                "Ada.": [good],
+            })
             _install_stub_llm(monkeypatch, llm)
 
             summary = extract_graph_for_source(
@@ -762,7 +796,7 @@ class TestFailedChunksAreReported:
         import logging
 
         store = self._fake_store(monkeypatch, ["c1"])
-        _install_stub_llm(monkeypatch, _StubLLM(["not json at all"]))
+        _install_stub_llm(monkeypatch, _StubLLM(["not json at all", "still not json"]))
 
         with caplog.at_level(logging.WARNING, logger="docsgpt.graphrag.extraction"):
             summary = extract_graph_for_source(
@@ -785,7 +819,10 @@ class TestFailedChunksAreReported:
         import logging
 
         self._fake_store(monkeypatch, ["c7"])
-        _install_stub_llm(monkeypatch, _StubLLM([RuntimeError("model exploded")]))
+        _install_stub_llm(
+            monkeypatch,
+            _StubLLM([RuntimeError("model exploded"), RuntimeError("model exploded again")]),
+        )
 
         with caplog.at_level(logging.WARNING, logger="docsgpt.graphrag.extraction"):
             extract_graph_for_source(
@@ -800,6 +837,126 @@ class TestFailedChunksAreReported:
         assert any("c7" in message for message in messages), messages
 
 
+@pytest.mark.unit
+class TestFailedChunksAreRetried:
+    """A chunk that fails once gets one more attempt before the build ends.
+
+    Failures are recorded as ``failed`` and the checkpoint treats them as
+    pending, but nothing ever ran the build again, so a single transient error
+    — a provider hiccup, one response that did not parse — left a permanent
+    hole in the graph until someone rebuilt the whole source. Retries run
+    after the rest of the build, which gives a burst of rate limiting time to
+    pass, and are bounded at one per chunk so a chunk that can never be
+    extracted costs at most two calls.
+    """
+
+    GOOD = _extraction_json(
+        entities=[{"name": "Ada", "type": "person", "description": "d"}],
+        relationships=[],
+    )
+
+    def _fake_store(self, monkeypatch, chunk_ids):
+        from unittest.mock import MagicMock
+
+        store = MagicMock(name="GraphStore")
+        store.pending_chunks.return_value = list(chunk_ids)
+        store.apply_chunk.return_value = (1, 0)
+        store.count_nodes.return_value = 1
+        monkeypatch.setattr(
+            "docsgpt.graphrag.store.GraphStore", lambda *a, **k: store
+        )
+        return store
+
+    def _run(self, chunks, progress=None):
+        return extract_graph_for_source(
+            str(uuid.uuid4()),
+            user="owner-1",
+            chunks=chunks,
+            config=SourceConfig(),
+            request_id="req-retry",
+            progress_cb=progress,
+        )
+
+    @staticmethod
+    def _marked_failed(store):
+        return [c.args[1] for c in store.mark_chunk.call_args_list if c.args[2] == "failed"]
+
+    def test_a_transient_failure_is_retried_and_written(self, monkeypatch, stub_embedding):
+        store = self._fake_store(monkeypatch, ["c1"])
+        llm = _ScriptedLLM({"flaky": [RuntimeError("rate limited"), self.GOOD]})
+        _install_stub_llm(monkeypatch, llm)
+
+        summary = self._run([_chunk("c1", "flaky")])
+
+        assert summary["failed_chunks"] == 0
+        assert summary["chunks_processed"] == 1
+        assert store.apply_chunk.call_args.args[1] == "c1"
+        assert self._marked_failed(store) == []
+
+    def test_an_unparseable_response_is_retried(self, monkeypatch, stub_embedding):
+        store = self._fake_store(monkeypatch, ["c1"])
+        _install_stub_llm(monkeypatch, _ScriptedLLM({"odd": ["not json", self.GOOD]}))
+
+        summary = self._run([_chunk("c1", "odd")])
+
+        assert summary["failed_chunks"] == 0
+        assert self._marked_failed(store) == []
+
+    def test_a_chunk_that_fails_again_is_marked_failed_once(self, monkeypatch, stub_embedding):
+        store = self._fake_store(monkeypatch, ["c1"])
+        llm = _ScriptedLLM({"broken": ["not json", "still not json"]})
+        _install_stub_llm(monkeypatch, llm)
+
+        summary = self._run([_chunk("c1", "broken")])
+
+        assert summary["failed_chunks"] == 1
+        assert summary["chunks_processed"] == 0
+        assert llm.calls_for("broken") == 2
+        assert self._marked_failed(store) == ["c1"]
+
+    def test_only_failed_chunks_are_retried(self, monkeypatch, stub_embedding):
+        from docsgpt.core.settings import settings
+
+        monkeypatch.setattr(settings, "GRAPHRAG_EXTRACTION_WORKERS", 4)
+        self._fake_store(monkeypatch, ["c1", "c2", "c3"])
+        llm = _ScriptedLLM({
+            "one": [self.GOOD],
+            "two": [RuntimeError("timeout"), self.GOOD],
+            "three": [self.GOOD],
+        })
+        _install_stub_llm(monkeypatch, llm)
+
+        summary = self._run([_chunk("c1", "one"), _chunk("c2", "two"), _chunk("c3", "three")])
+
+        assert summary["chunks_processed"] == 3
+        assert summary["failed_chunks"] == 0
+        assert (llm.calls_for("one"), llm.calls_for("two"), llm.calls_for("three")) == (1, 2, 1)
+
+    def test_a_failed_write_is_retried(self, monkeypatch, stub_embedding):
+        store = self._fake_store(monkeypatch, ["c1"])
+        store.apply_chunk.side_effect = [RuntimeError("write failed"), (1, 0)]
+        _install_stub_llm(monkeypatch, _ScriptedLLM({"text": [self.GOOD, self.GOOD]}))
+
+        summary = self._run([_chunk("c1", "text")])
+
+        assert summary["failed_chunks"] == 0
+        assert summary["chunks_processed"] == 1
+        assert self._marked_failed(store) == []
+
+    def test_progress_ends_at_the_total(self, monkeypatch, stub_embedding):
+        self._fake_store(monkeypatch, ["c1", "c2"])
+        _install_stub_llm(monkeypatch, _ScriptedLLM({
+            "fine": [self.GOOD],
+            "broken": ["not json", "still not json"],
+        }))
+        events = []
+
+        self._run([_chunk("c1", "fine"), _chunk("c2", "broken")], progress=events.append)
+
+        assert all(e["current"] <= e["total"] for e in events)
+        assert events[-1]["current"] == events[-1]["total"] == 2
+
+
 @pytest.mark.integration
 class TestSummaryNodeCount:
     """``nodes`` must describe the graph, not the number of upserts."""