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

Merge pull request #337 from ggozad/chore/search-performance

Browse source 
Performance improvements for large documents
This commit is contained in:
Yiorgis Gozadinos 2026-04-07 16:27:24 +03:00 committed by GitHub
commit 506fa050f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 250 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

@ -4,10 +4,16 @@
### Changed
- **Dependency updates**: lancedb 0.30.2, pydantic-ai-slim ≥1.77.0, docling ≥2.84.0, docling-core ≥2.71.0, haiku.skills ≥0.13.0, cachetools ≥7.0.5, pydantic-monty ≥0.0.9, cohere ≥5.21.1, textual ≥8.2.1, ty ≥0.0.28, ruff ≥0.15.9
- **Search result model**: `SearchResult` now includes `order` field propagated from chunk order
### Fixed
- **Type checking**: Fix 37 new ty 0.0.28 diagnostics with proper None guards, assertions, and specific ignore codes
- **Search performance**: Avoid loading full document blobs (docling_document, content) during search — use column projection to fetch only needed metadata (id, uri, title, metadata)
- **Context expansion performance**: Load only docling columns during expand_context (skip content blob), and only when doc_item_refs exist
- **Chunk expansion performance**: Fetch only chunks in the needed order range during context expansion instead of all chunks for a document
- **Embedding batching**: Batch embedding calls in groups of 512 to avoid request size limits and timeouts with large documents
- **DoclingDocument validation**: Strip page images before validation on the read path — pages are only needed for visualize_chunk and account for ~99% of the JSON size
## [0.36.3] - 2026-04-01

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

