Cleanup

docs/installation.md

@@ -31,5 +31,4 @@ uv pip install haiku.rag[mxbai]
 ## Requirements
 
 - Python 3.10+
-- SQLite 3.38+
 - Ollama (for default embeddings)

mkdocs.yml

@@ -1,5 +1,5 @@
 site_name: haiku.rag
-site_description: Retrieval-Augmented Generation (RAG) library on SQLite.
+site_description: Retrieval-Augmented Generation (RAG) library on LanceDB.
 site_url: https://ggozad.github.io/haiku.rag/
 theme:
   name: material

src/haiku/rag/client.py

@@ -334,27 +334,13 @@ class HaikuRAG:
 
         if reranker is None:
             # No reranking - return direct search results
-            if search_type == "vector":
+            return await self.chunk_repository.search(query, limit, search_type)
-                return await self.chunk_repository.search_chunks(query, limit)
-            elif search_type == "fts":
-                return await self.chunk_repository.search_chunks_fts(query, limit)
-            else:  # hybrid (default)
-                return await self.chunk_repository.search_chunks_hybrid(query, limit)
 
         # Get more initial results (3X) for reranking
         search_limit = limit * 3
-        if search_type == "vector":
+        search_results = await self.chunk_repository.search(
-            search_results = await self.chunk_repository.search_chunks(
+            query, search_limit, search_type
-                query, search_limit
+        )
-            )
-        elif search_type == "fts":
-            search_results = await self.chunk_repository.search_chunks_fts(
-                query, search_limit
-            )
-        else:  # hybrid (default)
-            search_results = await self.chunk_repository.search_chunks_hybrid(
-                query, search_limit
-            )
 
         # Apply reranking
         chunks = [chunk for chunk, _ in search_results]
@@ -564,8 +550,6 @@ class HaikuRAG:
                 )
                 yield doc.id
 
-        # LanceDB doesn't need explicit commits
-
     def close(self):
         """Close the underlying store connection."""
         self.store.close()

src/haiku/rag/store/engine.py

@@ -47,9 +47,7 @@ class Store:
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
 
-        # For file paths, create a LanceDB directory structure
+        self.db = lancedb.connect(db_path)
-        lance_path = str(db_path).replace(".sqlite", ".lancedb")
-        self.db = lancedb.connect(lance_path)
 
         self.create_or_update_db()
 

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

@@ -28,16 +28,14 @@ class ChunkRepository:
         """Create a chunk in the database."""
         assert entity.document_id, "Chunk must have a document_id to be created"
 
+        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)
 
-        # Generate new UUID
-        chunk_id = str(uuid4())
-
-        # Create chunk record
         chunk_record = self.store.ChunkRecord(
             id=chunk_id,
             document_id=entity.document_id,
@@ -46,7 +44,6 @@ class ChunkRepository:
             vector=embedding,
         )
 
-        # Add to table
         self.store.chunks_table.add([chunk_record])
 
         entity.id = chunk_id
@@ -76,10 +73,8 @@ class ChunkRepository:
         """Update an existing chunk."""
         assert entity.id, "Chunk ID is required for update"
 
-        # Generate new embedding
         embedding = await self.embedder.embed(entity.content)
 
-        # Update the record
         self.store.chunks_table.update(
             where=f"id = '{entity.id}'",
             values={
@@ -94,12 +89,10 @@ class ChunkRepository:
 
     async def delete(self, entity_id: str) -> bool:
         """Delete a chunk by its ID."""
-        # Check if chunk exists
         chunk = await self.get_by_id(entity_id)
         if chunk is None:
             return False
 
-        # Delete the chunk
         self.store.chunks_table.delete(f"id = '{entity_id}'")
         return True
 
@@ -145,7 +138,6 @@ class ChunkRepository:
     async def delete_all(self) -> bool:
         """Delete all chunks from the database."""
         try:
-            # Get count before deletion
             count = len(
                 list(
                     self.store.chunks_table.search()
@@ -171,74 +163,66 @@ class ChunkRepository:
         if not chunks:
             return False
 
-        # Delete chunks by document_id
         self.store.chunks_table.delete(f"document_id = '{document_id}'")
         return True
 
-    async def search_chunks(
+    async def search(
-        self, query: str, limit: int = 5
+        self, query: str, limit: int = 5, search_type: str = "hybrid"
     ) -> list[tuple[Chunk, float]]:
-        """Search for relevant chunks using vector similarity."""
+        """Search for relevant chunks using the specified search method.
-        # Generate embedding for the query
-        query_embedding = await self.embedder.embed(query)
 
-        # Perform vector search with proper query type
+        Args:
-        results = (
+            query: The search query string.
-            self.store.chunks_table.search(query_embedding, query_type="vector")
+            limit: Maximum number of results to return.
-            .limit(limit)
+            search_type: Type of search - "vector", "fts", or "hybrid" (default).
-            .to_pydantic(self.store.ChunkRecord)
-        )
 
-        return await self._process_search_results(results)
+        Returns:
-
+            List of (chunk, score) tuples ordered by relevance.
-    async def search_chunks_fts(
+        """
-        self, query: str, limit: int = 5
-    ) -> list[tuple[Chunk, float]]:
-        """Search for chunks using full-text search."""
         if not query.strip():
             return []
 
-        # Ensure FTS index exists
+        if search_type == "vector":
-        self._ensure_fts_index()
+            query_embedding = await self.embedder.embed(query)
+
+            results = (
+                self.store.chunks_table.search(query_embedding, query_type="vector")
+                .limit(limit)
+                .to_pydantic(self.store.ChunkRecord)
+            )
+
+            return await self._process_search_results(results)
+
+        elif search_type == "fts":
+            # Ensure FTS index exists
+            self._ensure_fts_index()
 
-        # Use LanceDB's native full-text search
-        try:
             results = (
                 self.store.chunks_table.search(query, query_type="fts")
                 .limit(limit)
                 .to_pydantic(self.store.ChunkRecord)
             )
             return await self._process_search_results(results)
-        except Exception:
-            # Fallback to vector search if FTS is not available or fails
-            return await self.search_chunks(query, limit)
 
-    async def search_chunks_hybrid(
+        else:  # hybrid (default)
-        self, query: str, limit: int = 5
+            # Ensure FTS index exists for hybrid search
-    ) -> list[tuple[Chunk, float]]:
+            self._ensure_fts_index()
-        """Hybrid search combining vector and full-text search with native LanceDB RRFReranker."""
-        if not query.strip():
-            return []
 
-        # Ensure FTS index exists for hybrid search
+            query_embedding = await self.embedder.embed(query)
-        self._ensure_fts_index()
 
-        # Generate embedding for the query since LanceDB doesn't have embedding function configured
+            # Create RRF reranker
-        query_embedding = await self.embedder.embed(query)
+            reranker = RRFReranker()
 
-        # Create RRF reranker (k parameter is handled internally)
+            # Perform native hybrid search with RRF reranking
-        reranker = RRFReranker()
+            results = (
-
+                self.store.chunks_table.search(query_type="hybrid")
-        # Perform native hybrid search with RRF reranking
+                .vector(query_embedding)
-        results = (
+                .text(query)
-            self.store.chunks_table.search(query_type="hybrid")
+                .rerank(reranker)
-            .vector(query_embedding)
+                .limit(limit)
-            .text(query)
+                .to_pydantic(self.store.ChunkRecord)
-            .rerank(reranker)
+            )
-            .limit(limit)
+            return await self._process_search_results(results)
-            .to_pydantic(self.store.ChunkRecord)
-        )
-        return await self._process_search_results(results)
 
     async def get_by_document_id(self, document_id: str) -> list[Chunk]:
         """Get all chunks for a specific document."""
@@ -272,7 +256,6 @@ class ChunkRepository:
             for chunk in results
         ]
 
-        # Sort by order if available
         chunks.sort(key=lambda c: c.metadata.get("order", 0))
         return chunks
 

tests/generate_benchmark_db.py

@@ -11,7 +11,7 @@ from haiku.rag.qa import get_qa_agent
 
 console = Console()
 
-db_path = Path(__file__).parent / "data" / "benchmark.sqlite"
+db_path = Path(__file__).parent / "data" / "benchmark.lancedb"
 
 
 async def populate_db():

tests/test_chunk.py

@@ -35,7 +35,7 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
     assert all(chunk.document_id == created_document.id for chunk in chunks)
 
     # Test chunk search
-    results = await chunk_repo.search_chunks("election", limit=2)
+    results = await chunk_repo.search("election", limit=2, search_type="vector")
     assert len(results) <= 2
     assert all(hasattr(chunk, "content") for chunk, _ in results)
 

tests/test_search.py

@@ -47,17 +47,21 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
         question = doc_data["question"]
 
         # Test vector search
-        vector_results = await chunk_repo.search_chunks(question, limit=5)
+        vector_results = await chunk_repo.search(
+            question, limit=5, search_type="vector"
+        )
         target_document_ids = {chunk.document_id for chunk, _ in vector_results}
         assert target_document.id in target_document_ids
 
         # Test FTS search
-        fts_results = await chunk_repo.search_chunks_fts(question, limit=5)
+        fts_results = await chunk_repo.search(question, limit=5, search_type="fts")
         target_document_ids = {chunk.document_id for chunk, _ in fts_results}
         assert target_document.id in target_document_ids
 
         # Test hybrid search
-        hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5)
+        hybrid_results = await chunk_repo.search(
+            question, limit=5, search_type="hybrid"
+        )
         target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
         assert target_document.id in target_document_ids
 
@@ -85,7 +89,7 @@ async def test_chunks_include_document_info(temp_db_path):
     created_document = await doc_repo._create_with_docling(document, docling_document)
 
     # Search for chunks
-    results = await chunk_repo.search_chunks_hybrid("test document", limit=1)
+    results = await chunk_repo.search("test document", limit=1, search_type="hybrid")
 
     assert len(results) > 0
     chunk, _ = results[0]