From 52712dbdb2e4736bb2d20fbb65a4ed3f0581fe2f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 11:04:40 +0300 Subject: [PATCH 1/5] Use .select() projection to fetch only id/uri/title/metadata in search --- .../haiku/rag/store/repositories/chunk.py | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index ed47af5f..20957846 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -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: @@ -429,18 +430,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 +453,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)) From 6fdb15e3b23ddcc2efc4e05f153e656ee9477902 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 11:19:04 +0300 Subject: [PATCH 2/5] Add DocumentRepository.get_docling_data() for lazy loading docling in expand_context --- haiku_rag_slim/haiku/rag/client.py | 19 +++---- .../haiku/rag/store/repositories/document.py | 24 +++++++++ tests/test_document.py | 50 +++++++++++++++++++ 3 files changed, 82 insertions(+), 11 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index aff7e91c..7b7d0dfc 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index f40a6ed3..c829bd57 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -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() diff --git a/tests/test_document.py b/tests/test_document.py index 3f43e12b..9e4c92bc 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -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 From 68c3fa0f79d56c52aefdf51860c14db00b6cd52a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 11:33:46 +0300 Subject: [PATCH 3/5] 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" From 32258d4e2a87e14c357d952762c3b9208dc9c424 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 12:06:28 +0300 Subject: [PATCH 4/5] Add batching to embeddings for huge documents --- CHANGELOG.md | 1 + .../haiku/rag/embeddings/__init__.py | 16 ++++++++-- tests/test_embedder.py | 31 +++++++++++++++++++ 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 917878d8..83352f45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ - **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 ## [0.36.3] - 2026-04-01 diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 30219977..75155484 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -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) ] diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 23083275..ee61557d 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -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.""" From 8a24606784fad9f995062092d22451c0661b564b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 12:34:39 +0300 Subject: [PATCH 5/5] Strip page images from DoclingDocument before validation, unless we use visualize_chunk() --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/client.py | 4 +- .../haiku/rag/store/models/document.py | 55 ++++++++++++++----- .../haiku/rag/store/repositories/chunk.py | 29 ---------- 4 files changed, 45 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 83352f45..faab4702 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ - **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 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 2521b1a7..4217b0f5 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1586,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 [] diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 02406238..05d819f4 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index b98f033a..1165950f 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -397,35 +397,6 @@ class ChunkRepository: 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" - - min_order = chunk.order - num_adjacent - max_order = chunk.order + num_adjacent - - 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" ) -> list[tuple[Chunk, float]]: