Accept both single chunks and lists for batch insertion in ChunkRepository.create()

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

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:

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:
+    async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]:
-        """Create a chunk in the database."""
+        """Create one or more chunks in the database."""
-        assert entity.document_id, "Chunk must have a document_id to be created"
+        # 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
+            # Generate embedding if not provided
-        if entity.embedding is not None:
+            if entity.embedding is not None:
-            embedding = entity.embedding
+                embedding = entity.embedding
-        else:
+            else:
-            embedding = await self.embedder.embed(entity.content)
+                embedding = await self.embedder.embed(entity.content)
-        order_val = int(entity.order)
+            order_val = int(entity.order)
 
-        chunk_record = self.store.ChunkRecord(
+            chunk_record = self.store.ChunkRecord(
-            id=chunk_id,
+                id=chunk_id,
-            document_id=entity.document_id,
+                document_id=entity.document_id,
-            content=entity.content,
+                content=entity.content,
-            metadata=json.dumps(
+                metadata=json.dumps(
-                {k: v for k, v in entity.metadata.items() if k != "order"}
+                    {k: v for k, v in entity.metadata.items() if k != "order"}
-            ),
+                ),
-            order=order_val,
+                order=order_val,
-            vector=embedding,
+                vector=embedding,
-        )
+            )
 
-        self.store.chunks_table.add([chunk_record])
+            self.store.chunks_table.add([chunk_record])
 
-        entity.id = chunk_id
+            entity.id = chunk_id
-        return entity
+            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."""

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())