have rebuild --rechunk re-chunk from the stored docling blob instead of re-converting from the markdown export
This commit is contained in:
parent
7831057339
commit
875c635a77
2 changed files with 33 additions and 8 deletions
|
|
@ -223,28 +223,27 @@ async def _flush_rebuild_batch(
|
||||||
async def _rebuild_rechunk(
|
async def _rebuild_rechunk(
|
||||||
client: "HaikuRAG", documents: list[Document]
|
client: "HaikuRAG", documents: list[Document]
|
||||||
) -> AsyncGenerator[str, None]:
|
) -> AsyncGenerator[str, None]:
|
||||||
"""Re-chunk and re-embed from existing document content."""
|
"""Re-chunk and re-embed each document from its stored docling blob."""
|
||||||
from haiku.rag.embeddings import embed_chunks
|
from haiku.rag.embeddings import embed_chunks
|
||||||
|
|
||||||
pending_chunks: list[Chunk] = []
|
pending_chunks: list[Chunk] = []
|
||||||
pending_docs: list[Document] = []
|
pending_docs: list[Document] = []
|
||||||
pending_doc_ids: list[str] = []
|
pending_doc_ids: list[str] = []
|
||||||
|
|
||||||
converter = get_converter(client._config)
|
|
||||||
|
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
assert doc.id is not None
|
assert doc.id is not None
|
||||||
|
|
||||||
# Convert stored markdown to DoclingDocument
|
docling_document = doc.get_docling_document()
|
||||||
docling_document = await converter.convert_text(doc.content, format="md")
|
if docling_document is None:
|
||||||
|
raise ValueError(
|
||||||
|
f"Document {doc.id} has no stored docling document; rechunk "
|
||||||
|
"requires it. Run a full rebuild (without --rechunk) instead."
|
||||||
|
)
|
||||||
|
|
||||||
# Chunk and embed
|
# Chunk and embed
|
||||||
chunks = await client.chunk(docling_document)
|
chunks = await client.chunk(docling_document)
|
||||||
embedded_chunks = await embed_chunks(chunks, client._config)
|
embedded_chunks = await embed_chunks(chunks, client._config)
|
||||||
|
|
||||||
# Update document fields
|
|
||||||
doc.set_docling(docling_document)
|
|
||||||
|
|
||||||
# Prepare chunks with document_id and order
|
# Prepare chunks with document_id and order
|
||||||
for order, chunk in enumerate(embedded_chunks):
|
for order, chunk in enumerate(embedded_chunks):
|
||||||
chunk.document_id = doc.id
|
chunk.document_id = doc.id
|
||||||
|
|
|
||||||
|
|
@ -166,6 +166,32 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
|
||||||
assert "picture" in out.labels
|
assert "picture" in out.labels
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_rechunk_preserves_picture_data(temp_db_path):
|
||||||
|
"""``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
|
||||||
|
from haiku.rag.client import RebuildMode
|
||||||
|
from haiku.rag.client.documents import _store_document_with_chunks
|
||||||
|
from haiku.rag.store.models.document import Document
|
||||||
|
from tests.store.test_document_items import _docling_doc_with_picture
|
||||||
|
|
||||||
|
docling_doc = _docling_doc_with_picture()
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||||
|
rag._config.processing.pictures = "image"
|
||||||
|
document = Document(content="x", uri="test://doc")
|
||||||
|
document.set_docling(docling_doc)
|
||||||
|
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||||
|
assert created.id is not None
|
||||||
|
before = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||||
|
assert before.get("#/pictures/0") is not None
|
||||||
|
|
||||||
|
async for _ in rag.rebuild_database(mode=RebuildMode.RECHUNK):
|
||||||
|
pass
|
||||||
|
|
||||||
|
after = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||||
|
assert after.get("#/pictures/0") == before.get("#/pictures/0")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
|
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
|
||||||
"""Switching to ``pictures="none"`` and re-running update_document
|
"""Switching to ``pictures="none"`` and re-running update_document
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue