diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a46598c..06c43009 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,9 @@ ### Changed +- **Chunk Creation**: `ChunkRepository.create()` now accepts both single chunks and lists for batch insertion + - Batch insertion reduces LanceDB version creation when adding multiple chunks with custom chunks + - Batch embedding generation for improved performance with multiple chunks - Updated core dependencies: ## [0.17.1] - 2025-11-18 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 96aebcc5..0ecff0fe 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -461,11 +461,12 @@ class HaikuRAG: # Update document metadata await self.document_repository.update(existing_doc) - # Add new chunks + # Set document_id and order for all chunks for order, chunk in enumerate(chunks): chunk.document_id = document_id chunk.order = order - await self.chunk_repository.create(chunk) + # Batch create all chunks in a single operation + await self.chunk_repository.create(chunks) return existing_doc else: diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 5818ef14..d5083808 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -37,34 +37,77 @@ class ChunkRepository: # Log the error but don't fail - FTS might already exist logger.debug(f"FTS index creation skipped: {e}") - async def create(self, entity: Chunk) -> Chunk: - """Create a chunk in the database.""" - assert entity.document_id, "Chunk must have a document_id to be created" + async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]: + """Create one or more chunks in the database.""" + # Handle single chunk + if isinstance(entity, Chunk): + assert entity.document_id, "Chunk must have a document_id to be created" - chunk_id = str(uuid4()) + chunk_id = str(uuid4()) - # Generate embedding if not provided - if entity.embedding is not None: - embedding = entity.embedding - else: - embedding = await self.embedder.embed(entity.content) - order_val = int(entity.order) + # Generate embedding if not provided + if entity.embedding is not None: + embedding = entity.embedding + else: + embedding = await self.embedder.embed(entity.content) + order_val = int(entity.order) - chunk_record = self.store.ChunkRecord( - id=chunk_id, - document_id=entity.document_id, - content=entity.content, - metadata=json.dumps( - {k: v for k, v in entity.metadata.items() if k != "order"} - ), - order=order_val, - vector=embedding, - ) + chunk_record = self.store.ChunkRecord( + id=chunk_id, + document_id=entity.document_id, + content=entity.content, + metadata=json.dumps( + {k: v for k, v in entity.metadata.items() if k != "order"} + ), + order=order_val, + vector=embedding, + ) - self.store.chunks_table.add([chunk_record]) + self.store.chunks_table.add([chunk_record]) - entity.id = chunk_id - return entity + entity.id = chunk_id + return entity + + # Handle batch of chunks + chunks = entity + if not chunks: + return [] + + # Validate all chunks have document_id + for chunk in chunks: + assert chunk.document_id, "All chunks must have a document_id to be created" + + # Batch generate embeddings for chunks that need them + texts_to_embed = [chunk.content for chunk in chunks if chunk.embedding is None] + embeddings = await self.embedder.embed(texts_to_embed) if texts_to_embed else [] + embedding_iter = iter(embeddings) + + # Prepare all chunk records + chunk_records = [] + for chunk in chunks: + chunk_id = str(uuid4()) + embedding = ( + chunk.embedding if chunk.embedding is not None else next(embedding_iter) + ) + + assert chunk.document_id is not None + chunk_record = self.store.ChunkRecord( + id=chunk_id, + document_id=chunk.document_id, + content=chunk.content, + metadata=json.dumps( + {k: v for k, v in chunk.metadata.items() if k != "order"} + ), + order=int(chunk.order), + vector=embedding, + ) + chunk_records.append(chunk_record) + chunk.id = chunk_id + + # Single batch insert for all chunks + self.store.chunks_table.add(chunk_records) + + return chunks async def get_by_id(self, entity_id: str) -> Chunk | None: """Get a chunk by its ID.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 8891c525..82062058 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -213,10 +213,12 @@ class DocumentRepository: assert created_doc.id is not None, ( "Document ID should not be None after creation" ) + # Set document_id and order for all chunks for order, chunk in enumerate(chunks): chunk.document_id = created_doc.id chunk.order = order - await self.chunk_repository.create(chunk) + # Batch create all chunks in a single operation + await self.chunk_repository.create(chunks) # Vacuum old versions in background (non-blocking) asyncio.create_task(self.store.vacuum())