From 68c3fa0f79d56c52aefdf51860c14db00b6cd52a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 11:33:46 +0300 Subject: [PATCH] Add order to SearchResult, add ChunkRepository.get_chunks_in_range(), to use them in _expand_with_chunks to fetch only nearby chunks --- CHANGELOG.md | 4 ++ haiku_rag_slim/haiku/rag/client.py | 38 +++++++---- .../haiku/rag/store/models/chunk.py | 2 + .../haiku/rag/store/repositories/chunk.py | 68 ++++++++++++++++--- 4 files changed, 86 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f95254e..917878d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,10 +4,14 @@ ### 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 ## [0.36.3] - 2026-04-01 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 7b7d0dfc..2521b1a7 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1400,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 @@ -1433,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, diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index d662970b..c11c4f50 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 20957846..b98f033a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -363,22 +363,68 @@ class ChunkRepository: ) return len(df) + 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. + + Args: + document_id: The document ID to get chunks for. + min_order: Minimum order value (inclusive). + max_order: Maximum order value (inclusive). + + 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 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" - chunk_order = chunk.order + min_order = chunk.order - num_adjacent + max_order = chunk.order + num_adjacent - # 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 + where = ( + f"document_id = '{chunk.document_id}'" + f" AND `order` >= {min_order}" + f" AND `order` <= {max_order}" + f" AND id != '{chunk.id}'" + ) + 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"