danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add docling_pages column to DocumentRecord/Document for separate page image storage.

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-07 16:22:45 +03:00
parent b614ef19e1
commit 53445b8722
No known key found for this signature in database
6 changed files with 195 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
haiku_rag_slim/haiku/rag/client.py
View file

@ -17,7 +17,6 @@ import httpx
from haiku.rag.config import AppConfig, Config
from haiku.rag.converters import get_converter
from haiku.rag.reranking import get_reranker
from haiku.rag.store.compression import compress_json
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document
@ -495,9 +494,8 @@ class HaikuRAG:
uri=uri,
title=title,
metadata=metadata or {},
docling_document=compress_json(docling_document.model_dump_json()),
docling_version=docling_document.version,
)
document.set_docling(docling_document)
# Store document and chunks
return await self._store_document_with_chunks(document, embedded_chunks)
@ -535,9 +533,8 @@ class HaikuRAG:
uri=uri,
title=title,
metadata=metadata or {},
docling_document=compress_json(docling_document.model_dump_json()),
docling_version=docling_document.version,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, chunks)
@ -672,10 +669,7 @@ class HaikuRAG:
# Update existing document and rechunk
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
)
existing_doc.docling_version = docling_document.version
existing_doc.set_docling(docling_document)
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
@ -694,9 +688,8 @@ class HaikuRAG:
uri=uri,
title=title,
metadata=metadata,
docling_document=compress_json(docling_document.model_dump_json()),
docling_version=docling_document.version,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
async def _create_or_update_document_from_url(
@ -789,10 +782,7 @@ class HaikuRAG:
# Update existing document and rechunk
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
)
existing_doc.docling_version = docling_document.version
existing_doc.set_docling(docling_document)
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
@ -811,9 +801,8 @@ class HaikuRAG:
uri=url,
title=title,
metadata=metadata,
docling_document=compress_json(docling_document.model_dump_json()),
docling_version=docling_document.version,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
def _get_extension_from_content_type_or_url(
@ -963,10 +952,7 @@ class HaikuRAG:
# Store docling data if provided
if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
)
existing_doc.docling_version = docling_document.version
existing_doc.set_docling(docling_document)
elif content is not None:
existing_doc.content = content
@ -975,10 +961,7 @@ class HaikuRAG:
# DoclingDocument provided without chunks - chunk and embed using primitives
if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
)
existing_doc.docling_version = docling_document.version
existing_doc.set_docling(docling_document)
new_chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(new_chunks, self._config)
@ -990,10 +973,7 @@ class HaikuRAG:
assert content is not None
existing_doc.content = content
converted_docling = await self.convert(existing_doc.content)
existing_doc.docling_document = compress_json(
converted_docling.model_dump_json()
)
existing_doc.docling_version = converted_docling.version
existing_doc.set_docling(converted_docling)
new_chunks = await self.chunk(converted_docling)
embedded_chunks = await embed_chunks(new_chunks, self._config)
@ -1578,16 +1558,15 @@ class HaikuRAG:
from PIL import ImageDraw
# Get the document
# Get the document structure (from cache if available)
if not chunk.document_id:
return []
doc = await self.document_repository.get_by_id(chunk.document_id)
doc = await self.document_repository.get_docling_data(chunk.document_id)
if not doc:
return []
# Get DoclingDocument with page images for rendering
docling_doc = doc.get_docling_document(include_pages=True)
docling_doc = doc.get_docling_document()
if not docling_doc:
return []
@ -1604,13 +1583,19 @@ class HaikuRAG:
boxes_by_page[bbox.page_no] = []
boxes_by_page[bbox.page_no].append(bbox)
# Load only the needed page images
pages_doc = await self.document_repository.get_pages_data(chunk.document_id)
if not pages_doc:
return []
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
# Render each page with its bounding boxes
images = []
for page_no in sorted(boxes_by_page.keys()):
if page_no not in docling_doc.pages:
if page_no not in page_images:
continue
page = docling_doc.pages[page_no]
page = page_images[page_no]
if page.image is None or page.image.pil_image is None:
continue
@ -1800,6 +1785,7 @@ class HaikuRAG:
title=doc.title,
metadata=json.dumps(doc.metadata),
docling_document=doc.docling_document,
docling_pages=doc.docling_pages,
docling_version=doc.docling_version,
created_at=doc.created_at.isoformat() if doc.created_at else now,
updated_at=now,
@ -1836,8 +1822,7 @@ class HaikuRAG:
embedded_chunks = await embed_chunks(chunks, self._config)
# Update document fields
doc.docling_document = compress_json(docling_document.model_dump_json())
doc.docling_version = docling_document.version
doc.set_docling(docling_document)
# Prepare chunks with document_id and order
for order, chunk in enumerate(embedded_chunks):
@ -1915,8 +1900,7 @@ class HaikuRAG:
chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, self._config)
doc.docling_document = compress_json(docling_document.model_dump_json())
doc.docling_version = docling_document.version
doc.set_docling(docling_document)
# Prepare chunks with document_id and order
for order, chunk in enumerate(embedded_chunks):

