refactor(extract-worker): fall back to doc content for embeddings

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
 

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

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()]

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):

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"),
     [
@@ -145,9 +179,10 @@ 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,
     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:

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