From 295fd4c00aa7ddb2255af9156b866b699e45a456 Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Wed, 28 Oct 2026 10:20:11 +0100 Subject: [PATCH 1/2] refactor(extract-worker): fall back to doc content for embeddings --- .../extract_worker/activities.py | 13 +- .../extract_worker/embeddings.py | 111 +++++++++++++++--- workers/extract-worker/tests/__init__.py | 2 + workers/extract-worker/tests/conftest.py | 20 +++- .../extract-worker/tests/test_embeddings.py | 42 ++++++- .../extract-worker/tests/test_workflows.py | 12 +- 6 files changed, 168 insertions(+), 32 deletions(-) diff --git a/workers/extract-worker/extract_worker/activities.py b/workers/extract-worker/extract_worker/activities.py index 19120dba..48ba7dcf 100644 --- a/workers/extract-worker/extract_worker/activities.py +++ b/workers/extract-worker/extract_worker/activities.py @@ -74,7 +74,7 @@ index_embeddings_act, make_chunker, make_md_converter, - with_structure_artifact, + with_embeddable_content, ) from .objects import ( DocId, @@ -347,11 +347,14 @@ async def create_embedding_batches_act( es_client: ESClient, ) -> AsyncIterable[Path]: query = build_embedding_doc_query(docs) - docs = with_structure_artifact( - _search_docs(es_client, project, query), artifacts_root + docs = _search_docs(es_client, project, query) + batches = with_embeddable_content( + async_batches(docs, target_n_docs_per_task), + es_client=es_client, + project=project, + artifacts_root=artifacts_root, ) - docs = (DatashareFile.from_parent(d) async for d in docs) - batches = async_batches(docs, target_n_docs_per_task) + batches = ([DatashareFile.from_parent(d) for d in b] async for b in batches) async for p in write_batches(batches, output_dir, prefix="embed_"): yield p diff --git a/workers/extract-worker/extract_worker/embeddings.py b/workers/extract-worker/extract_worker/embeddings.py index 49eeb38a..435ce60d 100644 --- a/workers/extract-worker/extract_worker/embeddings.py +++ b/workers/extract-worker/extract_worker/embeddings.py @@ -2,6 +2,7 @@ import contextlib import json import logging +import re from collections.abc import AsyncIterable, Iterator from io import BytesIO from pathlib import Path @@ -11,6 +12,7 @@ from datashare_python.objects import ( PROCESSED_FILE_TA, ByteRangesPagination, + DatashareFile, Document, ErrorReportWithPages, ProcessedFile, @@ -25,9 +27,12 @@ ) from elasticsearch._async.helpers import async_bulk from icij_common.es import ( + DOC_CONTENT, ES_DOCUMENT_TYPE, + HITS, ID_, QUERY, + SOURCE, ESClient, and_query, has_id, @@ -52,6 +57,7 @@ _EMBEDDING = "embedding" _CHUNK_INDEX = "chunkIndex" _DELETE_GROUP_SIZE = 256 +_PARAGRAPH_SEP = re.compile(r"\n\s*\n") class MissingEmbeddingsIndex(ValueError): @@ -309,11 +315,7 @@ def _embed_doc( dim: int, artifacts_root: Path, ) -> list[EmbeddedChunk]: - page_chunks = [ - (page, chunk) - for page, md in read_structure_pages(doc, artifacts_root) - for chunk in _chunk_markdown(md, chunker, md_converter) - ] + page_chunks = list(_page_chunks(doc, chunker, md_converter, artifacts_root)) if not page_chunks: return [] # embed chunks with their headings for context, but only store the chunk text @@ -337,6 +339,42 @@ def _embed_doc( ] +def _page_chunks( + doc: ProcessedFile, + chunker: "BaseChunker", # noqa: F821 + md_converter: "DocumentConverter", # noqa: F821 + artifacts_root: Path, +) -> Iterator[tuple[int, "BaseChunk"]]: # noqa: F821 + content = _fallback_content(doc) + if content is not None: + for chunk in _chunk_text(content, chunker): + yield 0, chunk + return + for page, md in read_structure_pages(doc, artifacts_root): + for chunk in _chunk_markdown(md, chunker, md_converter): + yield page, chunk + + +def _fallback_content(doc: ProcessedFile) -> str | None: + if isinstance(doc, DatashareFile): + return doc.parent.content + return None + + +def _chunk_text( + text: str, + chunker: "BaseChunker", # noqa: F821 +) -> Iterator["BaseChunk"]: # noqa: F821 + from docling_core.types.doc import DocItemLabel, DoclingDocument # noqa: PLC0415 + + dl_doc = DoclingDocument(name="content") + paragraphs = (p.strip() for p in _PARAGRAPH_SEP.split(text)) + for paragraph in paragraphs: + if paragraph: + dl_doc.add_text(label=DocItemLabel.TEXT, text=paragraph) + yield from chunker.chunk(dl_doc) + + def _chunk_markdown( md: bytes, chunker: "BaseChunker", # noqa: F821 @@ -418,14 +456,55 @@ def build_embedding_doc_query( raise ValueError(f"unsupported format {docs.__class__.__name__}") -async def with_structure_artifact( - docs: AsyncIterable[Document], artifacts_root: Path -) -> AsyncIterable[Document]: - async for d in docs: - structure_dir = artifact_path( - d.id, StructureArtifact, project=d.project, root=artifacts_root - ) - if _structure_md_path(structure_dir) is None: - logger.debug("skipping %s (no markdown structure)", d.id) - continue - yield d +async def with_embeddable_content( + batches: AsyncIterable[list[Document]], + *, + es_client: ESClient, + project: str, + artifacts_root: Path, +) -> AsyncIterable[list[Document]]: + """Fall back to doc context when no markdown present (skip if neither available). + + :param batches: batches of documents to filter + :param es_client: ES client + :param project: document project + :param artifacts_root: root of the artifacts directory + :return: batches of embeddable documents + """ + async for batch in batches: + missing = [d.id for d in batch if not _has_structure(d, artifacts_root)] + contents = await _fetch_contents(es_client, project, missing) + embeddable = [] + for d in batch: + if d.id not in contents: + embeddable.append(d) + continue + content = contents[d.id] + if content is None or not content.strip(): + logger.debug("skipping %s (no markdown/content)", d.id) + continue + embeddable.append(d.model_copy(update={"content": content})) + if embeddable: + yield embeddable + + +def _has_structure(doc: Document, artifacts_root: Path) -> bool: + structure_dir = artifact_path( + doc.id, StructureArtifact, project=doc.project, root=artifacts_root + ) + return _structure_md_path(structure_dir) is not None + + +async def _fetch_contents( + es_client: ESClient, project: str, doc_ids: list[DocId] +) -> dict[DocId, str | None]: + contents = dict.fromkeys(doc_ids) + if not doc_ids: + return contents + body = {QUERY: has_id(doc_ids), "size": len(doc_ids)} + res = await es_client.search( + index=project, body=body, _source_includes=[DOC_CONTENT] + ) + for hit in res[HITS][HITS]: + contents[hit[ID_]] = hit[SOURCE].get(DOC_CONTENT) + return contents diff --git a/workers/extract-worker/tests/__init__.py b/workers/extract-worker/tests/__init__.py index bbca3271..cd93456e 100644 --- a/workers/extract-worker/tests/__init__.py +++ b/workers/extract-worker/tests/__init__.py @@ -20,6 +20,8 @@ "## Beneficial owners\n\nThe owners were hidden behind nominees.\n", ] +DOC_3_CONTENT = "The vote is on Sunday.\n\n \n\nPolls close at eight.\n" + def write_structure(doc: ProcessedFile, artifacts_root: Path, pages: list[str]) -> None: contents = [(p + PAGE_SEP).encode() for p in pages[:-1]] + [pages[-1].encode()] diff --git a/workers/extract-worker/tests/conftest.py b/workers/extract-worker/tests/conftest.py index 95b0d0a8..dea36952 100644 --- a/workers/extract-worker/tests/conftest.py +++ b/workers/extract-worker/tests/conftest.py @@ -17,7 +17,6 @@ TEST_PROJECT, clear_dirs, dev_worker_context, - doc_3, index_docs, pytest_collection_modifyitems, test_es_client, @@ -52,7 +51,7 @@ from extract_worker.constants import TaskQueue from icij_common.es import ESClient -from tests import DOCS_PATH, PAGES, write_structure +from tests import DOC_3_CONTENT, DOCS_PATH, PAGES, write_structure mimetypes.init() @@ -139,13 +138,28 @@ def doc_2() -> Document: ) +@pytest.fixture(scope="session") +def doc_3() -> Document: + return Document( + id="doc-3", + root_document="root-3", + extraction_level=1, + index=TEST_PROJECT, + language=DatashareLanguage("ENGLISH"), + content=DOC_3_CONTENT, + content_type="message/rfc822", + path=Path("doc-3.eml"), + metadata={"tika_metadata_resourcename": "doc-3.eml"}, + ) + + @pytest.fixture async def populate_es( test_es_client: ESClient, # noqa: F811 doc_0: Document, doc_1: Document, doc_2: Document, - doc_3: Document, # noqa: F811 + doc_3: Document, ) -> list[Document]: docs = [doc_0, doc_1, doc_2, doc_3] async for _ in index_docs(test_es_client, docs=docs, index_name=TEST_PROJECT): diff --git a/workers/extract-worker/tests/test_embeddings.py b/workers/extract-worker/tests/test_embeddings.py index 5feb7070..0e5b6c3e 100644 --- a/workers/extract-worker/tests/test_embeddings.py +++ b/workers/extract-worker/tests/test_embeddings.py @@ -33,7 +33,7 @@ ) from icij_common.es import HITS, ESClient -from tests import PAGE_SEP, PAGES +from tests import DOC_3_CONTENT, PAGE_SEP, PAGES _DIM = 4 @@ -114,6 +114,40 @@ async def test_embed_documents_act( assert all(c.model == "fake-model" for c in chunks) +async def test_embed_documents_act_should_fallback_to_content( + doc_3: Document, test_worker_config: ExtractWorkerConfig, tmp_path: Path +) -> None: + # Given + doc = DatashareFile.from_parent(doc_3) + batch = tmp_path / "batch.jsonl" + batch.write_text(PROCESSED_FILE_TA.dump_json(doc).decode() + "\n") + output_path = tmp_path / "chunks.jsonl" + + # When + response, n_chunks_by_doc = await embed_documents_act( + batch, + chunker=HierarchicalChunker(), + md_converter=make_md_converter(), + embedder=_FakeEmbedder(), + dim=_DIM, + artifacts_root=test_worker_config.paths.artifacts, + output_path=output_path, + ) + + # Then + assert response.successes == ProcessingReportWithPages(n_docs=1, n_pages=1) + assert response.errors.n_docs == 0 + assert n_chunks_by_doc == {doc.doc_id: 2} + chunks = list(read_jsonl_as(output_path, EmbeddedChunk)) + assert [c.text for c in chunks] == [ + "The vote is on Sunday.", + "Polls close at eight.", + ] + assert [c.page for c in chunks] == [0, 0] + assert all(c.headings is None for c in chunks) + assert all(c.root_document == "root-3" for c in chunks) + + @pytest.mark.parametrize( ("docs", "expected"), [ @@ -148,6 +182,7 @@ def test_build_embedding_doc_query( async def test_create_embedding_batches_act( populate_es: list[Document], # noqa: ARG001 doc_with_structure: ProcessedFile, + doc_3: Document, test_es_client: ESClient, test_worker_config: ExtractWorkerConfig, tmp_path: Path, @@ -168,7 +203,10 @@ async def test_create_embedding_batches_act( # Then assert len(batches) == 1 docs = list(read_jsonl_as(batches[0], PROCESSED_FILE_TA)) - assert [d.doc_id for d in docs] == [doc_with_structure.doc_id] + by_id = {d.doc_id: d for d in docs} + assert sorted(by_id) == [doc_with_structure.doc_id, doc_3.id] + assert by_id[doc_with_structure.doc_id].parent.content is None + assert by_id[doc_3.id].parent.content == DOC_3_CONTENT def _chunk(doc_id: str, chunk_index: int) -> EmbeddedChunk: diff --git a/workers/extract-worker/tests/test_workflows.py b/workers/extract-worker/tests/test_workflows.py index ef77fe07..39af15f1 100644 --- a/workers/extract-worker/tests/test_workflows.py +++ b/workers/extract-worker/tests/test_workflows.py @@ -71,6 +71,7 @@ async def test_embed_workflow_e2e( # noqa: PLR0917 test_es_client: ESClient, test_worker_config: ExtractWorkerConfig, doc_with_structure: ProcessedFile, + doc_3: Document, ) -> None: # Given client = test_temporal_client @@ -89,15 +90,14 @@ async def test_embed_workflow_e2e( # noqa: PLR0917 ) # Then - assert response.successes == ProcessingReportWithPages(n_docs=1, n_pages=2) + assert response.successes == ProcessingReportWithPages(n_docs=2, n_pages=3) assert response.errors.n_docs == 0 - assert response.n_chunks == 2 + assert response.n_chunks == 3 res = await test_es_client.search(index=index, body={"query": {"match_all": {}}}) hits = res["hits"]["hits"] - assert len(hits) == 2 - for h in hits: - assert h["_source"]["docId"] == doc_with_structure.doc_id - assert len(h["_source"]["embedding"]) == 64 + doc_ids = sorted(h["_source"]["docId"] for h in hits) + assert doc_ids == sorted([doc_with_structure.doc_id] * 2 + [doc_3.id]) + assert all(len(h["_source"]["embedding"]) == 64 for h in hits) @pytest.mark.e2e From 3090fbefede24fac3a7bc61a0a800bd9aaeaba5f Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Wed, 28 Oct 2026 10:24:38 +0100 Subject: [PATCH 2/2] qa --- workers/extract-worker/tests/test_embeddings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/extract-worker/tests/test_embeddings.py b/workers/extract-worker/tests/test_embeddings.py index 0e5b6c3e..fac90afb 100644 --- a/workers/extract-worker/tests/test_embeddings.py +++ b/workers/extract-worker/tests/test_embeddings.py @@ -179,7 +179,7 @@ def test_build_embedding_doc_query( @pytest.mark.integration -async def test_create_embedding_batches_act( +async def test_create_embedding_batches_act( # noqa: PLR0917 populate_es: list[Document], # noqa: ARG001 doc_with_structure: ProcessedFile, doc_3: Document,