Batch operations & small optimization

src/haiku/rag/store/engine.py

@@ -1,4 +1,5 @@
 import json
+import logging
 from importlib import metadata
 from pathlib import Path
 from uuid import uuid4
@@ -10,6 +11,8 @@ from pydantic import Field
 from haiku.rag.config import Config
 from haiku.rag.embeddings import get_embedder
 
+logger = logging.getLogger(__name__)
+
 
 class DocumentRecord(LanceModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
@@ -21,14 +24,17 @@ class DocumentRecord(LanceModel):
 
 
 def create_chunk_model(vector_dim: int):
-    """Create a ChunkRecord model with the specified vector dimension."""
+    """Create a ChunkRecord model with the specified vector dimension.
+
+    This creates a model with proper vector typing for LanceDB.
+    """
 
     class ChunkRecord(LanceModel):
         id: str = Field(default_factory=lambda: str(uuid4()))
         document_id: str
         content: str
         metadata: str = Field(default="{}")
-        vector: Vector(vector_dim) = Field(default_factory=list)  # type: ignore
+        vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
 
     return ChunkRecord
 
@@ -46,27 +52,41 @@ class Store:
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
 
-        # Connect to LanceDB (local, cloud, or object storage)
+        # Connect to LanceDB
-        if Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION:
+        self.db = self._connect_to_lancedb(db_path)
-            self.db = lancedb.connect(
+
+        # Initialize tables
+        self.create_or_update_db()
+
+        # Validate config compatibility after connection is established
+        if not skip_validation:
+            self._validate_configuration()
+
+    def _connect_to_lancedb(self, db_path: Path):
+        """Establish connection to LanceDB (local, cloud, or object storage)."""
+        # Check if we have cloud configuration
+        if self._has_cloud_config():
+            return lancedb.connect(
                 uri=Config.LANCEDB_URI,
                 api_key=Config.LANCEDB_API_KEY,
                 region=Config.LANCEDB_REGION,
             )
         else:
             # Local file system connection
-            self.db = lancedb.connect(db_path)
+            return lancedb.connect(db_path)
 
-        self.create_or_update_db()
+    def _has_cloud_config(self) -> bool:
+        """Check if cloud configuration is complete."""
+        return bool(
+            Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION
+        )
 
-        # Validate config compatibility after connection is established
+    def _validate_configuration(self) -> None:
-        if not skip_validation:
+        """Validate that the configuration is compatible with the database."""
-            from haiku.rag.store.repositories.settings import (
+        from haiku.rag.store.repositories.settings import SettingsRepository
-                SettingsRepository,
-            )
 
-            settings_repo = SettingsRepository(self)
+        settings_repo = SettingsRepository(self)
-            settings_repo.validate_config_compatibility()
+        settings_repo.validate_config_compatibility()
 
     def create_or_update_db(self):
         """Create the database tables."""
@@ -114,54 +134,48 @@ class Store:
             )
             if existing_settings:
                 db_version = self.get_haiku_version()  # noqa: F841
-                # XXX Add upgrade logic here similar to SQLite version
+                # TODO: Add upgrade logic here similar to SQLite version when needed
-
         except Exception:
+            # Settings table might not exist yet in fresh databases
             pass
 
     def get_haiku_version(self) -> str:
         """Returns the user version stored in settings."""
-        try:
+        settings_records = list(
-            settings_records = list(
+            self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
-                self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
+        )
+        if settings_records:
+            settings = (
+                json.loads(settings_records[0].settings)
+                if settings_records[0].settings
+                else {}
             )
-            if settings_records:
+            return settings.get("version", "0.0.0")
-                settings = (
-                    json.loads(settings_records[0].settings)
-                    if settings_records[0].settings
-                    else {}
-                )
-                return settings.get("version", "0.0.0")
-        except Exception:
-            pass
         return "0.0.0"
 
     def set_haiku_version(self, version: str) -> None:
         """Updates the user version in settings."""
-        try:
+        settings_records = list(
-            settings_records = list(
+            self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
-                self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
+        )
+        if settings_records:
+            settings = (
+                json.loads(settings_records[0].settings)
+                if settings_records[0].settings
+                else {}
+            )
+            settings["version"] = version
+            # Update the record
+            self.settings_table.update(
+                where="id = 'settings'", values={"settings": json.dumps(settings)}
+            )
+        else:
+            # Create new settings record
+            settings_data = Config.model_dump(mode="json")
+            settings_data["version"] = version
+            self.settings_table.add(
+                [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
             )
-            if settings_records:
-                settings = (
-                    json.loads(settings_records[0].settings)
-                    if settings_records[0].settings
-                    else {}
-                )
-                settings["version"] = version
-                # Update the record
-                self.settings_table.update(
-                    where="id = 'settings'", values={"settings": json.dumps(settings)}
-                )
-            else:
-                # Create new settings record
-                settings_data = Config.model_dump(mode="json")
-                settings_data["version"] = version
-                self.settings_table.add(
-                    [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
-                )
-        except Exception:
-            pass
 
     def recreate_embeddings_table(self) -> None:
         """Recreate the chunks table with current vector dimensions."""

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

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import logging
 from uuid import uuid4
 
 from docling_core.types.doc.document import DoclingDocument
@@ -11,6 +12,8 @@ from haiku.rag.embeddings import get_embedder
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.chunk import Chunk
 
+logger = logging.getLogger(__name__)
+
 
 class ChunkRepository:
     """Repository for Chunk operations."""
@@ -24,8 +27,9 @@ class ChunkRepository:
         """Ensure FTS index exists on the content column."""
         try:
             self.store.chunks_table.create_fts_index("content", replace=True)
-        except Exception:
+        except Exception as e:
-            pass
+            # Log the error but don't fail - FTS might already exist
+            logger.debug(f"FTS index creation skipped: {e}")
 
     async def _optimize(self) -> None:
         """Optimize the chunks table to refresh indexes."""
@@ -36,9 +40,11 @@ class ChunkRepository:
         async with self._optimize_lock:
             try:
                 self.store.chunks_table.optimize()
-            except (RuntimeError, OSError):
+            except (RuntimeError, OSError) as e:
                 # Handle "too many open files" and other resource errors gracefully
-                pass
+                logger.debug(
+                    f"Table optimization skipped due to resource constraints: {e}"
+                )
 
     async def create(self, entity: Chunk) -> Chunk:
         """Create a chunk in the database."""
@@ -147,18 +153,21 @@ class ChunkRepository:
     ) -> list[Chunk]:
         """Create chunks and embeddings for a document from DoclingDocument."""
         chunk_texts = await chunker.chunk(document)
+
+        # Generate embeddings in parallel for all chunks
+        embeddings_tasks = []
+        for chunk_text in chunk_texts:
+            embeddings_tasks.append(self.embedder.embed(chunk_text))
+
+        # Wait for all embeddings to complete
+        embeddings = await asyncio.gather(*embeddings_tasks)
+
+        # Prepare all chunk records for batch insertion
+        chunk_records = []
         created_chunks = []
 
-        for order, chunk_text in enumerate(chunk_texts):
+        for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)):
-            chunk = Chunk(
-                document_id=document_id, content=chunk_text, metadata={"order": order}
-            )
-            # Use create but don't trigger individual optimizations
             chunk_id = str(uuid4())
-            if chunk.embedding is not None:
-                embedding = chunk.embedding
-            else:
-                embedding = await self.embedder.embed(chunk.content)
 
             chunk_record = self.store.ChunkRecord(
                 id=chunk_id,
@@ -167,11 +176,20 @@ class ChunkRepository:
                 metadata=json.dumps({"order": order}),
                 vector=embedding,
             )
-            self.store.chunks_table.add([chunk_record])
+            chunk_records.append(chunk_record)
 
-            chunk.id = chunk_id
+            chunk = Chunk(
+                id=chunk_id,
+                document_id=document_id,
+                content=chunk_text,
+                metadata={"order": order},
+            )
             created_chunks.append(chunk)
 
+        # Batch insert all chunks at once
+        if chunk_records:
+            self.store.chunks_table.add(chunk_records)
+
         # Force optimization once at the end for bulk operations
         await self._optimize()
         return created_chunks
@@ -216,7 +234,7 @@ class ChunkRepository:
             query_embedding = await self.embedder.embed(query)
 
             results = self.store.chunks_table.search(
-                query_embedding, query_type="vector"
+                query_embedding, query_type="vector", vector_column_name="vector"
             ).limit(limit)
 
             return await self._process_search_results(results)
@@ -305,6 +323,7 @@ class ChunkRepository:
         # Get both arrow and pydantic results to access scores
         arrow_result = query_result.to_arrow()
         pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
+
         # Extract scores from arrow result based on search type
         scores = []
         column_names = arrow_result.column_names
@@ -322,17 +341,26 @@ class ChunkRepository:
         else:
             raise ValueError("Unknown search result format, cannot extract scores")
 
-        for i, chunk_record in enumerate(pydantic_results):
+        # Collect all unique document IDs for batch lookup
-            # Get document info
+        document_ids = list(set(chunk.document_id for chunk in pydantic_results))
+
+        # Batch fetch all documents at once
+        documents_map = {}
+        if document_ids:
+            # Create a WHERE clause for all document IDs
+            where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
             doc_results = list(
                 self.store.documents_table.search()
-                .where(f"id = '{chunk_record.document_id}'")
+                .where(where_clause)
-                .limit(1)
                 .to_pydantic(DocumentRecord)
             )
+            documents_map = {doc.id: doc for doc in doc_results}
 
-            doc_uri = doc_results[0].uri if doc_results else None
+        for i, chunk_record in enumerate(pydantic_results):
-            doc_meta = doc_results[0].metadata if doc_results else "{}"
+            # Get document info from pre-fetched map
+            doc = documents_map.get(chunk_record.document_id)
+            doc_uri = doc.uri if doc else None
+            doc_meta = doc.metadata if doc else "{}"
 
             chunk = Chunk(
                 id=chunk_record.id,

src/haiku/rag/store/repositories/document.py

@@ -17,11 +17,31 @@ class DocumentRepository:
 
     def __init__(self, store: Store) -> None:
         self.store = store
+        self._chunk_repository = None
 
-        from haiku.rag.store.repositories.chunk import ChunkRepository
+    @property
+    def chunk_repository(self):
+        """Lazy-load ChunkRepository when needed."""
+        if self._chunk_repository is None:
+            from haiku.rag.store.repositories.chunk import ChunkRepository
 
-        chunk_repository = ChunkRepository(store)
+            self._chunk_repository = ChunkRepository(self.store)
-        self.chunk_repository = chunk_repository
+        return self._chunk_repository
+
+    def _record_to_document(self, record: DocumentRecord) -> Document:
+        """Convert a DocumentRecord to a Document model."""
+        return Document(
+            id=record.id,
+            content=record.content,
+            uri=record.uri,
+            metadata=json.loads(record.metadata) if record.metadata else {},
+            created_at=datetime.fromisoformat(record.created_at)
+            if record.created_at
+            else datetime.now(),
+            updated_at=datetime.fromisoformat(record.updated_at)
+            if record.updated_at
+            else datetime.now(),
+        )
 
     async def create(self, entity: Document) -> Document:
         """Create a document in the database."""
@@ -61,19 +81,7 @@ class DocumentRepository:
         if not results:
             return None
 
-        doc_record = results[0]
+        return self._record_to_document(results[0])
-        return Document(
-            id=doc_record.id,
-            content=doc_record.content,
-            uri=doc_record.uri,
-            metadata=json.loads(doc_record.metadata) if doc_record.metadata else {},
-            created_at=datetime.fromisoformat(doc_record.created_at)
-            if doc_record.created_at
-            else datetime.now(),
-            updated_at=datetime.fromisoformat(doc_record.updated_at)
-            if doc_record.updated_at
-            else datetime.now(),
-        )
 
     async def update(self, entity: Document) -> Document:
         """Update an existing document."""
@@ -104,10 +112,7 @@ class DocumentRepository:
             return False
 
         # Delete associated chunks first
-        from haiku.rag.store.repositories.chunk import ChunkRepository
+        await self.chunk_repository.delete_by_document_id(entity_id)
-
-        chunk_repo = ChunkRepository(self.store)
-        await chunk_repo.delete_by_document_id(entity_id)
 
         # Delete the document
         self.store.documents_table.delete(f"id = '{entity_id}'")
@@ -125,22 +130,7 @@ class DocumentRepository:
             query = query.limit(limit)
 
         results = list(query.to_pydantic(DocumentRecord))
-
+        return [self._record_to_document(doc) for doc in results]
-        return [
-            Document(
-                id=doc.id,
-                content=doc.content,
-                uri=doc.uri,
-                metadata=json.loads(doc.metadata) if doc.metadata else {},
-                created_at=datetime.fromisoformat(doc.created_at)
-                if doc.created_at
-                else datetime.now(),
-                updated_at=datetime.fromisoformat(doc.updated_at)
-                if doc.updated_at
-                else datetime.now(),
-            )
-            for doc in results
-        ]
 
     async def get_by_uri(self, uri: str) -> Document | None:
         """Get a document by its URI."""
@@ -154,27 +144,12 @@ class DocumentRepository:
         if not results:
             return None
 
-        doc_record = results[0]
+        return self._record_to_document(results[0])
-        return Document(
-            id=doc_record.id,
-            content=doc_record.content,
-            uri=doc_record.uri,
-            metadata=json.loads(doc_record.metadata) if doc_record.metadata else {},
-            created_at=datetime.fromisoformat(doc_record.created_at)
-            if doc_record.created_at
-            else datetime.now(),
-            updated_at=datetime.fromisoformat(doc_record.updated_at)
-            if doc_record.updated_at
-            else datetime.now(),
-        )
 
     async def delete_all(self) -> None:
         """Delete all documents from the database."""
         # Delete all chunks first
-        from haiku.rag.store.repositories.chunk import ChunkRepository
+        await self.chunk_repository.delete_all()
-
-        chunk_repo = ChunkRepository(self.store)
-        await chunk_repo.delete_all()
 
         # Get count before deletion
         count = len(