diff --git a/workers/extract-worker/extract_worker/activities.py b/workers/extract-worker/extract_worker/activities.py index bbf1e05e..19120dba 100644 --- a/workers/extract-worker/extract_worker/activities.py +++ b/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,11 +266,15 @@ 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 @@ -278,6 +282,7 @@ async def index_embeddings(self, chunks: Path, project: str, dim: int) -> int: 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, diff --git a/workers/extract-worker/extract_worker/embeddings.py b/workers/extract-worker/extract_worker/embeddings.py index c4ec8109..49eeb38a 100644 --- a/workers/extract-worker/extract_worker/embeddings.py +++ b/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,6 +188,7 @@ 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: @@ -192,13 +196,15 @@ async def embed_documents_act( 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, @@ -206,35 +212,55 @@ async def index_embeddings_act( 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,13 +310,15 @@ 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, @@ -298,12 +326,12 @@ def _embed_doc( 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) ) ] diff --git a/workers/extract-worker/extract_worker/objects.py b/workers/extract-worker/extract_worker/objects.py index 69f71671..a8d39d0f 100644 --- a/workers/extract-worker/extract_worker/objects.py +++ b/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 diff --git a/workers/extract-worker/extract_worker/workflows.py b/workers/extract-worker/extract_worker/workflows.py index 6a99c70e..771ec565 100644 --- a/workers/extract-worker/extract_worker/workflows.py +++ b/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, ) diff --git a/workers/extract-worker/tests/test_embeddings.py b/workers/extract-worker/tests/test_embeddings.py index 5dd0fc1f..5feb7070 100644 --- a/workers/extract-worker/tests/test_embeddings.py +++ b/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"]