6
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -26,6 +26,7 @@ class DocumentRecord(LanceModel):
title: str | None = None
metadata: str = Field(default="{}")
docling_document: bytes | None = None
docling_pages: bytes | None = None
docling_version: str | None = None
created_at: str = Field(default_factory=lambda: "")
updated_at: str = Field(default_factory=lambda: "")
@ -43,10 +44,11 @@ def get_documents_arrow_schema() -> pa.Schema:
which has 64-bit offsets and no practical size limit.
"""
base_schema = DocumentRecord.to_arrow_schema()
large_binary_columns = {"docling_document", "docling_pages"}
fields = []
for field in base_schema:
if field.name == "docling_document":
fields.append(pa.field("docling_document", pa.large_binary()))
if field.name in large_binary_columns:
fields.append(pa.field(field.name, pa.large_binary()))
else:
fields.append(field)
return pa.schema(fields)

71
haiku_rag_slim/haiku/rag/store/models/document.py
View file

@ -5,10 +5,10 @@ from typing import TYPE_CHECKING
from cachetools import LRUCache
from pydantic import BaseModel, Field
from haiku.rag.store.compression import decompress_json
from haiku.rag.store.compression import compress_docling_split, decompress_json
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.document import DoclingDocument, PageItem
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
@ -30,8 +30,7 @@ def _get_cached_docling_document(
"""Get or parse DoclingDocument with LRU caching by document ID.
Strips page images before validation for performance — cached documents
do not contain page data. Use _parse_full_docling_document for
operations that need page images (e.g. visualize_chunk).
do not contain page data.
"""
if document_id in _docling_document_cache:
return _docling_document_cache[document_id]
@ -41,14 +40,6 @@ def _get_cached_docling_document(
return doc
def _parse_full_docling_document(compressed_data: bytes) -> "DoclingDocument":
"""Parse DoclingDocument with full page data (no caching, no stripping)."""
from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data)
return DoclingDocument.model_validate_json(json_str)
def invalidate_docling_document_cache(document_id: str) -> None:
"""Remove a document from the DoclingDocument cache."""
_docling_document_cache.pop(document_id, None)
@ -65,34 +56,62 @@ class Document(BaseModel):
title: str | None = None
metadata: dict = {}
docling_document: bytes | None = Field(default=None, exclude=True)
docling_pages: bytes | None = Field(default=None, exclude=True)
docling_version: str | None = Field(default=None, exclude=True)
created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now)
def get_docling_document(
self, *, include_pages: bool = False
) -> "DoclingDocument | None":
"""Parse and return the stored DoclingDocument.
def set_docling(self, docling_doc: "DoclingDocument") -> None:
"""Serialize and store a DoclingDocument, splitting structure and pages.
Sets docling_document (zstd-compressed structure without pages),
docling_pages (zstd-compressed page images), and docling_version.
"""
structure, pages = compress_docling_split(docling_doc.model_dump_json())
self.docling_document = structure
self.docling_pages = pages
self.docling_version = docling_doc.version
def get_docling_document(self) -> "DoclingDocument | None":
"""Parse and return the stored DoclingDocument (without page images).
By default, strips page images before parsing for performance.
Uses LRU cache (keyed by document ID) to avoid repeated parsing.
Args:
include_pages: If True, parse with full page data (slower,
bypasses cache). Only needed for operations that access
page images (e.g. visualize_chunk).
Returns:
The parsed DoclingDocument, or None if not stored or no ID.
The parsed DoclingDocument, or None if not stored.
"""
if self.docling_document is None:
return None
if include_pages:
return _parse_full_docling_document(self.docling_document)
# No caching for documents without ID
if self.id is None:
return _validate_without_pages(self.docling_document)
return _get_cached_docling_document(self.id, self.docling_document)
def get_page_images(self, page_numbers: list[int]) -> "dict[int, PageItem]":
"""Decompress and return page images for the requested page numbers.
Loads only the docling_pages blob — does not need the structure.
Validates only the requested pages through Pydantic (for pil_image property).
Args:
page_numbers: Page numbers to retrieve.
Returns:
Dict mapping page number to validated PageItem.
"""
if self.docling_pages is None:
return {}
from docling_core.types.doc.document import PageItem
pages_json = decompress_json(self.docling_pages)
all_pages = json.loads(pages_json)
result: dict[int, PageItem] = {}
for page_no in page_numbers:
page_data = all_pages.get(str(page_no))
if page_data is not None:
result[page_no] = PageItem.model_validate(page_data)
return result

