Move embedding logic from ChunkRepository to client._ensure_chunks_embedded()

haiku_rag_slim/haiku/rag/client.py

@@ -193,6 +193,37 @@ class HaikuRAG:
 
         return chunks
 
+    async def _ensure_chunks_embedded(self, chunks: list[Chunk]) -> list[Chunk]:
+        """Ensure all chunks have embeddings, embedding any that don't.
+
+        Args:
+            chunks: List of chunks, some may have embeddings already.
+
+        Returns:
+            List of chunks with all embeddings populated.
+        """
+        from haiku.rag.embeddings import embed_chunks
+
+        # Find chunks that need embedding
+        chunks_to_embed = [c for c in chunks if c.embedding is None]
+
+        if not chunks_to_embed:
+            return chunks
+
+        # Embed chunks that don't have embeddings (returns new Chunk objects)
+        embedded = await embed_chunks(chunks_to_embed, self._config)
+
+        # Build result maintaining original order
+        embedded_map = {(c.content, c.order): c for c in embedded}
+        result = []
+        for chunk in chunks:
+            if chunk.embedding is not None:
+                result.append(chunk)
+            else:
+                result.append(embedded_map[(chunk.content, chunk.order)])
+
+        return result
+
     async def _store_document_with_chunks(
         self,
         document: Document,
@@ -211,6 +242,9 @@ class HaikuRAG:
         """
         import asyncio
 
+        # Ensure all chunks have embeddings before storing
+        chunks = await self._ensure_chunks_embedded(chunks)
+
         # Snapshot table versions for versioned rollback (if supported)
         versions = self.store.current_table_versions()
 
@@ -258,6 +292,9 @@ class HaikuRAG:
 
         assert document.id is not None, "Document ID is required for update"
 
+        # Ensure all chunks have embeddings before storing
+        chunks = await self._ensure_chunks_embedded(chunks)
+
         # Snapshot table versions for versioned rollback
         versions = self.store.current_table_versions()
 

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -37,20 +37,18 @@ class ChunkRepository:
             logger.debug(f"FTS index creation skipped: {e}")
 
     async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]:
-        """Create one or more chunks in the database."""
+        """Create one or more chunks in the database.
+
+        Chunks must have embeddings set before calling this method.
+        Use client._ensure_chunks_embedded() to embed chunks if needed.
+        """
         # Handle single chunk
         if isinstance(entity, Chunk):
             assert entity.document_id, "Chunk must have a document_id to be created"
+            assert entity.embedding is not None, "Chunk must have an embedding"
 
             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)
-
             chunk_record = self.store.ChunkRecord(
                 id=chunk_id,
                 document_id=entity.document_id,
@@ -58,8 +56,8 @@ class ChunkRepository:
                 metadata=json.dumps(
                     {k: v for k, v in entity.metadata.items() if k != "order"}
                 ),
-                order=order_val,
+                order=int(entity.order),
-                vector=embedding,
+                vector=entity.embedding,
             )
 
             self.store.chunks_table.add([chunk_record])
@@ -72,22 +70,15 @@ class ChunkRepository:
         if not chunks:
             return []
 
-        # Validate all chunks have document_id
+        # Validate all chunks have document_id and embedding
         for chunk in chunks:
             assert chunk.document_id, "All chunks must have a document_id to be created"
-
+            assert chunk.embedding is not None, "All chunks must have embeddings"
-        # 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(
@@ -98,7 +89,7 @@ class ChunkRepository:
                     {k: v for k, v in chunk.metadata.items() if k != "order"}
                 ),
                 order=int(chunk.order),
-                vector=embedding,
+                vector=chunk.embedding,
             )
             chunk_records.append(chunk_record)
             chunk.id = chunk_id
@@ -131,11 +122,12 @@ class ChunkRepository:
         )
 
     async def update(self, entity: Chunk) -> Chunk:
-        """Update an existing chunk."""
+        """Update an existing chunk.
-        assert entity.id, "Chunk ID is required for update"
 
-        embedding = await self.embedder.embed(entity.content)
+        Chunk must have embedding set before calling this method.
-        order_val = int(entity.order)
+        """
+        assert entity.id, "Chunk ID is required for update"
+        assert entity.embedding is not None, "Chunk must have an embedding"
 
         self.store.chunks_table.update(
             where=f"id = '{entity.id}'",
@@ -145,8 +137,8 @@ class ChunkRepository:
                 "metadata": json.dumps(
                     {k: v for k, v in entity.metadata.items() if k != "order"}
                 ),
-                "order": order_val,
+                "order": int(entity.order),
-                "vector": embedding,
+                "vector": entity.embedding,
             },
         )
         return entity

tests/test_chunk.py

@@ -3,11 +3,7 @@ from datasets import Dataset
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
-from haiku.rag.store.engine import Store
 from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
-from haiku.rag.store.models.document import Document
-from haiku.rag.store.repositories.chunk import ChunkRepository
-from haiku.rag.store.repositories.document import DocumentRepository
 
 
 @pytest.mark.asyncio
@@ -77,111 +73,6 @@ async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path):
             assert chunk.order == i
 
 
-@pytest.mark.asyncio
-async def test_chunk_repository_crud(temp_db_path):
-    """Test basic CRUD operations in ChunkRepository."""
-    # Create a store
-    store = Store(temp_db_path, create=True)
-    chunk_repo = ChunkRepository(store)
-    doc_repo = DocumentRepository(store)
-
-    # First create a document to reference
-    document = Document(content="Test document content", metadata={})
-    created_document = await doc_repo.create(document)
-    document_id = created_document.id
-
-    assert document_id is not None, "Document ID should not be None"
-
-    # Test create chunk manually
-    chunk = Chunk(
-        document_id=document_id,
-        content="Test chunk content",
-        metadata={"test": "value"},
-    )
-
-    created_chunk = await chunk_repo.create(chunk)
-    assert isinstance(created_chunk, Chunk)
-    assert created_chunk.id is not None
-    assert created_chunk.content == "Test chunk content"
-
-    # Test get by ID
-    retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
-    assert retrieved_chunk is not None
-    assert retrieved_chunk.content == "Test chunk content"
-    assert retrieved_chunk.metadata["test"] == "value"
-
-    # Test update
-    retrieved_chunk.content = "Updated chunk content"
-    updated_chunk = await chunk_repo.update(retrieved_chunk)
-    assert updated_chunk.content == "Updated chunk content"
-
-    # Test list all
-    all_chunks = await chunk_repo.list_all()
-    assert len(all_chunks) >= 1
-    assert any(chunk.id == created_chunk.id for chunk in all_chunks)
-
-    # Test delete
-    deleted = await chunk_repo.delete(created_chunk.id)
-    assert deleted is True
-
-    # Verify chunk is gone
-    retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
-    assert retrieved_chunk is None
-
-    store.close()
-
-
-@pytest.mark.asyncio
-async def test_adjacent_chunks(temp_db_path):
-    """Test the get_adjacent_chunks repository method."""
-    store = Store(temp_db_path, create=True)
-    doc_repo = DocumentRepository(store)
-    chunk_repo = ChunkRepository(store)
-
-    # Create a simple document first
-    document_content = "Test document for chunking"
-    document = Document(content=document_content)
-    created_document = await doc_repo.create(document)
-
-    # Manually create multiple chunks with order metadata
-    chunks_data = [
-        ("First chunk content", 0),
-        ("Second chunk content", 1),
-        ("Third chunk content", 2),
-        ("Fourth chunk content", 3),
-        ("Fifth chunk content", 4),
-    ]
-
-    created_chunks = []
-    for content, order in chunks_data:
-        chunk = Chunk(document_id=created_document.id, content=content, order=order)
-        created_chunk = await chunk_repo.create(chunk)
-        created_chunks.append(created_chunk)
-
-    # Test with the middle chunk (index 2, order 2)
-    middle_chunk = created_chunks[2]
-
-    # Get adjacent chunks (1 before and after)
-    adjacent_chunks = await chunk_repo.get_adjacent_chunks(middle_chunk, 1)
-
-    # Should have 2 chunks (one before, one after)
-    assert len(adjacent_chunks) == 2
-
-    # Should not include the original chunk
-    assert middle_chunk.id not in [chunk.id for chunk in adjacent_chunks]
-
-    # Should include chunks with order 1 and 3
-    orders = [chunk.order for chunk in adjacent_chunks]
-    assert 1 in orders
-    assert 3 in orders
-
-    # All adjacent chunks should be from the same document
-    for chunk in adjacent_chunks:
-        assert chunk.document_id == created_document.id
-
-    store.close()
-
-
 def test_chunk_metadata_parsing():
     """Test ChunkMetadata parsing from chunk metadata dict."""
     metadata_dict = {