refactor(extract-worker): stream embedding writes + handle chunking failures

workers/extract-worker/extract_worker/activities.py

@@ -254,7 +254,7 @@ async def embed_documents(
         with cache.get_or_cache_resource(cache_key, factory) as resources:
             chunker, md_converter, embedder = resources
             logger.debug("embedding documents from %s...", batch)
-            response = await embed_documents_act(
+            response, n_chunks_by_doc = await embed_documents_act(
                 batch,
                 chunker=chunker,
                 md_converter=md_converter,
@@ -266,18 +266,23 @@ async def embed_documents(
             )
         logger.debug("embedded %s chunks !", response.n_chunks)
         return EmbedBatchResult(
-            chunks=output_path.relative_to(workdir), response=response
+            chunks=output_path.relative_to(workdir),
+            n_chunks_by_doc=n_chunks_by_doc,
+            response=response,
         )
 
     @activity_defn(name=Activity.INDEX_EMBEDDINGS)
-    async def index_embeddings(self, chunks: Path, project: str, dim: int) -> int:
+    async def index_embeddings(
+        self, chunks: Path, n_chunks_by_doc: dict[DocId, int], project: str, dim: int
+    ) -> int:
         es_client = lifespan_es_client()
         w_config = cast(ExtractWorkerConfig, lifespan_worker_config())
         embedding_config = w_config.embedding
         index = embedding_config.index_name(project)
         logger.debug("indexing embedded chunks into %s...", index)
         n_indexed = await index_embeddings_act(
             w_config.paths.workdir / chunks,
+            n_chunks_by_doc,
             es_client=es_client,
             index=index,
             dim=dim,

workers/extract-worker/extract_worker/embeddings.py

@@ -50,6 +50,8 @@
 
 _DOC_ID = "docId"
 _EMBEDDING = "embedding"
+_CHUNK_INDEX = "chunkIndex"
+_DELETE_GROUP_SIZE = 256
 
 
 class MissingEmbeddingsIndex(ValueError):
@@ -153,13 +155,14 @@ async def embed_documents_act(
     artifacts_root: Path,
     output_path: Path,
     progress: AsyncProgressRateHandler | None = None,
-) -> EmbeddingResponse:
+) -> tuple[EmbeddingResponse, dict[DocId, int]]:
     docs = [d async for d in async_read_jsonl_as(batch, PROCESSED_FILE_TA)]
     n_docs = len(docs)
     if progress is not None:
         progress = to_raw_async_progress(progress, max_progress=n_docs)
     n_pages, n_successes, n_successes_pages, n_chunks = 0, 0, 0, 0
     errors = []
+    n_chunks_by_doc = dict()
     with output_path.open("w") as f:
         for doc_i, doc in enumerate(docs):
             # to avoid heartbeat timeout
@@ -185,56 +188,79 @@ async def embed_documents_act(
                 for chunk in chunks:
                     f.write(f"{chunk.model_dump_json()}\n")
                 n_chunks += len(chunks)
+                n_chunks_by_doc[doc.doc_id] = len(chunks)
                 n_successes += 1
                 n_successes_pages += doc.n_pages
             if progress is not None:
                 await progress(doc_i + 1)
     processed = ProcessingReportWithPages(n_docs=n_docs, n_pages=n_pages)
     successes = ProcessingReportWithPages(n_docs=n_successes, n_pages=n_successes_pages)
     errors = ErrorReportWithPages.from_errors(*errors)
-    return EmbeddingResponse(
+    response = EmbeddingResponse(
         processed=processed, successes=successes, errors=errors, n_chunks=n_chunks
     )
+    return response, n_chunks_by_doc
 
 
 async def index_embeddings_act(
     chunks_path: Path,
+    n_chunks_by_doc: dict[DocId, int],
     *,
     es_client: ESClient,
     index: str,
     dim: int,
     bulk_size: int,
 ) -> int:
     await ensure_embeddings_index(es_client, index, dim=dim)
-    chunks = [c async for c in async_read_jsonl_as(chunks_path, EmbeddedChunk)]
-    if not chunks:
-        return 0
-    # chunking might have changed, we don't want stale chunks to remain,
-    # remove previous chunks
-    doc_ids = sorted({c.doc_id for c in chunks})
-    delete_query = {QUERY: {"terms": {_DOC_ID: doc_ids}}}
-    await es_client.delete_by_query(
-        index=index, body=delete_query, conflicts="proceed", refresh=True
+    actions = (
+        _index_action(c, index)
+        async for c in async_read_jsonl_as(chunks_path, EmbeddedChunk)
+    )
+    n_indexed, _ = await async_bulk(
+        es_client,
+        actions,
+        chunk_size=bulk_size,
+        stats_only=True,
+        raise_on_error=True,
+        refresh="wait_for",
+    )
+    await _delete_stale_chunks(es_client, index, n_chunks_by_doc)
+    logger.debug(
+        "indexed %s chunks for %s docs in %s", n_indexed, len(n_chunks_by_doc), index
     )
-    n_indexed = 0
-    for start in range(0, len(chunks), bulk_size):
-        bulk = chunks[start : start + bulk_size]
-        actions = (_index_action(c, index) for c in bulk)
-        n_success, _ = await async_bulk(
-            es_client, actions, raise_on_error=True, refresh="wait_for"
-        )
-        n_indexed += n_success
-    logger.debug("indexed %s chunks for %s docs in %s", n_indexed, len(doc_ids), index)
     return n_indexed
 
 
+async def _delete_stale_chunks(
+    es_client: ESClient, index: str, n_chunks_by_doc: dict[DocId, int]
+) -> None:
+    docs = sorted(n_chunks_by_doc.items())
+    for start in range(0, len(docs), _DELETE_GROUP_SIZE):
+        group = docs[start : start + _DELETE_GROUP_SIZE]
+        stale = [
+            {
+                "bool": {
+                    "filter": [
+                        {"term": {_DOC_ID: doc_id}},
+                        {"range": {_CHUNK_INDEX: {"gte": n_chunks}}},
+                    ]
+                }
+            }
+            for doc_id, n_chunks in group
+        ]
+        query = {"bool": {"should": stale, "minimum_should_match": 1}}
+        await es_client.delete_by_query(
+            index=index, body={QUERY: query}, conflicts="proceed", refresh=True
+        )
+
+
 def embeddings_index_body(dim: int) -> dict[str, Any]:
     keyword = {"type": "keyword"}
     properties = {
         _DOC_ID: keyword,
         "rootDocument": keyword,
         "project": keyword,
-        "chunkIndex": {"type": "integer"},
+        _CHUNK_INDEX: {"type": "integer"},
         "page": {"type": "integer"},
         "text": {"type": "text"},
         "headings": {"type": "text"},
@@ -284,26 +310,28 @@ def _embed_doc(
     artifacts_root: Path,
 ) -> list[EmbeddedChunk]:
     page_chunks = [
-        (page, chunker.contextualize(chunk), chunk.meta.headings)
+        (page, chunk)
         for page, md in read_structure_pages(doc, artifacts_root)
         for chunk in _chunk_markdown(md, chunker, md_converter)
     ]
     if not page_chunks:
         return []
-    embeddings = embedder.embed([text for _, text, _ in page_chunks], dim)
+    # embed chunks with their headings for context, but only store the chunk text
+    texts = [chunker.contextualize(chunk) for _, chunk in page_chunks]
+    embeddings = embedder.embed(texts, dim)
     return [
         EmbeddedChunk(
             doc_id=doc.doc_id,
             root_document=doc.root_document,
             project=doc.project,
             chunk_index=chunk_i,
             page=page,
-            text=text,
-            headings=headings,
+            text=chunk.text,
+            headings=chunk.meta.headings,
             model=embedder.model,
             embedding=embedding,
         )
-        for chunk_i, ((page, text, headings), embedding) in enumerate(
+        for chunk_i, ((page, chunk), embedding) in enumerate(
             zip(page_chunks, embeddings, strict=True)
         )
     ]

workers/extract-worker/extract_worker/objects.py

@@ -122,6 +122,6 @@ def __add__(self, other: "EmbeddingResponse") -> "EmbeddingResponse":
 
 
 class EmbedBatchResult(DatashareModel):
-    # path to the embedded chunks JSONL relative to workdir
     chunks: Path
+    n_chunks_by_doc: dict[DocId, int]
     response: EmbeddingResponse

workers/extract-worker/extract_worker/workflows.py

@@ -111,7 +111,7 @@ async def _embed_and_index(
         )
         n_indexed = await execute_activity(
             DocumentEmbedding.index_embeddings,
-            args=[res.chunks, project, args.dim],
+            args=[res.chunks, res.n_chunks_by_doc, project, args.dim],
             task_queue=TaskQueues.IO,
             timeouts=worker_config.timeouts.index_embeddings,
         )

workers/extract-worker/tests/test_embeddings.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from datashare_python.conftest import TEST_PROJECT
@@ -11,6 +12,7 @@
 )
 from datashare_python.utils import read_jsonl_as
 from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
+from extract_worker import embeddings
 from extract_worker.activities import create_embedding_batches_act
 from extract_worker.config import ExtractWorkerConfig
 from extract_worker.embeddings import (
@@ -82,7 +84,7 @@ async def test_embed_documents_act(
     output_path = tmp_path / "chunks.jsonl"
 
     # When
-    response = await embed_documents_act(
+    response, n_chunks_by_doc = await embed_documents_act(
         batch,
         chunker=HierarchicalChunker(),
         md_converter=make_md_converter(),
@@ -98,13 +100,15 @@ async def test_embed_documents_act(
     assert response.errors.n_docs == 1
     assert response.errors.errors[0].doc_id == missing.doc_id
     assert response.n_chunks == 2
+    assert n_chunks_by_doc == {doc_with_structure.doc_id: 2}
     chunks = list(read_jsonl_as(output_path, EmbeddedChunk))
     assert [c.chunk_index for c in chunks] == [0, 1]
     assert [c.page for c in chunks] == [0, 1]
     assert all(c.doc_id == doc_with_structure.doc_id for c in chunks)
     assert all(c.root_document == "root-0" for c in chunks)
     assert chunks[0].headings == ["Offshore leaks"]
-    assert "Shell companies" in chunks[0].text
+    # stored text shouldn't contain the headings used as embedding context
+    assert chunks[0].text == "Shell companies registered in the islands."
     assert chunks[1].headings == ["Beneficial owners"]
     assert all(len(c.embedding) == _DIM for c in chunks)
     assert all(c.model == "fake-model" for c in chunks)
@@ -193,11 +197,20 @@ async def test_index_embeddings_act(test_es_client: ESClient, tmp_path: Path) ->
     kwargs = {"es_client": test_es_client, "index": index, "dim": _DIM}
 
     # When
-    await index_embeddings_act(stale, bulk_size=2, **kwargs)
-    n_indexed = await index_embeddings_act(chunks_path, bulk_size=2, **kwargs)
+    await index_embeddings_act(stale, {"doc-0": 3}, bulk_size=2, **kwargs)
+    with patch(
+        "extract_worker.embeddings.async_bulk", wraps=embeddings.async_bulk
+    ) as bulk:
+        n_indexed = await index_embeddings_act(
+            chunks_path, {"doc-0": 2, "doc-1": 1}, bulk_size=2, **kwargs
+        )
 
     # Then
     assert n_indexed == 3
+    bulk.assert_awaited_once()
+    actions = bulk.await_args.args[1]
+    assert hasattr(actions, "__aiter__")
+    assert bulk.await_args.kwargs["chunk_size"] == 2
     res = await test_es_client.search(
         index=index, body={"query": {"match_all": {}}, "sort": ["_id"]}
     )
@@ -220,7 +233,12 @@ async def test_index_embeddings_act_should_raise_for_missing_index(
     # When/Then
     with pytest.raises(MissingEmbeddingsIndex):
         await index_embeddings_act(
-            chunks_path, es_client=test_es_client, index=index, dim=_DIM, bulk_size=2
+            chunks_path,
+            {"doc-0": 1},
+            es_client=test_es_client,
+            index=index,
+            dim=_DIM,
+            bulk_size=2,
         )
     assert not await test_es_client.indices.exists(index=index)
 
@@ -263,3 +281,64 @@ async def test_ensure_embeddings_index_should_raise_for_missing_vector_field(
     with pytest.raises(EmbeddingSizeMismatch) as exc:
         await ensure_embeddings_index(test_es_client, index, dim=_DIM)
     assert exc.value.actual is None
+
+
+async def _indexed_ids(es_client: ESClient, index: str) -> list[str]:
+    body = {"query": {"match_all": {}}, "sort": ["_id"]}
+    res = await es_client.search(index=index, body=body)
+    return [h["_id"] for h in res[HITS][HITS]]
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("delete_group_size", [1, 256])
+async def test_index_embeddings_act_should_only_replace_embedded_docs_chunks(
+    test_es_client: ESClient, tmp_path: Path, delete_group_size: int
+) -> None:
+    # Given
+    index = f"{TEST_PROJECT}-embeddings"
+    await test_es_client.indices.create(index=index, body=embeddings_index_body(_DIM))
+    kwargs = {"es_client": test_es_client, "index": index, "dim": _DIM}
+    previous = tmp_path / "previous.jsonl"
+    previous_chunks = [_chunk("doc-0", 0), _chunk("doc-1", 0), _chunk("doc-2", 0)]
+    previous.write_text("\n".join(c.model_dump_json() for c in previous_chunks))
+    await index_embeddings_act(
+        previous, {"doc-0": 1, "doc-1": 1, "doc-2": 1}, bulk_size=10, **kwargs
+    )
+    chunks_path = tmp_path / "chunks.jsonl"
+    chunks_path.write_text(_chunk("doc-0", 0).model_dump_json())
+
+    # When
+    with patch("extract_worker.embeddings._DELETE_GROUP_SIZE", delete_group_size):
+        await index_embeddings_act(
+            chunks_path, {"doc-0": 1, "doc-1": 0}, bulk_size=10, **kwargs
+        )
+
+    # Then
+    assert await _indexed_ids(test_es_client, index) == ["doc-0-0", "doc-2-0"]
+
+
+@pytest.mark.integration
+async def test_index_embeddings_act_should_keep_previous_chunks_on_write_failure(
+    test_es_client: ESClient, tmp_path: Path
+) -> None:
+    # Given
+    index = f"{TEST_PROJECT}-embeddings"
+    await test_es_client.indices.create(index=index, body=embeddings_index_body(_DIM))
+    kwargs = {"es_client": test_es_client, "index": index, "dim": _DIM}
+    previous = tmp_path / "previous.jsonl"
+    previous.write_text(
+        "\n".join(_chunk("doc-0", i).model_dump_json() for i in range(2))
+    )
+    await index_embeddings_act(previous, {"doc-0": 2}, bulk_size=10, **kwargs)
+    chunks_path = tmp_path / "chunks.jsonl"
+    chunks_path.write_text(_chunk("doc-0", 0).model_dump_json())
+
+    # When
+    with (
+        patch("extract_worker.embeddings.async_bulk", side_effect=RuntimeError("boom")),
+        pytest.raises(RuntimeError),
+    ):
+        await index_embeddings_act(chunks_path, {"doc-0": 1}, bulk_size=10, **kwargs)
+
+    # Then
+    assert await _indexed_ids(test_es_client, index) == ["doc-0-0", "doc-0-1"]