diff --git a/docsgpt/graphrag/extraction.py b/docsgpt/graphrag/extraction.py index db8ea27a..9539f4fd 100644 --- a/docsgpt/graphrag/extraction.py +++ b/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: diff --git a/tests/graphrag/test_extraction.py b/tests/graphrag/test_extraction.py index eda8dd7c..0573cb7c 100644 --- a/tests/graphrag/test_extraction.py +++ b/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("\n").removesuffix("\n") + 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."""