From 09acc845c05f5a52193c499408d543661b047788 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 10 Jan 2026 15:06:56 +0200 Subject: [PATCH 1/6] Update docling, docling-core, lancedb deps --- haiku_rag_slim/pyproject.toml | 6 ++--- uv.lock | 49 ++++++++++++++--------------------- 2 files changed, 23 insertions(+), 32 deletions(-) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 3edd37c7..8f163aa2 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -22,9 +22,9 @@ classifiers = [ ] dependencies = [ - "docling-core==2.57.0", + "docling-core==2.58.1", "httpx>=0.28.1", - "lancedb==0.26.0", + "lancedb==0.26.1", "pathspec>=0.12.1", "pydantic>=2.12.5", "pydantic-ai-slim[openai,fastmcp,logfire,ag-ui]>=1.39.0", @@ -37,7 +37,7 @@ dependencies = [ [project.optional-dependencies] # Document processing -docling = ["docling==2.65.0", "opencv-python-headless>=4.12.0.88"] +docling = ["docling==2.67.0", "opencv-python-headless>=4.12.0.88"] # Embedding providers voyageai = ["voyageai>=0.3.7"] # Rerankers diff --git a/uv.lock b/uv.lock index 15db6020..532b5f62 100644 --- a/uv.lock +++ b/uv.lock @@ -741,7 +741,7 @@ wheels = [ [[package]] name = "docling" -version = "2.65.0" +version = "2.67.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accelerate" }, @@ -773,14 +773,14 @@ dependencies = [ { name = "tqdm" }, { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bb/dc/d42fb4dbe39578ff81d9c52b7a1faaf115c014a38a6e335e6db7ced80658/docling-2.65.0.tar.gz", hash = "sha256:64c6feccf808e19f6100aa9e2a6c71c992cd66e897cdfb33ccf15d0b62873b50", size = 257037, upload-time = "2025-12-15T16:56:16.609Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a7/eb/1900eaf576ab935d5b630a5ba6eb0e7063a2e1d36542df2f9ba572505d57/docling-2.67.0.tar.gz", hash = "sha256:d8c1992bbc090cee8c3f602fa63ac0c3127a65e3cabdf509d23ebb4e14feaf38", size = 259923, upload-time = "2026-01-09T08:37:37.984Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/94/71/ba084670cc4cc3679f7ab650a32ff3ddbf61ad7a1eb559104ee819c028cf/docling-2.65.0-py3-none-any.whl", hash = "sha256:02b00096e2785c08924c9fa89034b918d90f0ee6625aacf20435d1d847c89d09", size = 275389, upload-time = "2025-12-15T16:56:15.003Z" }, + { url = "https://files.pythonhosted.org/packages/8d/99/fa9ac2cffaa77cf46b6414d0a32490b2580c8f3753ab1c65d983ee27f3a0/docling-2.67.0-py3-none-any.whl", hash = "sha256:d3ea4c0d7a64fce4b8f8b38415e075f6637384eaadb27acbf5621673e0da7669", size = 277794, upload-time = "2026-01-09T08:37:36.52Z" }, ] [[package]] name = "docling-core" -version = "2.57.0" +version = "2.58.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonref" }, @@ -794,9 +794,9 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/08/fa/541931a448ddc9609a3e21ca176baf28bd0a89cbbd7be23380637a87d427/docling_core-2.57.0.tar.gz", hash = "sha256:24b71172e8bc4fbe2b4213b29f51878114f4c13f25fc0bc9920168505c76c1a3", size = 230768, upload-time = "2025-12-18T14:18:14.928Z" } +sdist = { url = "https://files.pythonhosted.org/packages/72/06/7a6b5886c1afc7580669b7930c73597050dab149b5fd7ceaf99d7b925e6e/docling_core-2.58.1.tar.gz", hash = "sha256:1ca7b261195c4fc49ce99abbe74b62a07321f00f4ac2eb4faaff095acc8dac82", size = 232575, upload-time = "2026-01-09T16:32:29.304Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/24/1c/03cc087b360b7815c5472f5cfaaf9b09fa6d0397c05c120a7c623654bf21/docling_core-2.57.0-py3-none-any.whl", hash = "sha256:398768b4d275ad2e11703ad9446fd0315a9a5fc8d09bc45ed207302d844d6ed5", size = 223002, upload-time = "2025-12-18T14:18:13.558Z" }, + { url = "https://files.pythonhosted.org/packages/88/fe/6f749c0549c273d03c7d43d75ab132d2eb7ebf70ed3356a609a48bccee90/docling_core-2.58.1-py3-none-any.whl", hash = "sha256:d12590b301d8c9cd60491be64749d02fdde11a5f6923616183a5f6143239bbd8", size = 223970, upload-time = "2026-01-09T16:32:27.76Z" }, ] [package.optional-dependencies] @@ -805,7 +805,7 @@ chunking = [ { name = "transformers" }, { name = "tree-sitter" }, { name = "tree-sitter-c" }, - { name = "tree-sitter-java" }, + { name = "tree-sitter-java-orchard" }, { name = "tree-sitter-javascript" }, { name = "tree-sitter-python" }, { name = "tree-sitter-typescript" }, @@ -1421,10 +1421,10 @@ zeroentropy = [ [package.metadata] requires-dist = [ { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.20.1" }, - { name = "docling", marker = "extra == 'docling'", specifier = "==2.65.0" }, - { name = "docling-core", specifier = "==2.57.0" }, + { name = "docling", marker = "extra == 'docling'", specifier = "==2.67.0" }, + { name = "docling-core", specifier = "==2.58.1" }, { name = "httpx", specifier = ">=0.28.1" }, - { name = "lancedb", specifier = "==0.26.0" }, + { name = "lancedb", specifier = "==0.26.1" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, { name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.12.0.88" }, { name = "pathspec", specifier = ">=0.12.1" }, @@ -1842,7 +1842,7 @@ wheels = [ [[package]] name = "lancedb" -version = "0.26.0" +version = "0.26.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "deprecation" }, @@ -1854,12 +1854,12 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/91/fe585b2181bd61efc65e1da410ae8ab7b29a26f156e4ca7d7d616b1234de/lancedb-0.26.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:3a0d435fff1392f056c173f695f71d495c691c555daa9802c056ea23f6a3900e", size = 41174270, upload-time = "2025-12-16T17:16:30.699Z" }, - { url = "https://files.pythonhosted.org/packages/ce/fc/e47e092f4fc97a8810b37dbee07996689bca42f0817f3f3c38d7fb51dd9d/lancedb-0.26.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a2206320fd0f33c01e264960afd768987646133cf152c4d3a8b7faf81b3017bf", size = 42936720, upload-time = "2025-12-16T17:24:43.527Z" }, - { url = "https://files.pythonhosted.org/packages/b5/d7/323897d22a7c00ef1dc4f5b76df1a11df549fe887d8e05d689c2224e47b8/lancedb-0.26.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ca0322cb4b62d526748f6f29e5b43cce4251c7f693e111897eb1f77e7f1ec2b", size = 45846184, upload-time = "2025-12-16T17:27:33.802Z" }, - { url = "https://files.pythonhosted.org/packages/3a/0b/7671c94b27a5aa267b9f1d6db759c9e08070cb8f783828ade04da9dc7d79/lancedb-0.26.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7f2b8d69a647265b8753576b501354333c3edfd47d12ec9f47e665e8574c92fe", size = 42954293, upload-time = "2025-12-16T17:24:30.335Z" }, - { url = "https://files.pythonhosted.org/packages/52/2e/9f720d6ae7bd3a94d096f320a0ec2f277735423af9d16cf5c61c4a70e6ca/lancedb-0.26.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8e5cc334686a389cf2f28d1c239d13a205098ed98f3914226d3966858e58b957", size = 45896935, upload-time = "2025-12-16T17:27:30.156Z" }, - { url = "https://files.pythonhosted.org/packages/00/0e/4b292c24a9e25ee2cd081d2da930fcdc672ee0eea531fc453c19c73addb5/lancedb-0.26.0-cp39-abi3-win_amd64.whl", hash = "sha256:2fc9b48a11f526de87388002eb3838329db7279241eefb3166c1c6c3b194a3cf", size = 50615000, upload-time = "2025-12-16T17:53:34.409Z" }, + { url = "https://files.pythonhosted.org/packages/45/b5/110651418ceb1fa4ff2eb74ce4bad911ecf49dc765b134f0201d5564aab8/lancedb-0.26.1-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:b1c4389134ede49e4be0497b9719f573f447e627426bb9e6fc1b642db11fb22d", size = 43416143, upload-time = "2026-01-02T17:57:07.232Z" }, + { url = "https://files.pythonhosted.org/packages/81/8a/b48a14281d7875e5bfccf22d911d9e1fa019c1fe7b805d290a4449e3cf60/lancedb-0.26.1-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:07abd18e0aa4730442d0361bab4491ad469de14f9087c3542e56ca6d7fcda473", size = 45302392, upload-time = "2026-01-02T18:04:55.963Z" }, + { url = "https://files.pythonhosted.org/packages/4b/d0/8f6bc531f290206c7a0061236928710506598a2591ff1fcaea477fc52e7f/lancedb-0.26.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df8eb631519c6ede9975099bea187ea25a09e4617a421fe19e5e1613651cd62f", size = 48372676, upload-time = "2026-01-02T18:08:12.373Z" }, + { url = "https://files.pythonhosted.org/packages/f5/13/d8db83335ddf28afe1fb814ca995da7f67826f337d547e54471d7d425dd1/lancedb-0.26.1-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:2941c9f8aa22244002307c4da5d19f12ab77dcb0569eb4f8a48b60e9c4fdee79", size = 45318771, upload-time = "2026-01-02T18:04:26.429Z" }, + { url = "https://files.pythonhosted.org/packages/08/94/10e9d4b5ba49eeba72024d310dc42e0c24feb8d5676f48e989198121a8a0/lancedb-0.26.1-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0d5f125b98836a49095c492085f5ecf3a78906fafcab59c367d9347eb372a4cc", size = 48425627, upload-time = "2026-01-02T18:11:57.443Z" }, + { url = "https://files.pythonhosted.org/packages/17/5d/d7a834ce8dd9c5e6ef7a0e308c7de5f87bb8f04c0944a1bea617d9d42dc7/lancedb-0.26.1-cp39-abi3-win_amd64.whl", hash = "sha256:9338d34c6e7472c97e49fd6b2638b29d3d087e8b002d92cafdbb46a8b0b1480e", size = 53214501, upload-time = "2026-01-02T22:34:06.836Z" }, ] [[package]] @@ -4935,19 +4935,10 @@ wheels = [ ] [[package]] -name = "tree-sitter-java" -version = "0.23.5" +name = "tree-sitter-java-orchard" +version = "0.4.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fa/dc/eb9c8f96304e5d8ae1663126d89967a622a80937ad2909903569ccb7ec8f/tree_sitter_java-0.23.5.tar.gz", hash = "sha256:f5cd57b8f1270a7f0438878750d02ccc79421d45cca65ff284f1527e9ef02e38", size = 138121, upload-time = "2024-12-21T18:24:26.936Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/67/21/b3399780b440e1567a11d384d0ebb1aea9b642d0d98becf30fa55c0e3a3b/tree_sitter_java-0.23.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:355ce0308672d6f7013ec913dee4a0613666f4cda9044a7824240d17f38209df", size = 58926, upload-time = "2024-12-21T18:24:12.53Z" }, - { url = "https://files.pythonhosted.org/packages/57/ef/6406b444e2a93bc72a04e802f4107e9ecf04b8de4a5528830726d210599c/tree_sitter_java-0.23.5-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:24acd59c4720dedad80d548fe4237e43ef2b7a4e94c8549b0ca6e4c4d7bf6e69", size = 62288, upload-time = "2024-12-21T18:24:14.634Z" }, - { url = "https://files.pythonhosted.org/packages/4e/6c/74b1c150d4f69c291ab0b78d5dd1b59712559bbe7e7daf6d8466d483463f/tree_sitter_java-0.23.5-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9401e7271f0b333df39fc8a8336a0caf1b891d9a2b89ddee99fae66b794fc5b7", size = 85533, upload-time = "2024-12-21T18:24:16.695Z" }, - { url = "https://files.pythonhosted.org/packages/29/09/e0d08f5c212062fd046db35c1015a2621c2631bc8b4aae5740d7adb276ad/tree_sitter_java-0.23.5-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:370b204b9500b847f6d0c5ad584045831cee69e9a3e4d878535d39e4a7e4c4f1", size = 84033, upload-time = "2024-12-21T18:24:18.758Z" }, - { url = "https://files.pythonhosted.org/packages/43/56/7d06b23ddd09bde816a131aa504ee11a1bbe87c6b62ab9b2ed23849a3382/tree_sitter_java-0.23.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:aae84449e330363b55b14a2af0585e4e0dae75eb64ea509b7e5b0e1de536846a", size = 82564, upload-time = "2024-12-21T18:24:20.493Z" }, - { url = "https://files.pythonhosted.org/packages/da/d6/0528c7e1e88a18221dbd8ccee3825bf274b1fa300f745fd74eb343878043/tree_sitter_java-0.23.5-cp39-abi3-win_amd64.whl", hash = "sha256:1ee45e790f8d31d416bc84a09dac2e2c6bc343e89b8a2e1d550513498eedfde7", size = 60650, upload-time = "2024-12-21T18:24:22.902Z" }, - { url = "https://files.pythonhosted.org/packages/72/57/5bab54d23179350356515526fff3cc0f3ac23bfbc1a1d518a15978d4880e/tree_sitter_java-0.23.5-cp39-abi3-win_arm64.whl", hash = "sha256:402efe136104c5603b429dc26c7e75ae14faaca54cfd319ecc41c8f2534750f4", size = 59059, upload-time = "2024-12-21T18:24:24.934Z" }, -] +sdist = { url = "https://files.pythonhosted.org/packages/45/80/71a5a4620e667a7d117931bd8bb8d38bbe28ef40687ad1d64b68766f7671/tree_sitter_java_orchard-0.4.2.tar.gz", hash = "sha256:dcddfe590d035a6285c5aa0a01e6a6fb9cf87726386e47056927e1865a2e0835", size = 145077, upload-time = "2025-11-17T15:39:50.364Z" } [[package]] name = "tree-sitter-javascript" From face6f95bc557837ee30d2ded2d4f681f77b6a81 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 10 Jan 2026 15:11:58 +0200 Subject: [PATCH 2/6] compression utils --- haiku_rag_slim/haiku/rag/store/compression.py | 11 +++++++++++ tests/store/test_compression.py | 15 +++++++++++++++ 2 files changed, 26 insertions(+) create mode 100644 haiku_rag_slim/haiku/rag/store/compression.py create mode 100644 tests/store/test_compression.py diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py new file mode 100644 index 00000000..881efb89 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/compression.py @@ -0,0 +1,11 @@ +import gzip + + +def compress_json(json_str: str) -> bytes: + """Compress a JSON string with gzip.""" + return gzip.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") diff --git a/tests/store/test_compression.py b/tests/store/test_compression.py new file mode 100644 index 00000000..068e756b --- /dev/null +++ b/tests/store/test_compression.py @@ -0,0 +1,15 @@ +from haiku.rag.store.compression import compress_json, decompress_json + + +class TestJsonCompression: + def test_compress_json_roundtrip(self): + json_str = '{"key": "value", "number": 42, "nested": {"a": 1}}' + compressed = compress_json(json_str) + decompressed = decompress_json(compressed) + assert decompressed == json_str + + def test_compress_json_with_unicode(self): + json_str = '{"message": "Hello, δΈ–η•Œ! 🌍"}' + compressed = compress_json(json_str) + decompressed = decompress_json(compressed) + assert decompressed == json_str From afc5b02285972fa47c0f398a939ba04be66541bd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 10 Jan 2026 15:17:27 +0200 Subject: [PATCH 3/6] Update schema to store docling_document as compressed bytes - Change DocumentRecord.docling_document_json (str) to docling_document (bytes) - Use large_binary Arrow type (64-bit offsets) to avoid 2GB column limit - Decompress in Document.get_docling_document() using gzip - Update DocumentRepository field mappings --- haiku_rag_slim/haiku/rag/store/engine.py | 26 +++++++++++++++++-- .../haiku/rag/store/models/document.py | 16 ++++++++---- .../haiku/rag/store/repositories/document.py | 10 +++---- tests/test_document.py | 18 ++++++++----- 4 files changed, 52 insertions(+), 18 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 6e2b9ace..ed085bdb 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -8,6 +8,7 @@ from typing import Any from uuid import uuid4 import lancedb +import pyarrow as pa from lancedb.pydantic import LanceModel, Vector from pydantic import Field @@ -24,12 +25,33 @@ class DocumentRecord(LanceModel): uri: str | None = None title: str | None = None metadata: str = Field(default="{}") - docling_document_json: str | None = None + docling_document: 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() -> pa.Schema: + """Generate Arrow schema for documents table with large_binary for docling_document. + + LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets + and is limited to ~2GB per column in a fragment. When many large documents + (with embedded page images) are grouped in a single fragment, this limit is + exceeded, causing "byte array offset overflow" panics. + + This function overrides the default mapping to use `large_binary` instead, + which has 64-bit offsets and no practical size limit. + """ + base_schema = DocumentRecord.to_arrow_schema() + fields = [] + for field in base_schema: + if field.name == "docling_document": + fields.append(pa.field("docling_document", pa.large_binary())) + else: + fields.append(field) + return pa.schema(fields) + + def create_chunk_model(vector_dim: int): """Create a ChunkRecord model with the specified vector dimension. @@ -281,7 +303,7 @@ class Store: self.documents_table = self.db.open_table("documents") else: self.documents_table = self.db.create_table( - "documents", schema=DocumentRecord + "documents", schema=get_documents_arrow_schema() ) # Create or get chunks table diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 5dd6105a..0d98ce9d 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -4,6 +4,8 @@ from typing import TYPE_CHECKING from cachetools import LRUCache from pydantic import BaseModel, Field +from haiku.rag.store.compression import decompress_json + if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -11,13 +13,16 @@ if TYPE_CHECKING: _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) -def _get_cached_docling_document(document_id: str, json_str: str) -> "DoclingDocument": +def _get_cached_docling_document( + document_id: str, compressed_data: bytes +) -> "DoclingDocument": """Get or parse DoclingDocument with LRU caching by document ID.""" if document_id in _docling_document_cache: return _docling_document_cache[document_id] from docling_core.types.doc.document import DoclingDocument + json_str = decompress_json(compressed_data) doc = DoclingDocument.model_validate_json(json_str) _docling_document_cache[document_id] = doc return doc @@ -38,7 +43,7 @@ class Document(BaseModel): uri: str | None = None title: str | None = None metadata: dict = {} - docling_document_json: str | None = None + docling_document: bytes | None = None docling_version: str | None = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) @@ -51,13 +56,14 @@ class Document(BaseModel): Returns: The parsed DoclingDocument, or None if not stored or no ID. """ - if self.docling_document_json is None: + if self.docling_document is None: return None # No caching for documents without ID if self.id is None: from docling_core.types.doc.document import DoclingDocument - return DoclingDocument.model_validate_json(self.docling_document_json) + json_str = decompress_json(self.docling_document) + return DoclingDocument.model_validate_json(json_str) - return _get_cached_docling_document(self.id, self.docling_document_json) + return _get_cached_docling_document(self.id, self.docling_document) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 3e97042f..6daef6b0 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -2,7 +2,7 @@ import json from datetime import datetime from uuid import uuid4 -from haiku.rag.store.engine import DocumentRecord, Store +from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema from haiku.rag.store.models.document import Document @@ -35,7 +35,7 @@ class DocumentRepository: uri=record.uri, title=record.title, metadata=json.loads(record.metadata), - docling_document_json=record.docling_document_json, + docling_document=record.docling_document, docling_version=record.docling_version, created_at=datetime.fromisoformat(record.created_at) if record.created_at @@ -61,7 +61,7 @@ class DocumentRepository: uri=entity.uri, title=entity.title, metadata=json.dumps(entity.metadata), - docling_document_json=entity.docling_document_json, + docling_document=entity.docling_document, docling_version=entity.docling_version, created_at=now, updated_at=now, @@ -111,7 +111,7 @@ class DocumentRepository: "uri": entity.uri, "title": entity.title, "metadata": json.dumps(entity.metadata), - "docling_document_json": entity.docling_document_json, + "docling_document": entity.docling_document, "docling_version": entity.docling_version, "updated_at": now, }, @@ -198,5 +198,5 @@ class DocumentRepository: # Drop and recreate table to clear all data self.store.db.drop_table("documents") self.store.documents_table = self.store.db.create_table( - "documents", schema=DocumentRecord + "documents", schema=get_documents_arrow_schema() ) diff --git a/tests/test_document.py b/tests/test_document.py index 1cf25678..e062ca41 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -70,9 +70,11 @@ def test_document_get_docling_document(): import json + from haiku.rag.store.compression import compress_json + document = Document( content="Test content", - docling_document_json=json.dumps(doc_json), + docling_document=compress_json(json.dumps(doc_json)), docling_version="1.3.0", ) @@ -88,7 +90,7 @@ def test_document_get_docling_document_none(): """Test get_docling_document returns None when not stored.""" document = Document(content="Test content") - assert document.docling_document_json is None + assert document.docling_document is None assert document.get_docling_document() is None @@ -118,13 +120,15 @@ def test_document_get_docling_document_caching(): import json - json_str = json.dumps(doc_json) + from haiku.rag.store.compression import compress_json + + compressed = compress_json(json.dumps(doc_json)) # Clear cache to get clean state _docling_document_cache.clear() document = Document( - id="test-doc-id", content="Test content", docling_document_json=json_str + id="test-doc-id", content="Test content", docling_document=compressed ) # First call - not in cache @@ -157,13 +161,15 @@ def test_document_get_docling_document_no_id_no_cache(): import json - json_str = json.dumps(doc_json) + from haiku.rag.store.compression import compress_json + + compressed = compress_json(json.dumps(doc_json)) # Clear cache _docling_document_cache.clear() # Document without ID - document = Document(content="Test content", docling_document_json=json_str) + document = Document(content="Test content", docling_document=compressed) doc1 = document.get_docling_document() doc2 = document.get_docling_document() From be715dd456dd1db63e27c9e253eec3e341e32120 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 11 Jan 2026 09:23:19 +0200 Subject: [PATCH 4/6] Revert docling core back to 1.8.0 for compatibility with docling-serve --- haiku_rag_slim/pyproject.toml | 2 +- uv.lock | 25 +++++++++++++++++-------- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 8f163aa2..27267c3b 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -22,7 +22,7 @@ classifiers = [ ] dependencies = [ - "docling-core==2.58.1", + "docling-core==2.57.0", "httpx>=0.28.1", "lancedb==0.26.1", "pathspec>=0.12.1", diff --git a/uv.lock b/uv.lock index 532b5f62..7cac008d 100644 --- a/uv.lock +++ b/uv.lock @@ -780,7 +780,7 @@ wheels = [ [[package]] name = "docling-core" -version = "2.58.1" +version = "2.57.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonref" }, @@ -794,9 +794,9 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/72/06/7a6b5886c1afc7580669b7930c73597050dab149b5fd7ceaf99d7b925e6e/docling_core-2.58.1.tar.gz", hash = "sha256:1ca7b261195c4fc49ce99abbe74b62a07321f00f4ac2eb4faaff095acc8dac82", size = 232575, upload-time = "2026-01-09T16:32:29.304Z" } +sdist = { url = "https://files.pythonhosted.org/packages/08/fa/541931a448ddc9609a3e21ca176baf28bd0a89cbbd7be23380637a87d427/docling_core-2.57.0.tar.gz", hash = "sha256:24b71172e8bc4fbe2b4213b29f51878114f4c13f25fc0bc9920168505c76c1a3", size = 230768, upload-time = "2025-12-18T14:18:14.928Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/88/fe/6f749c0549c273d03c7d43d75ab132d2eb7ebf70ed3356a609a48bccee90/docling_core-2.58.1-py3-none-any.whl", hash = "sha256:d12590b301d8c9cd60491be64749d02fdde11a5f6923616183a5f6143239bbd8", size = 223970, upload-time = "2026-01-09T16:32:27.76Z" }, + { url = "https://files.pythonhosted.org/packages/24/1c/03cc087b360b7815c5472f5cfaaf9b09fa6d0397c05c120a7c623654bf21/docling_core-2.57.0-py3-none-any.whl", hash = "sha256:398768b4d275ad2e11703ad9446fd0315a9a5fc8d09bc45ed207302d844d6ed5", size = 223002, upload-time = "2025-12-18T14:18:13.558Z" }, ] [package.optional-dependencies] @@ -805,7 +805,7 @@ chunking = [ { name = "transformers" }, { name = "tree-sitter" }, { name = "tree-sitter-c" }, - { name = "tree-sitter-java-orchard" }, + { name = "tree-sitter-java" }, { name = "tree-sitter-javascript" }, { name = "tree-sitter-python" }, { name = "tree-sitter-typescript" }, @@ -1422,7 +1422,7 @@ zeroentropy = [ requires-dist = [ { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.20.1" }, { name = "docling", marker = "extra == 'docling'", specifier = "==2.67.0" }, - { name = "docling-core", specifier = "==2.58.1" }, + { name = "docling-core", specifier = "==2.57.0" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "lancedb", specifier = "==0.26.1" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, @@ -4935,10 +4935,19 @@ wheels = [ ] [[package]] -name = "tree-sitter-java-orchard" -version = "0.4.2" +name = "tree-sitter-java" +version = "0.23.5" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/45/80/71a5a4620e667a7d117931bd8bb8d38bbe28ef40687ad1d64b68766f7671/tree_sitter_java_orchard-0.4.2.tar.gz", hash = "sha256:dcddfe590d035a6285c5aa0a01e6a6fb9cf87726386e47056927e1865a2e0835", size = 145077, upload-time = "2025-11-17T15:39:50.364Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fa/dc/eb9c8f96304e5d8ae1663126d89967a622a80937ad2909903569ccb7ec8f/tree_sitter_java-0.23.5.tar.gz", hash = "sha256:f5cd57b8f1270a7f0438878750d02ccc79421d45cca65ff284f1527e9ef02e38", size = 138121, upload-time = "2024-12-21T18:24:26.936Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/21/b3399780b440e1567a11d384d0ebb1aea9b642d0d98becf30fa55c0e3a3b/tree_sitter_java-0.23.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:355ce0308672d6f7013ec913dee4a0613666f4cda9044a7824240d17f38209df", size = 58926, upload-time = "2024-12-21T18:24:12.53Z" }, + { url = "https://files.pythonhosted.org/packages/57/ef/6406b444e2a93bc72a04e802f4107e9ecf04b8de4a5528830726d210599c/tree_sitter_java-0.23.5-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:24acd59c4720dedad80d548fe4237e43ef2b7a4e94c8549b0ca6e4c4d7bf6e69", size = 62288, upload-time = "2024-12-21T18:24:14.634Z" }, + { url = "https://files.pythonhosted.org/packages/4e/6c/74b1c150d4f69c291ab0b78d5dd1b59712559bbe7e7daf6d8466d483463f/tree_sitter_java-0.23.5-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9401e7271f0b333df39fc8a8336a0caf1b891d9a2b89ddee99fae66b794fc5b7", size = 85533, upload-time = "2024-12-21T18:24:16.695Z" }, + { url = "https://files.pythonhosted.org/packages/29/09/e0d08f5c212062fd046db35c1015a2621c2631bc8b4aae5740d7adb276ad/tree_sitter_java-0.23.5-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:370b204b9500b847f6d0c5ad584045831cee69e9a3e4d878535d39e4a7e4c4f1", size = 84033, upload-time = "2024-12-21T18:24:18.758Z" }, + { url = "https://files.pythonhosted.org/packages/43/56/7d06b23ddd09bde816a131aa504ee11a1bbe87c6b62ab9b2ed23849a3382/tree_sitter_java-0.23.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:aae84449e330363b55b14a2af0585e4e0dae75eb64ea509b7e5b0e1de536846a", size = 82564, upload-time = "2024-12-21T18:24:20.493Z" }, + { url = "https://files.pythonhosted.org/packages/da/d6/0528c7e1e88a18221dbd8ccee3825bf274b1fa300f745fd74eb343878043/tree_sitter_java-0.23.5-cp39-abi3-win_amd64.whl", hash = "sha256:1ee45e790f8d31d416bc84a09dac2e2c6bc343e89b8a2e1d550513498eedfde7", size = 60650, upload-time = "2024-12-21T18:24:22.902Z" }, + { url = "https://files.pythonhosted.org/packages/72/57/5bab54d23179350356515526fff3cc0f3ac23bfbc1a1d518a15978d4880e/tree_sitter_java-0.23.5-cp39-abi3-win_arm64.whl", hash = "sha256:402efe136104c5603b429dc26c7e75ae14faaca54cfd319ecc41c8f2534750f4", size = 59059, upload-time = "2024-12-21T18:24:24.934Z" }, +] [[package]] name = "tree-sitter-javascript" From 9c4fd08c47e25cc76b2fbc2582cdb6517a671de3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 11 Jan 2026 09:23:56 +0200 Subject: [PATCH 5/6] Update client serialization for compressed docling documents --- haiku_rag_slim/haiku/rag/client.py | 37 +++++++++++++++-------- tests/test_client.py | 48 +++++++++++++++++++----------- tests/test_context_enhancement.py | 2 +- tests/test_converters.py | 16 +++++----- tests/test_rebuild.py | 12 ++++---- 5 files changed, 69 insertions(+), 46 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 48592048..fe8fd11f 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -17,6 +17,7 @@ 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 @@ -372,7 +373,7 @@ class HaikuRAG: embedded_chunks = await embed_chunks(chunks, self._config) # Store markdown export as content for better display/readability - # The original content is preserved in docling_document_json + # The original content is preserved in docling_document stored_content = docling_document.export_to_markdown() # Create document model @@ -381,7 +382,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -417,7 +418,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -550,7 +551,9 @@ class HaikuRAG: # Update existing document and rechunk existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata - existing_doc.docling_document_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title @@ -564,7 +567,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -657,7 +660,9 @@ class HaikuRAG: # Update existing document and rechunk existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata - existing_doc.docling_document_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title @@ -671,7 +676,7 @@ class HaikuRAG: uri=url, title=title, metadata=metadata, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -788,7 +793,9 @@ 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_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version elif content is not None: existing_doc.content = content @@ -798,7 +805,9 @@ 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_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version new_chunks = await self.chunk(docling_document) @@ -810,7 +819,9 @@ class HaikuRAG: # Content provided without chunks - convert, chunk, and embed using primitives existing_doc.content = content # type: ignore[assignment] converted_docling = await self.convert(existing_doc.content) - existing_doc.docling_document_json = converted_docling.model_dump_json() + existing_doc.docling_document = compress_json( + converted_docling.model_dump_json() + ) existing_doc.docling_version = converted_docling.version new_chunks = await self.chunk(converted_docling) @@ -1485,7 +1496,7 @@ class HaikuRAG: uri=doc.uri, title=doc.title, metadata=json.dumps(doc.metadata), - docling_document_json=doc.docling_document_json, + docling_document=doc.docling_document, docling_version=doc.docling_version, created_at=doc.created_at.isoformat() if doc.created_at else now, updated_at=now, @@ -1523,7 +1534,7 @@ class HaikuRAG: embedded_chunks = await embed_chunks(chunks, self._config) # Update document fields - doc.docling_document_json = docling_document.model_dump_json() + doc.docling_document = compress_json(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order @@ -1602,7 +1613,7 @@ class HaikuRAG: chunks = await self.chunk(docling_document) embedded_chunks = await embed_chunks(chunks, self._config) - doc.docling_document_json = docling_document.model_dump_json() + doc.docling_document = compress_json(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order diff --git a/tests/test_client.py b/tests/test_client.py index bfe27c85..b8a8e877 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -8,6 +8,7 @@ from datasets import Dataset from haiku.rag.client import HaikuRAG from haiku.rag.config import Config +from haiku.rag.store.compression import decompress_json from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -790,13 +791,15 @@ async def test_client_create_document_stores_docling_json(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None assert doc.docling_version is not None # Verify JSON is valid and can be parsed import json - parsed = json.loads(doc.docling_document_json) + from haiku.rag.store.compression import decompress_json + + parsed = json.loads(decompress_json(doc.docling_document)) assert "version" in parsed assert parsed["version"] == doc.docling_version @@ -824,7 +827,8 @@ 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_json == docling_doc.model_dump_json() + 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 @@ -840,13 +844,13 @@ async def test_client_create_document_from_file_stores_docling_json(temp_db_path assert isinstance(doc, Document) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None assert doc.docling_version is not None # Verify the stored document also has the JSON retrieved = await client.get_document_by_id(doc.id) assert retrieved is not None - assert retrieved.docling_document_json == doc.docling_document_json + assert retrieved.docling_document == doc.docling_document assert retrieved.docling_version == doc.docling_version @@ -857,17 +861,17 @@ async def test_client_update_document_stores_docling_json(temp_db_path): # Create initial document doc = await client.create_document(content="Initial content") assert doc.id is not None - original_json = doc.docling_document_json + original_json = doc.docling_document # Update content via update_document updated_doc = await client.update_document( document_id=doc.id, content="New content via fields update" ) - assert updated_doc.docling_document_json is not None + assert updated_doc.docling_document is not None assert updated_doc.docling_version is not None # JSON should be different because content changed - assert updated_doc.docling_document_json != original_json + assert updated_doc.docling_document != original_json @pytest.mark.vcr() @@ -879,7 +883,7 @@ async def test_client_update_document_with_custom_chunks_no_docling_json( # Create initial document doc = await client.create_document(content="Initial content") assert doc.id is not None - original_json = doc.docling_document_json + original_json = doc.docling_document # Update with custom chunks custom_chunks = [Chunk(content="Custom chunk", order=0)] @@ -888,7 +892,7 @@ async def test_client_update_document_with_custom_chunks_no_docling_json( ) # Docling JSON should remain unchanged (no conversion when custom chunks provided) - assert updated_doc.docling_document_json == original_json + assert updated_doc.docling_document == original_json @pytest.mark.vcr() @@ -942,7 +946,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_json == docling_doc.model_dump_json() + 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 # Chunks should be regenerated @@ -981,7 +989,11 @@ 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_json == docling_doc.model_dump_json() + assert updated_doc.docling_document is not None + assert ( + decompress_json(updated_doc.docling_document) + == docling_doc.model_dump_json() + ) # Custom chunks should be used (not rechunked from docling) chunks = await client.chunk_repository.get_by_document_id(doc.id) @@ -1001,7 +1013,7 @@ async def test_client_file_update_stores_docling_json(temp_db_path): # Create initial document doc1 = await client.create_document_from_source(temp_path) assert isinstance(doc1, Document) - original_json = doc1.docling_document_json + original_json = doc1.docling_document original_version = doc1.docling_version # Modify file @@ -1013,8 +1025,8 @@ async def test_client_file_update_stores_docling_json(temp_db_path): assert doc2.id == doc1.id # Same document # Docling JSON should be updated - assert doc2.docling_document_json is not None - assert doc2.docling_document_json != original_json + assert doc2.docling_document is not None + assert doc2.docling_document != original_json assert doc2.docling_version == original_version # Version stays same @@ -1038,7 +1050,7 @@ async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None chunks = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks) >= 1 @@ -1095,7 +1107,7 @@ async def test_client_visualize_chunk_with_pdf(temp_db_path): doc = await client.create_document_from_source(pdf_path) assert isinstance(doc, Document) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None chunks = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks) > 0 @@ -1345,7 +1357,7 @@ async def test_client_create_document_with_html_format(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None # Verify the DoclingDocument has proper structure docling_doc = doc.get_docling_document() diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 8ec6abe8..54f48f54 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -613,7 +613,7 @@ This is paragraph four about topic C. ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None # Get chunks which should have doc_item_refs chunks = await client.chunk_repository.get_by_document_id(doc.id) diff --git a/tests/test_converters.py b/tests/test_converters.py index 1e832c70..f8281f60 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -20,7 +20,7 @@ def vcr_cassette_dir(): return str(Path(__file__).parent / "cassettes" / "test_converters") -def create_mock_docling_document_json(name: str = "test") -> dict: +def create_mock_docling_document(name: str = "test") -> dict: """Create a minimal valid DoclingDocument JSON structure for mocking.""" return { "schema_name": "DoclingDocument", @@ -476,7 +476,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_text_success(self, converter): """Test successful text conversion via docling-serve async workflow.""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -498,7 +498,7 @@ class TestDoclingServeConverter: config.providers.docling_serve.api_key = "test-key" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -527,7 +527,7 @@ class TestDoclingServeConverter: config.processing.conversion_options.images_scale = 3.0 converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -622,7 +622,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_file_pdf(self, converter): """Test converting PDF file via docling-serve async workflow.""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -645,7 +645,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_file_text(self, converter): """Test converting text file (reads locally, sends to docling-serve).""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -734,7 +734,7 @@ class TestDoclingServeConverterPictureDescription: config.prompts.picture_description = "Test prompt for picture description" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -767,7 +767,7 @@ class TestDoclingServeConverterPictureDescription: """Test that picture description is disabled by default.""" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index b9737d6b..912a5fc1 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -10,7 +10,7 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -23,7 +23,7 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): # Verify DoclingDocument JSON is preserved after rebuild doc_after = await client.document_repository.get_by_id(doc.id) assert doc_after is not None - assert doc_after.docling_document_json is not None + assert doc_after.docling_document is not None assert doc_after.docling_version is not None chunks_after = await client.chunk_repository.get_by_document_id(doc.id) @@ -40,7 +40,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None - original_docling_json = doc.docling_document_json + original_docling_json = doc.docling_document chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -57,7 +57,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): # DoclingDocument JSON should be unchanged (embed-only doesn't touch documents) doc_after = await client.document_repository.get_by_id(doc.id) assert doc_after is not None - assert doc_after.docling_document_json == original_docling_json + assert doc_after.docling_document == original_docling_json chunks_after = await client.chunk_repository.get_by_document_id(doc.id) chunk_ids_after = {c.id for c in chunks_after} @@ -112,7 +112,7 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None # Set a fake URI to simulate a document that came from a file doc.uri = "file:///nonexistent/path.txt" @@ -133,7 +133,7 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): doc_after = await client.document_repository.get_by_id(doc.id) assert doc_after is not None assert doc_after.content == content_before - assert doc_after.docling_document_json is not None + assert doc_after.docling_document is not None assert doc_after.docling_version is not None chunks_after = await client.chunk_repository.get_by_document_id(doc.id) From 29c682d8eb4ceb1a1af88c62fa7e7b303e58f1db Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 11 Jan 2026 11:39:46 +0200 Subject: [PATCH 6/6] Add migration --- CHANGELOG.md | 13 ++ .../haiku/rag/store/upgrades/__init__.py | 4 + .../haiku/rag/store/upgrades/v0_25_0.py | 200 ++++++++++++++++++ haiku_rag_slim/pyproject.toml | 2 +- uv.lock | 2 +- 5 files changed, 219 insertions(+), 2 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d995a26..ba9606b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,19 @@ # Changelog ## [Unreleased] +### Fixed + +- **Large Document Storage Overflow**: Fixed "byte array offset overflow" panic when vacuuming/rebuilding databases with many large PDF documents ([#225](https://github.com/ggozad/haiku.rag/issues/225)) + - Root cause: Arrow's 32-bit string column offsets limited to ~2GB per fragment + - Changed `docling_document_json` (string) to `docling_document` (bytes) with `large_binary` Arrow type (64-bit offsets) + - Added gzip compression for DoclingDocument JSON (~1.4x compression ratio) + - Migration automatically compresses existing documents in batches to avoid memory issues + - **Breaking**: Migration is destructive - all table version history is lost after upgrade + +### Changed + +- **Dependencies**: Updated lancedb 0.26.0 β†’ 0.26.1, docling 2.65.0 β†’ 2.67.0 + ## [0.24.2] - 2026-01-08 ### Fixed diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 43a93dc4..3dbd02e6 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -69,6 +69,9 @@ from haiku.rag.store.upgrades.v0_20_0 import ( from haiku.rag.store.upgrades.v0_23_1 import ( upgrade_contextualize_chunks as upgrade_0_23_1_contextualize, ) +from haiku.rag.store.upgrades.v0_25_0 import ( + upgrade_compress_docling_document as upgrade_0_25_0_compress, +) upgrades.append(upgrade_0_9_3_order) upgrades.append(upgrade_0_9_3_fts) @@ -76,3 +79,4 @@ upgrades.append(upgrade_0_10_1_add_title) upgrades.append(upgrade_0_19_6_embeddings) upgrades.append(upgrade_0_20_0_docling) upgrades.append(upgrade_0_23_1_contextualize) +upgrades.append(upgrade_0_25_0_compress) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py new file mode 100644 index 00000000..8c0a8eac --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py @@ -0,0 +1,200 @@ +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_json, decompress_json +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades import Upgrade + +logger = logging.getLogger(__name__) + +BATCH_SIZE = 10 + + +def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover + """Migrate docling_document_json (str) to docling_document (compressed bytes).""" + + class DocumentRecordV4(LanceModel): + id: str + content: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + docling_document: 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_v4() -> pa.Schema: + """Generate Arrow schema with large_binary for docling_document.""" + base_schema = DocumentRecordV4.to_arrow_schema() + fields = [] + for field in base_schema: + if field.name == "docling_document": + fields.append(pa.field("docling_document", pa.large_binary())) + else: + fields.append(field) + return pa.schema(fields) + + def migrate_row(row: dict) -> DocumentRecordV4: + """Migrate a single row, compressing docling_document.""" + docling_json = row.get("docling_document_json") or row.get("docling_document") + docling_bytes: bytes | None = None + + if docling_json: + if isinstance(docling_json, str): + docling_bytes = compress_json(docling_json) + elif isinstance(docling_json, bytes): + try: + decompress_json(docling_json) + docling_bytes = docling_json # Already compressed + except Exception: + docling_bytes = compress_json(docling_json.decode("utf-8")) + + metadata_raw = row.get("metadata") + metadata_str = ( + metadata_raw + if isinstance(metadata_raw, str) + else json.dumps(metadata_raw or {}) + ) + + return DocumentRecordV4( + id=row.get("id") or "", + content=row.get("content", ""), + uri=row.get("uri"), + title=row.get("title"), + metadata=metadata_str, + docling_document=docling_bytes, + docling_version=row.get("docling_version"), + created_at=row.get("created_at", ""), + updated_at=row.get("updated_at", ""), + ) + + # 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 Exception: + ids = [] + + if not ids: + # No documents, just recreate table with new schema + try: + store.db.drop_table("documents") + except Exception: + pass + store.documents_table = store.db.create_table( + "documents", schema=get_documents_arrow_schema_v4() + ) + return + + # Create staging table with new schema + try: + store.db.drop_table("documents_v4_staging") + except Exception: + pass + staging_table = store.db.create_table( + "documents_v4_staging", schema=get_documents_arrow_schema_v4() + ) + + # Migrate in batches: read from old, compress, write to staging + total_docs = len(ids) + total_batches = (total_docs + BATCH_SIZE - 1) // BATCH_SIZE + logger.info("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"'{id}'" for 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( + "Compressed batch %d/%d (%d documents)", + batch_num, + total_batches, + len(migrated_batch), + ) + + # Replace old table with staging table + try: + store.db.drop_table("documents") + except Exception: + pass + + store.documents_table = store.db.create_table( + "documents", schema=get_documents_arrow_schema_v4() + ) + + # 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"'{id}'" for id in batch_ids) + + batch = ( + staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist() + ) + records = [ + DocumentRecordV4( + id=row["id"], + content=row["content"], + uri=row["uri"], + title=row["title"], + metadata=row["metadata"], + docling_document=row["docling_document"], + docling_version=row["docling_version"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + for row in batch + ] + if records: + store.documents_table.add(records) + logger.info("Copied batch %d/%d", batch_num, total_batches) + + # Cleanup staging table + try: + store.db.drop_table("documents_v4_staging") + except Exception: + pass + + # Vacuum all tables (destructive migration, no history preserved) + 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_compress_docling_document = Upgrade( + version="0.25.0", + apply=_apply_compress_docling_document, + description="Compress docling_document with gzip and use large_binary type", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 27267c3b..17e1e4ce 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-slim" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" -version = "0.24.2" +version = "0.25.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } diff --git a/uv.lock b/uv.lock index 7cac008d..caa987cd 100644 --- a/uv.lock +++ b/uv.lock @@ -1362,7 +1362,7 @@ requires-dist = [ [[package]] name = "haiku-rag-slim" -version = "0.24.2" +version = "0.25.0" source = { editable = "haiku_rag_slim" } dependencies = [ { name = "docling-core" },