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

Merge pull request #338 from ggozad/feat/zstd-image-separation

Browse source 
Separate page images into dedicated column and migrate to zstd compression
This commit is contained in:
Yiorgis Gozadinos 2026-04-08 14:16:12 +03:00 committed by GitHub
commit 23d2aee955
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 636 additions and 98 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
CHANGELOG.md
View file

@ -1,6 +1,24 @@
# Changelog
## [Unreleased]
## [0.38.0] - 2026-04-07
### Added
- **Separate page storage**: Page images stored in dedicated `docling_pages` column — search/expand never loads page data
- **zstd compression**: Switch from gzip to zstd for docling document storage (Python 3.14 stdlib, zstandard package for older versions)
- **`Document.set_docling()`**: Helper method that handles split compression and version assignment, replacing 11 manual call sites
- **`Document.get_page_images()`**: Load page images without the document structure, for visualize_chunk
- **`DocumentRepository.get_pages_data()`**: Load only page data column for a document
### Changed
- **Database migration required**: Run `haiku-rag migrate` to split existing docling blobs into structure + pages and re-compress with zstd
### Fixed
- **Generated skill `domain_preamble`**: Apply `config.prompts.domain_preamble` to instructions in generated skill packages
## [0.37.0] - 2026-04-07
### Changed
@ -1304,7 +1322,8 @@ Existing documents without DoclingDocument data will work but won't have provena
- Initial version tracking
[Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.37.0...HEAD
[Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.38.0...HEAD
[0.38.0]: https://github.com/ggozad/haiku.rag/compare/0.37.0...0.38.0
[0.37.0]: https://github.com/ggozad/haiku.rag/compare/0.36.3...0.37.0
[0.36.3]: https://github.com/ggozad/haiku.rag/compare/0.36.2...0.36.3
[0.36.2]: https://github.com/ggozad/haiku.rag/compare/0.36.1...0.36.2

2
app/backend/pyproject.toml
View file

@ -8,7 +8,7 @@ dependencies = [
"uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=1.70.0",
"python-dotenv>=1.2.1",
"haiku.rag-slim>=0.37.0",
"haiku.rag-slim>=0.38.0",
"logfire[pydantic-ai]>=3.17.0",
]

2
docs/agents/rlm.md
View file

@ -62,7 +62,7 @@ The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https:
| `list_documents(limit, offset)` | List documents in the knowledge base |
| `get_document(id_or_title)` | Get full text content of a document |
| `get_chunk(chunk_id)` | Get a chunk with metadata (headings, page numbers, labels) for citations |
| `get_docling_document(document_id)` | Get the full DoclingDocument structure as a dict (texts, tables, pictures, pages) |
| `get_docling_document(document_id)` | Get the DoclingDocument structure as a dict (texts, tables, pictures) |
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
When documents are pre-loaded via the `documents` parameter, they are injected as a `documents` variable accessible in the sandbox code.

3
docs/cli.md
View file

@ -375,10 +375,11 @@ Error: Database requires migration from 0.19.0 to 0.26.5. 3 migration(s) pending
Run `haiku-rag migrate` to apply the pending migrations. The command shows which migrations were applied:
```
Applied 3 migration(s):
Applied 4 migration(s):
- 0.20.0: Add 'docling_document_json' and 'docling_version' columns
- 0.23.1: Add content_fts column for contextualized FTS search
- 0.25.0: Compress docling_document with gzip
- 0.38.0: Split docling_document pages into separate column and re-compress with zstd
Migration completed successfully.
```

2
docs/python.md
View file

@ -64,7 +64,7 @@ The `format` parameter controls how text content is parsed:
- `"plain"` - Plain text, no parsing (creates a simple text document)
!!! note
The document's `content` field stores the markdown export of the parsed document for consistent display. The original input is preserved in the `docling_document_json` field.
The document's `content` field stores the markdown export of the parsed document for consistent display. The original DoclingDocument structure is preserved in the `docling_document` field (zstd-compressed, without page images). Page images are stored separately in `docling_pages`.
From file:
```python

2
evaluations/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-evals"
description = "Benchmarking and evaluation scripts for haiku.rag"
version = "0.37.0"
version = "0.38.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
requires-python = ">=3.12"

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

49
haiku_rag_slim/haiku/rag/store/compression.py
View file

@ -1,11 +1,50 @@
import gzip
import json
try: # pragma: no cover
from compression.zstd import ( # ty: ignore[unresolved-import]
compress as _zstd_compress, # type: ignore[import-not-found]
)
from compression.zstd import ( # ty: ignore[unresolved-import]
decompress as _zstd_decompress, # type: ignore[import-not-found]
)
except ImportError:
from zstandard import ZstdCompressor, ZstdDecompressor, get_frame_parameters
_zstd_compressor = ZstdCompressor()
_zstd_decompressor = ZstdDecompressor()
def _zstd_compress(data: bytes) -> bytes:
return _zstd_compressor.compress(data)
def _zstd_decompress(data: bytes) -> bytes:
content_size = get_frame_parameters(data).content_size
return _zstd_decompressor.decompress(data, max_output_size=content_size)
def compress_json(json_str: str) -> bytes:
"""Compress a JSON string with gzip."""
return gzip.compress(json_str.encode("utf-8"))
"""Compress a JSON string with zstd."""
return _zstd_compress(json_str.encode("utf-8"))
def decompress_json(data: bytes) -> str:
"""Decompress gzip-compressed data to a JSON string."""
return gzip.decompress(data).decode("utf-8")
"""Decompress zstd-compressed data to a JSON string."""
return _zstd_decompress(data).decode("utf-8")
def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
"""Split a DoclingDocument JSON into structure and pages, compress both with zstd.
Returns:
Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the
document has no page images.
"""
data = json.loads(json_str)
pages = data.pop("pages", None)
structure_bytes = _zstd_compress(json.dumps(data).encode("utf-8"))
pages_bytes = None
if pages:
pages_bytes = _zstd_compress(json.dumps(pages).encode("utf-8"))
return structure_bytes, pages_bytes

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)

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

@ -5,23 +5,21 @@ 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)
def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument":
"""Decompress and validate DoclingDocument, stripping page images."""
"""Decompress and validate DoclingDocument."""
from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data)
data = json.loads(json_str)
data.pop("pages", None)
return DoclingDocument.model_validate(data)
return DoclingDocument.model_validate_json(json_str)
def _get_cached_docling_document(
@ -30,8 +28,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 +38,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 +54,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,
},

4
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -78,7 +78,11 @@ from haiku.rag.store.upgrades.v0_23_1 import (
from haiku.rag.store.upgrades.v0_25_0 import (
upgrade_compress_docling_document as upgrade_0_25_0_compress,
)
from haiku.rag.store.upgrades.v0_38_0 import (
upgrade_split_pages_zstd as upgrade_0_38_0_split_pages,
)
upgrades.append(upgrade_0_20_0_docling)
upgrades.append(upgrade_0_23_1_contextualize)
upgrades.append(upgrade_0_25_0_compress)
upgrades.append(upgrade_0_38_0_split_pages)

244
haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py Normal file
View file

@ -0,0 +1,244 @@
import gzip
import json
import logging
from datetime import timedelta
import pyarrow as pa
from lancedb.pydantic import LanceModel
from pydantic import Field
from haiku.rag.store.compression import compress_docling_split
from haiku.rag.store.engine import Store
from haiku.rag.store.upgrades import Upgrade
logger = logging.getLogger(__name__)
BATCH_SIZE = 5
def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
"""Split docling_document into structure + pages and re-compress with zstd."""
class DocumentRecordV5(LanceModel):
id: str
content: str
uri: str | None = None
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: "")
def get_documents_arrow_schema_v5() -> pa.Schema:
"""Generate Arrow schema with large_binary for both docling columns."""
base_schema = DocumentRecordV5.to_arrow_schema()
large_binary_columns = {"docling_document", "docling_pages"}
fields = []
for field in base_schema:
if field.name in large_binary_columns:
fields.append(pa.field(field.name, pa.large_binary()))
else:
fields.append(field)
return pa.schema(fields)
def migrate_row(row: dict) -> DocumentRecordV5:
"""Migrate a single row: decompress gzip, split pages, re-compress with zstd."""
docling_blob = row.get("docling_document")
structure_bytes: bytes | None = None
pages_bytes: bytes | None = None
if docling_blob and isinstance(docling_blob, bytes):
# Decompress from gzip
try:
json_str = gzip.decompress(docling_blob).decode("utf-8")
except Exception:
# May already be zstd or uncompressed — try as-is
json_str = docling_blob.decode("utf-8")
# Split structure and pages, re-compress with zstd
structure_bytes, pages_bytes = compress_docling_split(json_str)
metadata_raw = row.get("metadata")
metadata_str = (
metadata_raw
if isinstance(metadata_raw, str)
else json.dumps(metadata_raw or {})
)
return DocumentRecordV5(
id=row.get("id") or "",
content=row.get("content", ""),
uri=row.get("uri"),
title=row.get("title"),
metadata=metadata_str,
docling_document=structure_bytes,
docling_pages=pages_bytes,
docling_version=row.get("docling_version"),
created_at=row.get("created_at", ""),
updated_at=row.get("updated_at", ""),
)
def copy_staging_row(row: dict) -> DocumentRecordV5:
"""Copy a row from the staging table (already migrated)."""
return DocumentRecordV5(
id=row["id"],
content=row["content"],
uri=row["uri"],
title=row["title"],
metadata=row["metadata"],
docling_document=row["docling_document"],
docling_pages=row["docling_pages"],
docling_version=row["docling_version"],
created_at=row["created_at"],
updated_at=row["updated_at"],
)
staging_name = "documents_v5_staging"
# First pass: collect document IDs to process
try:
ids = [
row["id"]
for row in store.documents_table.search()
.select(["id"])
.to_arrow()
.to_pylist()
]
except (pa.ArrowInvalid, pa.ArrowNotImplementedError, OSError):
ids = []
if not ids:
# Check for staging table from a failed migration
if staging_name in store.db.list_tables().tables:
staging_table = store.db.open_table(staging_name)
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
if staging_ids:
logger.info(
"Recovering %d documents from failed migration", len(staging_ids)
)
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
total_batches = (len(staging_ids) + BATCH_SIZE - 1) // BATCH_SIZE
for batch_num, i in enumerate(
range(0, len(staging_ids), BATCH_SIZE), 1
):
batch_ids = staging_ids[i : i + BATCH_SIZE]
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
staging_table.search()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
records = [copy_staging_row(row) for row in batch]
if records:
store.documents_table.add(records)
logger.info("Recovered batch %d/%d", batch_num, total_batches)
store.db.drop_table(staging_name)
logger.info("Recovery complete")
return
# No documents and no staging — recreate table with new schema
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
return
# Create staging table with new schema
if staging_name in store.db.list_tables().tables:
store.db.drop_table(staging_name)
staging_table = store.db.create_table(
staging_name, schema=get_documents_arrow_schema_v5()
)
# Migrate in batches: read from old, split+recompress, write to staging
total_docs = len(ids)
total_batches = (total_docs + BATCH_SIZE - 1) // BATCH_SIZE
logger.info(
"Splitting pages and re-compressing %d documents in %d batches",
total_docs,
total_batches,
)
for batch_num, i in enumerate(range(0, len(ids), BATCH_SIZE), 1):
batch_ids = ids[i : i + BATCH_SIZE]
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
store.documents_table.search()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
migrated_batch = [migrate_row(row) for row in batch]
if migrated_batch:
staging_table.add(migrated_batch)
logger.info(
"Migrated batch %d/%d (%d documents)",
batch_num,
total_batches,
len(migrated_batch),
)
# Replace old table with staging table
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
# Copy from staging to final table in batches
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
logger.info("Copying %d documents to new table", len(staging_ids))
for batch_num, i in enumerate(range(0, len(staging_ids), BATCH_SIZE), 1):
batch_ids = staging_ids[i : i + BATCH_SIZE]
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist()
)
records = [copy_staging_row(row) for row in batch]
if records:
store.documents_table.add(records)
logger.info("Copied batch %d/%d", batch_num, total_batches)
# Cleanup staging table
if staging_name in store.db.list_tables().tables:
store.db.drop_table(staging_name)
# Vacuum all tables
logger.info("Vacuuming database")
for table in [store.documents_table, store.chunks_table, store.settings_table]:
try:
table.optimize(cleanup_older_than=timedelta(seconds=0))
except Exception:
pass
logger.info("Migration complete")
upgrade_split_pages_zstd = Upgrade(
version="0.38.0",
apply=_apply_split_pages_zstd,
description="Split docling_document pages into separate column and re-compress with zstd",
)

3
haiku_rag_slim/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.37.0"
version = "0.38.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -38,6 +38,7 @@ dependencies = [
"rich>=14.3.3",
"typer>=0.21.0,<0.22.0",
"watchfiles>=1.1.1",
"zstandard>=0.23.0; python_version<'3.14'",
]
[project.optional-dependencies]

4
pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.37.0"
version = "0.38.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -30,7 +30,7 @@ classifiers = [
]
dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.37.0",
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.38.0",
]
[project.scripts]

56
tests/store/test_compression.py
View file

@ -1,4 +1,10 @@
from haiku.rag.store.compression import compress_json, decompress_json
import json
from haiku.rag.store.compression import (
compress_docling_split,
compress_json,
decompress_json,
)
class TestJsonCompression:
@ -13,3 +19,51 @@ class TestJsonCompression:
compressed = compress_json(json_str)
decompressed = decompress_json(compressed)
assert decompressed == json_str
def test_compress_json_produces_zstd(self):
compressed = compress_json('{"test": true}')
assert compressed[:4] == b"\x28\xb5\x2f\xfd"
class TestDoclingCompressionSplit:
def test_split_with_pages(self):
data = {
"name": "test_doc",
"texts": [{"text": "hello"}],
"pages": {"1": {"image": "base64data"}, "2": {"image": "more"}},
}
json_str = json.dumps(data)
structure_bytes, pages_bytes = compress_docling_split(json_str)
assert structure_bytes is not None
assert pages_bytes is not None
# Structure should not contain pages
structure = json.loads(decompress_json(structure_bytes))
assert "pages" not in structure
assert structure["name"] == "test_doc"
assert structure["texts"] == [{"text": "hello"}]
# Pages should contain only pages
pages = json.loads(decompress_json(pages_bytes))
assert "1" in pages
assert "2" in pages
def test_split_without_pages(self):
data = {"name": "test_doc", "texts": []}
json_str = json.dumps(data)
structure_bytes, pages_bytes = compress_docling_split(json_str)
assert structure_bytes is not None
assert pages_bytes is None
structure = json.loads(decompress_json(structure_bytes))
assert structure["name"] == "test_doc"
def test_split_with_empty_pages(self):
data = {"name": "test_doc", "texts": [], "pages": {}}
json_str = json.dumps(data)
structure_bytes, pages_bytes = compress_docling_split(json_str)
assert structure_bytes is not None
assert pages_bytes is None

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)

145
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
@ -279,6 +367,63 @@ async def test_get_docling_data_loads_only_docling_columns(
store.close()
@pytest.mark.asyncio
async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_db_path):
"""get_pages_data returns only page image data for a document."""
import json
from haiku.rag.store.compression import compress_json
pages_blob = compress_json(
json.dumps({"1": {"size": {"width": 612, "height": 792}, "page_no": 1}})
)
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_pages=pages_blob,
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_pages == pages_blob
# Non-existent ID returns None
assert await doc_repo.get_pages_data("nonexistent-id") is None
store.close()
@pytest.mark.asyncio
async def test_get_pages_data_none_for_markdown_document(
qa_corpus: Dataset, temp_db_path
):
"""Markdown documents have no page images — get_pages_data returns None pages."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.md",
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.docling_pages is None
store.close()
@pytest.mark.asyncio
async def test_document_get_by_uri_with_special_characters(
qa_corpus: Dataset, temp_db_path

8
uv.lock
View file

@ -1418,7 +1418,7 @@ wheels = [
[[package]]
name = "haiku-rag"
version = "0.37.0"
version = "0.38.0"
source = { editable = "." }
dependencies = [
{ name = "haiku-rag-slim", extra = ["cohere", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@ -1475,7 +1475,7 @@ dev = [
[[package]]
name = "haiku-rag-evals"
version = "0.37.0"
version = "0.38.0"
source = { editable = "evaluations" }
dependencies = [
{ name = "datasets" },
@ -1500,7 +1500,7 @@ requires-dist = [
[[package]]
name = "haiku-rag-slim"
version = "0.37.0"
version = "0.38.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "cachetools" },
@ -1519,6 +1519,7 @@ dependencies = [
{ name = "rich" },
{ name = "typer" },
{ name = "watchfiles" },
{ name = "zstandard", marker = "python_full_version < '3.14'" },
]
[package.optional-dependencies]
@ -1599,6 +1600,7 @@ requires-dist = [
{ name = "typer", specifier = ">=0.21.0,<0.22.0" },
{ name = "watchfiles", specifier = ">=1.1.1" },
{ name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a11" },
{ name = "zstandard", marker = "python_full_version < '3.14'", specifier = ">=0.23.0" },
]
provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]