@ -1114,19 +1114,16 @@ class HaikuRAG:
expanded_results.extend(doc_results)
continue
# Fetch the document to get DoclingDocument
doc = await self.get_document_by_id(doc_id)
if doc is None:
expanded_results.extend(doc_results)
continue
docling_doc = doc.get_docling_document()
# Check if we can use DoclingDocument-based expansion
has_docling = docling_doc is not None
has_refs = any(r.doc_item_refs for r in doc_results)
docling_doc = None
if has_docling and has_refs:
if has_refs:
# Only load docling data when refs exist (skips content blob)
doc = await self.document_repository.get_docling_data(doc_id)
if doc is not None:
docling_doc = doc.get_docling_document()
if docling_doc is not None and has_refs:
# Use DoclingDocument-based expansion
expanded = await self._expand_with_docling(
doc_results,
@ -1403,32 +1400,40 @@ class HaikuRAG:
radius: int,
) -> list[SearchResult]:
"""Expand results using chunk-based adjacency."""
all_chunks = await self.chunk_repository.get_by_document_id(doc_id)
if not all_chunks:
return results
content_to_chunk = {c.content: c for c in all_chunks}
chunk_by_order = {c.order: c for c in all_chunks}
min_order, max_order = min(chunk_by_order.keys()), max(chunk_by_order.keys())
# Build ranges
# Build ranges from result orders
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
for result in results:
chunk = content_to_chunk.get(result.content)
if chunk is None:
if result.chunk_id is None:
passthrough.append(result)
continue
start = max(min_order, chunk.order - radius)
end = min(max_order, chunk.order + radius)
start = result.order - radius
end = result.order + radius
ranges.append((start, end, result))
if not ranges:
return results
# Compute the full order range needed and fetch only those chunks
all_starts = [s for s, _, _ in ranges]
all_ends = [e for _, e, _ in ranges]
range_min = min(all_starts)
range_max = max(all_ends)
chunks_in_range = await self.chunk_repository.get_chunks_in_range(
doc_id, range_min, range_max
)
if not chunks_in_range:
return results
chunk_by_order = {c.order: c for c in chunks_in_range}
# Merge and build results
final_results: list[SearchResult] = []
for min_idx, max_idx, original_results in self._merge_ranges(ranges):
# Collect chunks in order
chunks_in_range = [
merged_chunks = [
chunk_by_order[o]
for o in range(min_idx, max_idx + 1)
if o in chunk_by_order
@ -1436,7 +1441,7 @@ class HaikuRAG:
first = original_results[0]
final_results.append(
SearchResult(
content="".join(c.content for c in chunks_in_range),
content="".join(c.content for c in merged_chunks),
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
document_id=first.document_id,
@ -1581,8 +1586,8 @@ class HaikuRAG:
if not doc:
return []
# Get DoclingDocument
docling_doc = doc.get_docling_document()
# Get DoclingDocument with page images for rendering
docling_doc = doc.get_docling_document(include_pages=True)
if not docling_doc:
return []

16
haiku_rag_slim/haiku/rag/embeddings/__init__.py
View file

@ -53,6 +53,9 @@ def contextualize(chunks: list["Chunk"]) -> list[str]:
return texts
EMBEDDING_BATCH_SIZE = 512
async def embed_chunks(
chunks: list["Chunk"], config: AppConfig = Config
) -> list["Chunk"]:
@ -61,6 +64,9 @@ async def embed_chunks(
Contextualizes chunks (prepends headings) before embedding for better
semantic search. Returns new Chunk objects with embeddings set.
Embeddings are generated in batches to avoid request size limits
and timeouts with large document sets.
Args:
chunks: List of chunks to embed.
config: Configuration for embedder selection.
@ -75,7 +81,13 @@ async def embed_chunks(
embedder = get_embedder(config)
texts = contextualize(chunks)
embeddings = await embedder.embed_documents(texts)
# Batch embedding calls to avoid request size limits
all_embeddings: list[list[float]] = []
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
batch = texts[i : i + EMBEDDING_BATCH_SIZE]
batch_embeddings = await embedder.embed_documents(batch)
all_embeddings.extend(batch_embeddings)
return [
Chunk(
@ -89,7 +101,7 @@ async def embed_chunks(
document_meta=chunk.document_meta,
embedding=embedding,
)
for chunk, embedding in zip(chunks, embeddings)
for chunk, embedding in zip(chunks, all_embeddings)
]

2
haiku_rag_slim/haiku/rag/store/models/chunk.py
View file

@ -117,6 +117,7 @@ class SearchResult(BaseModel):
document_id: str | None = None
document_uri: str | None = None
document_title: str | None = None
order: int = 0
doc_item_refs: list[str] = []
page_numbers: list[int] = []
headings: list[str] | None = None
@ -137,6 +138,7 @@ class SearchResult(BaseModel):
document_id=chunk.document_id,
document_uri=chunk.document_uri,
document_title=chunk.document_title,
order=chunk.order,
doc_item_refs=meta.doc_item_refs,
page_numbers=meta.page_numbers,
headings=meta.headings,

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

@ -1,3 +1,4 @@
import json
from datetime import datetime
from typing import TYPE_CHECKING
@ -13,21 +14,41 @@ if TYPE_CHECKING:
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
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]
def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument":
"""Decompress and validate DoclingDocument, stripping page images."""
from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data)
doc = DoclingDocument.model_validate_json(json_str)
data = json.loads(json_str)
data.pop("pages", None)
return DoclingDocument.model_validate(data)
def _get_cached_docling_document(
document_id: str, compressed_data: bytes
) -> "DoclingDocument":
"""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).
"""
if document_id in _docling_document_cache:
return _docling_document_cache[document_id]
doc = _validate_without_pages(compressed_data)
_docling_document_cache[document_id] = doc
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)
@ -48,22 +69,30 @@ class Document(BaseModel):
created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now)
def get_docling_document(self) -> "DoclingDocument | None":
def get_docling_document(
self, *, include_pages: bool = False
) -> "DoclingDocument | None":
"""Parse and return the stored DoclingDocument.
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.
"""
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:
from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(self.docling_document)
return DoclingDocument.model_validate_json(json_str)
return _validate_without_pages(self.docling_document)
return _get_cached_docling_document(self.id, self.docling_document)

78
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -13,7 +13,7 @@ if TYPE_CHECKING:
from lancedb.rerankers import RRFReranker
from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk
logger = logging.getLogger(__name__)
@ -321,17 +321,18 @@ class ChunkRepository:
results = list(query.to_pydantic(self.store.ChunkRecord))
# Get document info
doc_results = list(
# Get document info (only metadata columns, skip content/docling blobs)
doc_rows = list(
self.store.documents_table.search()
.select(["id", "uri", "title", "metadata"])
.where(f"id = '{document_id}'")
.limit(1)
.to_pydantic(DocumentRecord)
.to_list()
)
doc_uri = doc_results[0].uri if doc_results else None
doc_title = doc_results[0].title if doc_results else None
doc_meta = doc_results[0].metadata if doc_results else "{}"
doc_uri = doc_rows[0]["uri"] if doc_rows else None
doc_title = doc_rows[0]["title"] if doc_rows else None
doc_meta = doc_rows[0].get("metadata", "{}") if doc_rows else "{}"
chunks: list[Chunk] = []
for rec in results:
@ -362,22 +363,39 @@ class ChunkRepository:
)
return len(df)
async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
"""Get adjacent chunks before and after the given chunk within the same document."""
assert chunk.document_id, "Document id is required for adjacent chunk finding"
async def get_chunks_in_range(
self, document_id: str, min_order: int, max_order: int
) -> list[Chunk]:
"""Get chunks for a document within an order range.
chunk_order = chunk.order
Args:
document_id: The document ID to get chunks for.
min_order: Minimum order value (inclusive).
max_order: Maximum order value (inclusive).
# Fetch chunks for the same document and filter by order proximity
all_chunks = await self.get_by_document_id(chunk.document_id)
adjacent_chunks: list[Chunk] = []
for c in all_chunks:
c_order = c.order
if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
adjacent_chunks.append(c)
return adjacent_chunks
Returns:
List of chunks within the order range.
"""
where = (
f"document_id = '{document_id}'"
f" AND `order` >= {min_order}"
f" AND `order` <= {max_order}"
)
results = list(
self.store.chunks_table.search()
.where(where)
.to_pydantic(self.store.ChunkRecord)
)
return [
Chunk(
id=rec.id,
document_id=rec.document_id,
content=rec.content,
metadata=json.loads(rec.metadata),
order=rec.order,
)
for rec in results
]
async def _process_search_results(
self, query_result: "pd.DataFrame | LanceQueryBuilder"
@ -429,18 +447,18 @@ class ChunkRepository:
# Collect all unique document IDs for batch lookup
document_ids = list(set(chunk.document_id for chunk in pydantic_results))
# Batch fetch all documents at once
documents_map = {}
# Batch fetch document metadata (skip content/docling blobs)
documents_map: dict[str, dict] = {}
if document_ids:
# Use IN clause for efficient batch lookup
id_list = "', '".join(document_ids)
where_clause = f"id IN ('{id_list}')"
doc_results = list(
doc_rows = list(
self.store.documents_table.search()
.select(["id", "uri", "title", "metadata"])
.where(where_clause)
.to_pydantic(DocumentRecord)
.to_list()
)
documents_map = {doc.id: doc for doc in doc_results}
documents_map = {str(row["id"]): row for row in doc_rows}
# Build final results with document info
chunks_with_scores = []
@ -452,9 +470,9 @@ class ChunkRepository:
content=chunk_record.content,
metadata=json.loads(chunk_record.metadata),
order=chunk_record.order,
document_uri=doc.uri if doc else None,
document_title=doc.title if doc else None,
document_meta=json.loads(doc.metadata if doc else "{}"),
document_uri=doc["uri"] if doc else None,
document_title=doc["title"] if doc else None,
document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"),
)
score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score))

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

@ -90,6 +90,30 @@ class DocumentRepository:
return self._record_to_document(results[0])
_DOCLING_COLUMNS = ["id", "docling_document", "docling_version"]
async def get_docling_data(self, entity_id: str) -> Document | None:
"""Get a document with only docling data loaded (skips content blob)."""
safe_id = _escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.select(self._DOCLING_COLUMNS)
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not results:
return None
row = results[0]
return Document(
id=row["id"],
content="",
docling_document=row.get("docling_document"),
docling_version=row.get("docling_version"),
)
async def update(self, entity: Document) -> Document:
"""Update an existing document."""
self.store._assert_writable()

50
tests/test_document.py
View file

@ -229,6 +229,56 @@ def test_document_get_docling_document_no_id_no_cache():
assert doc1 is not doc2
@pytest.mark.asyncio
async def test_get_docling_data_loads_only_docling_columns(
qa_corpus: Dataset, temp_db_path
):
"""get_docling_data returns docling blob without loading content."""
import json
from haiku.rag.store.compression import compress_json
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
doc_json = {
"name": "test_doc",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
compressed = compress_json(json.dumps(doc_json))
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_document=compressed,
docling_version="2.1.0",
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_docling_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_document == compressed
assert result.docling_version == "2.1.0"
# Verify docling document can be parsed
docling_doc = result.get_docling_document()
assert docling_doc is not None
assert docling_doc.name == "test_doc"
# Non-existent ID returns None
assert await doc_repo.get_docling_data("nonexistent-id") is None
store.close()
@pytest.mark.asyncio
async def test_document_get_by_uri_with_special_characters(
qa_corpus: Dataset, temp_db_path

31
tests/test_embedder.py
View file

@ -160,6 +160,37 @@ async def test_embed_chunks_empty_list():
assert result == []
async def test_embed_chunks_batches_large_inputs(monkeypatch):
"""Test that embed_chunks batches calls when chunk count exceeds batch size."""
from haiku.rag.embeddings import EMBEDDING_BATCH_SIZE, EmbedderWrapper
call_sizes: list[int] = []
async def tracking_embed(self, texts):
call_sizes.append(len(texts))
return [[0.1] * 10 for _ in texts]
monkeypatch.setattr(EmbedderWrapper, "embed_documents", tracking_embed)
# Create more chunks than one batch
num_chunks = EMBEDDING_BATCH_SIZE + 100
chunks = [
Chunk(id=f"chunk-{i}", content=f"Content {i}", order=i)
for i in range(num_chunks)
]
result = await embed_chunks(chunks)
assert len(result) == num_chunks
assert len(call_sizes) == 2
assert call_sizes[0] == EMBEDDING_BATCH_SIZE
assert call_sizes[1] == 100
# Verify order is preserved
assert result[0].id == "chunk-0"
assert result[-1].id == f"chunk-{num_chunks - 1}"
assert all(r.embedding == [0.1] * 10 for r in result)
@pytest.mark.vcr()
async def test_embed_chunks_preserves_all_fields(allow_model_requests):
"""Test that embed_chunks preserves all chunk fields."""