24
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -36,6 +36,7 @@ class DocumentRepository:
title=record.title,
metadata=json.loads(record.metadata),
docling_document=record.docling_document,
docling_pages=record.docling_pages,
docling_version=record.docling_version,
created_at=datetime.fromisoformat(record.created_at)
if record.created_at
@ -62,6 +63,7 @@ class DocumentRepository:
title=entity.title,
metadata=json.dumps(entity.metadata),
docling_document=entity.docling_document,
docling_pages=entity.docling_pages,
docling_version=entity.docling_version,
created_at=now,
updated_at=now,
@ -114,6 +116,27 @@ class DocumentRepository:
docling_version=row.get("docling_version"),
)
async def get_pages_data(self, entity_id: str) -> Document | None:
"""Get a document with only page image data loaded."""
safe_id = _escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.select(["id", "docling_pages"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not results:
return None
row = results[0]
return Document(
id=row["id"],
content="",
docling_pages=row.get("docling_pages"),
)
async def update(self, entity: Document) -> Document:
"""Update an existing document."""
self.store._assert_writable()
@ -138,6 +161,7 @@ class DocumentRepository:
"title": entity.title,
"metadata": json.dumps(entity.metadata),
"docling_document": entity.docling_document,
"docling_pages": entity.docling_pages,
"docling_version": entity.docling_version,
"updated_at": now,
},

20
tests/test_client.py
View file

@ -1,3 +1,4 @@
import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, patch
@ -864,8 +865,11 @@ async def test_client_import_document_stores_docling_data(temp_db_path):
assert doc.id is not None
assert "Content from docling document" in doc.content
assert doc.docling_document is not None
assert decompress_json(doc.docling_document) == docling_doc.model_dump_json()
assert doc.docling_version == docling_doc.version
# Structure is stored without pages
structure = json.loads(decompress_json(doc.docling_document))
assert "pages" not in structure
assert structure["name"] == "test"
@pytest.mark.vcr()
@ -983,11 +987,11 @@ async def test_client_update_document_with_docling_rechunks(temp_db_path):
# Content should be extracted from docling document
assert "Completely different text" in updated_doc.content
assert updated_doc.docling_document is not None
assert (
decompress_json(updated_doc.docling_document)
== docling_doc.model_dump_json()
)
assert updated_doc.docling_version == docling_doc.version
# Structure is stored without pages
structure = json.loads(decompress_json(updated_doc.docling_document))
assert "pages" not in structure
assert structure["name"] == "updated"
# Chunks should be regenerated
new_chunks = await client.chunk_repository.get_by_document_id(doc.id)
@ -1026,10 +1030,8 @@ async def test_client_update_document_docling_with_chunks(temp_db_path):
# Content should be extracted from docling (since content wasn't provided)
assert "Text from docling" in updated_doc.content
assert updated_doc.docling_document is not None
assert (
decompress_json(updated_doc.docling_document)
== docling_doc.model_dump_json()
)
structure = json.loads(decompress_json(updated_doc.docling_document))
assert "pages" not in structure
# Custom chunks should be used (not rechunked from docling)
chunks = await client.chunk_repository.get_by_document_id(doc.id)

88
tests/test_document.py
View file

@ -229,6 +229,94 @@ def test_document_get_docling_document_no_id_no_cache():
assert doc1 is not doc2
def test_set_docling_splits_structure_and_pages():
"""set_docling stores structure and pages separately."""
import json
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.compression import decompress_json
docling_doc = DoclingDocument(name="split_test")
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
document = Document(content="test")
document.set_docling(docling_doc)
assert document.docling_document is not None
assert document.docling_version == docling_doc.version
# Structure should not contain pages
structure = json.loads(decompress_json(document.docling_document))
assert "pages" not in structure
assert structure["name"] == "split_test"
# get_docling_document should work from the split structure
parsed = document.get_docling_document()
assert parsed is not None
assert parsed.name == "split_test"
assert len(list(parsed.iterate_items())) > 0
def test_set_docling_with_page_images():
"""set_docling stores page images in docling_pages."""
import json
from docling_core.types.doc.base import Size
from docling_core.types.doc.document import DoclingDocument, PageItem
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.compression import decompress_json
docling_doc = DoclingDocument(name="pages_test")
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Content")
docling_doc.pages[1] = PageItem(
size=Size(width=612, height=792),
page_no=1,
)
document = Document(content="test")
document.set_docling(docling_doc)
assert document.docling_pages is not None
# Pages blob should contain page data
pages = json.loads(decompress_json(document.docling_pages))
assert "1" in pages
def test_get_page_images():
"""get_page_images returns requested pages from docling_pages blob."""
import json
from haiku.rag.store.compression import compress_json
pages_data = {
"1": {"size": {"width": 612, "height": 792}, "page_no": 1},
"2": {"size": {"width": 612, "height": 792}, "page_no": 2},
"3": {"size": {"width": 612, "height": 792}, "page_no": 3},
}
document = Document(
content="test",
docling_pages=compress_json(json.dumps(pages_data)),
)
result = document.get_page_images([1, 3])
assert len(result) == 2
assert 1 in result
assert 3 in result
assert 2 not in result
# Missing pages are skipped
result = document.get_page_images([99])
assert len(result) == 0
# None docling_pages returns empty
doc_no_pages = Document(content="test")
assert doc_no_pages.get_page_images([1]) == {}
@pytest.mark.asyncio
async def test_get_docling_data_loads_only_docling_columns(
qa_corpus: Dataset, temp_db_path