diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 9ea36c43..836c38b6 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -180,11 +180,18 @@ class HaikuRAG: Returns: List of Chunk objects (without embeddings, without document_id). + Each chunk has its `order` field set to its position in the list. """ from haiku.rag.chunkers import get_chunker chunker = get_chunker(self._config) - return await chunker.chunk(docling_document) + chunks = await chunker.chunk(docling_document) + + # Set order for each chunk + for i, chunk in enumerate(chunks): + chunk.order = i + + return chunks async def _store_document_with_chunks( self, @@ -1298,37 +1305,34 @@ class HaikuRAG: self, documents: list[Document] ) -> AsyncGenerator[str, None]: """Re-embed all chunks without changing chunk boundaries.""" + from haiku.rag.embeddings import contextualize + for doc in documents: assert doc.id is not None - # Get raw chunk records directly from LanceDB - chunk_records = list( - self.store.chunks_table.search() - .where(f"document_id = '{doc.id}'") - .to_pydantic(self.store.ChunkRecord) - ) - if not chunk_records: + # Get existing chunks + chunks = await self.chunk_repository.get_by_document_id(doc.id) + if not chunks: continue - # Batch embed all chunk contents - contents = [rec.content for rec in chunk_records] - embeddings = await self.chunk_repository.embedder.embed(contents) + # Generate new embeddings using contextualize for consistency + texts = contextualize(chunks) + embeddings = await self.chunk_repository.embedder.embed(texts) - # Build updated records only for chunks with changed embeddings + # Build updated records updated_records = [ self.store.ChunkRecord( - id=rec.id, - document_id=rec.document_id, - content=rec.content, - metadata=rec.metadata, - order=rec.order, + id=chunk.id, # type: ignore[arg-type] + document_id=chunk.document_id, # type: ignore[arg-type] + content=chunk.content, + metadata=json.dumps(chunk.metadata), + order=chunk.order, vector=embedding, ) - for rec, embedding in zip(chunk_records, embeddings) - if rec.vector != embedding + for chunk, embedding in zip(chunks, embeddings) ] - # Batch update chunks with changed embeddings + # Batch update all chunks if updated_records: self.store.chunks_table.merge_insert( "id" @@ -1340,76 +1344,68 @@ class HaikuRAG: self, documents: list[Document] ) -> AsyncGenerator[str, None]: """Re-chunk and re-embed from existing document content.""" - converter = get_converter(self._config) + from haiku.rag.embeddings import embed_chunks for doc in documents: assert doc.id is not None - docling_document = await converter.convert_text(doc.content) - # Update document with docling JSON + # Convert content to DoclingDocument + docling_document = await self.convert(doc.content) + + # Chunk and embed + chunks = await self.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks, self._config) + + # Update document with docling JSON and store new chunks doc.docling_document_json = docling_document.model_dump_json() doc.docling_version = docling_document.version - await self.document_repository.update(doc) + await self._update_document_with_chunks(doc, embedded_chunks) - await self.chunk_repository.create_chunks_for_document( - doc.id, docling_document - ) yield doc.id async def _rebuild_full( self, documents: list[Document] ) -> AsyncGenerator[str, None]: """Full rebuild: re-convert from source, re-chunk, re-embed.""" - converter = get_converter(self._config) + from haiku.rag.embeddings import embed_chunks for doc in documents: assert doc.id is not None + + # Try to rebuild from source if available + if doc.uri and self._check_source_accessible(doc.uri): + try: + await self.delete_document(doc.id) + new_doc = await self.create_document_from_source( + source=doc.uri, metadata=doc.metadata or {} + ) + assert isinstance(new_doc, Document) + assert new_doc.id is not None + yield new_doc.id + continue + except Exception as e: + logger.error( + "Error recreating document from source %s: %s", + doc.uri, + e, + ) + continue + + # Fallback: rebuild from stored content if doc.uri: - source_accessible = self._check_source_accessible(doc.uri) - - if source_accessible: - try: - await self.delete_document(doc.id) - new_doc = await self.create_document_from_source( - source=doc.uri, metadata=doc.metadata or {} - ) - assert isinstance(new_doc, Document) - assert new_doc.id is not None - yield new_doc.id - except Exception as e: - logger.error( - "Error recreating document from source %s: %s", - doc.uri, - e, - ) - continue - else: - logger.warning( - "Source missing for %s, re-embedding from content", doc.uri - ) - docling_document = await converter.convert_text(doc.content) - - # Update document with docling JSON - doc.docling_document_json = docling_document.model_dump_json() - doc.docling_version = docling_document.version - await self.document_repository.update(doc) - - await self.chunk_repository.create_chunks_for_document( - doc.id, docling_document - ) - yield doc.id - else: - docling_document = await converter.convert_text(doc.content) - - # Update document with docling JSON - doc.docling_document_json = docling_document.model_dump_json() - doc.docling_version = docling_document.version - await self.document_repository.update(doc) - - await self.chunk_repository.create_chunks_for_document( - doc.id, docling_document + logger.warning( + "Source missing for %s, re-embedding from content", doc.uri ) - yield doc.id + + docling_document = await self.convert(doc.content) + 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_version = docling_document.version + await self._update_document_with_chunks(doc, embedded_chunks) + + yield doc.id def _check_source_accessible(self, uri: str) -> bool: """Check if a document's source URI is accessible.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 04ddcd2a..a8929f07 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -16,9 +16,6 @@ from lancedb.rerankers import RRFReranker from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.models.chunk import Chunk -if TYPE_CHECKING: - from docling_core.types.doc.document import DoclingDocument - logger = logging.getLogger(__name__) @@ -190,55 +187,6 @@ class ChunkRepository: ) return chunks - async def create_chunks_for_document( - self, document_id: str, document: "DoclingDocument" - ) -> list[Chunk]: - """Create chunks and embeddings for a document from DoclingDocument.""" - from haiku.rag.chunkers import get_chunker - - chunker = get_chunker(self.store._config) - chunks = await chunker.chunk(document) - - # Build embedding texts with headings prepended for better semantic search - # The stored content stays raw, but embeddings capture section context - embedding_texts = [] - for chunk in chunks: - chunk_meta = chunk.get_chunk_metadata() - if chunk_meta.headings: - embedding_text = "\n".join(chunk_meta.headings) + "\n" + chunk.content - else: - embedding_text = chunk.content - embedding_texts.append(embedding_text) - embeddings = await self.embedder.embed(embedding_texts) - - # Prepare all chunk records for batch insertion - chunk_records = [] - created_chunks = [] - - for order, (chunk, embedding) in enumerate(zip(chunks, embeddings)): - chunk_id = str(uuid4()) - - chunk_record = self.store.ChunkRecord( - id=chunk_id, - document_id=document_id, - content=chunk.content, - metadata=json.dumps(chunk.metadata), - order=order, - vector=embedding, - ) - chunk_records.append(chunk_record) - - chunk.id = chunk_id - chunk.document_id = document_id - chunk.order = order - created_chunks.append(chunk) - - # Batch insert all chunks at once - if chunk_records: - self.store.chunks_table.add(chunk_records) - - return created_chunks - async def delete_all(self) -> None: """Delete all chunks from the database.""" # Drop and recreate table to clear all data diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 380f8980..c081046e 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -3,7 +3,6 @@ from datasets import Dataset from haiku.rag.client import HaikuRAG from haiku.rag.config import Config -from haiku.rag.converters import get_converter from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk, ChunkMetadata from haiku.rag.store.models.document import Document @@ -53,45 +52,29 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path): @pytest.mark.asyncio -async def test_create_chunks_for_document(qa_corpus: Dataset, temp_db_path): - """Test creating chunks for a document.""" - # Create a store and repositories - store = Store(temp_db_path, create=True) - chunk_repo = ChunkRepository(store) - doc_repo = DocumentRepository(store) +async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path): + """Test document chunking using client primitives.""" + from haiku.rag.client import HaikuRAG + from haiku.rag.embeddings import embed_chunks - # Get the first document from the corpus - first_doc = qa_corpus[0] - document_text = first_doc["document_extracted"] + async with HaikuRAG(db_path=temp_db_path, create=True) as client: + # Get the first document from the corpus + first_doc = qa_corpus[0] + document_text = first_doc["document_extracted"] - # Create a document first (without chunks) - document = Document(content=document_text, metadata={"source": "test"}) - created_document = await doc_repo.create(document) - document_id = created_document.id + # Use client primitives: convert → chunk → embed + docling_document = await client.convert(document_text) + chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks) - assert document_id is not None, "Document ID should not be None" + # Verify chunks were created with embeddings + assert len(chunks) > 0 + assert all(chunk.embedding is None for chunk in chunks) # Before embedding + assert all(chunk.embedding is not None for chunk in embedded_chunks) # After - # Convert text to DoclingDocument - converter = get_converter(Config) - docling_document = await converter.convert_text(document_text, name="test.md") - - # Test creating chunks for the document - chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document) - - # Verify chunks were created - assert len(chunks) > 0 - assert all(chunk.document_id == document_id for chunk in chunks) - assert all(chunk.id is not None for chunk in chunks) - - # Verify chunk order - for i, chunk in enumerate(chunks): - assert chunk.order == i - - # Verify chunks exist in database - db_chunks = await chunk_repo.get_by_document_id(document_id) - assert len(db_chunks) == len(chunks) - - store.close() + # Verify chunk order + for i, chunk in enumerate(chunks): + assert chunk.order == i @pytest.mark.asyncio