fix _patch_picture_descriptions silently dropping docling_pages
This commit is contained in:
parent
ce7201271f
commit
e55e3d13c8
2 changed files with 33 additions and 3 deletions
|
|
@ -358,7 +358,17 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
|
||||||
pic.meta = PictureMeta()
|
pic.meta = PictureMeta()
|
||||||
pic.meta.description = DescriptionMetaField(text=text)
|
pic.meta.description = DescriptionMetaField(text=text)
|
||||||
|
|
||||||
doc.set_docling(docling_doc)
|
# Update only docling_document — set_docling would also overwrite
|
||||||
|
# docling_pages by routing through compress_docling_split, which
|
||||||
|
# extracts pages from the in-memory JSON and finds none (the pages
|
||||||
|
# blob is stored separately and is not loaded by get_docling_document).
|
||||||
|
# That would silently destroy page rasters for every doc with at
|
||||||
|
# least one undescribed picture.
|
||||||
|
from haiku.rag.store.compression import compress_docling_split
|
||||||
|
|
||||||
|
structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json())
|
||||||
|
doc.docling_document = structure_bytes
|
||||||
|
doc.docling_version = docling_doc.version
|
||||||
return len(descriptions)
|
return len(descriptions)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -440,8 +440,17 @@ async def test_rebuild_descriptions_requires_enabled(temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkeypatch):
|
async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkeypatch):
|
||||||
"""End-to-end: ingest a doc with a picture (no VLM at ingest), then run
|
"""End-to-end: ingest a doc with a picture (no VLM at ingest), then run
|
||||||
rebuild --descriptions with the VLM mocked. The docling blob should gain
|
rebuild --descriptions with the VLM mocked. After the rebuild:
|
||||||
the description in meta, and the chunk text should pick it up."""
|
|
||||||
|
- the docling blob has the description in meta;
|
||||||
|
- the chunk text picks it up;
|
||||||
|
- the docling_pages blob is preserved untouched.
|
||||||
|
|
||||||
|
The pages assertion guards against a foot-gun in compress_docling_split:
|
||||||
|
the docling document loaded via get_docling_document() never carries
|
||||||
|
pages (they live in a separate column), so calling set_docling() after
|
||||||
|
patching would write pages_bytes=None and silently destroy page rasters
|
||||||
|
on disk — breaking visualize_chunk for the affected docs."""
|
||||||
from haiku.rag.client.documents import _store_document_with_chunks
|
from haiku.rag.client.documents import _store_document_with_chunks
|
||||||
from haiku.rag.config import AppConfig
|
from haiku.rag.config import AppConfig
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
@ -458,6 +467,14 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
|
||||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||||
assert created.id is not None
|
assert created.id is not None
|
||||||
|
|
||||||
|
# _docling_doc_with_picture has no PageItems, so set_docling leaves
|
||||||
|
# docling_pages as None. Inject sentinel bytes to stand in for what
|
||||||
|
# a real ingest with generate_page_images=True would store.
|
||||||
|
sentinel_pages = b"\x80SENTINEL_PAGE_BYTES"
|
||||||
|
await rag.store.documents_table.update(
|
||||||
|
{"docling_pages": sentinel_pages}, where=f"id = '{created.id}'"
|
||||||
|
)
|
||||||
|
|
||||||
from_blob = (
|
from_blob = (
|
||||||
await rag.document_repository.get_by_id(created.id)
|
await rag.document_repository.get_by_id(created.id)
|
||||||
).get_docling_document() # type: ignore[union-attr]
|
).get_docling_document() # type: ignore[union-attr]
|
||||||
|
|
@ -503,6 +520,9 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
|
||||||
chunks = await rag.chunk_repository.get_by_document_id(created.id)
|
chunks = await rag.chunk_repository.get_by_document_id(created.id)
|
||||||
assert any("A red square (mocked)." in (c.content or "") for c in chunks)
|
assert any("A red square (mocked)." in (c.content or "") for c in chunks)
|
||||||
|
|
||||||
|
# docling_pages must survive untouched — see docstring.
|
||||||
|
assert after.docling_pages == sentinel_pages
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkeypatch):
|
async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkeypatch):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue