Adapt search to lancedb

pyproject.toml

@@ -56,7 +56,7 @@ dev = [
     "mkdocs>=1.6.1",
     "mkdocs-material>=9.6.14",
     "pre-commit>=4.2.0",
-    "pyright>=1.1.403",
+    "pyright>=1.1.404",
     "pytest>=8.4.0",
     "pytest-asyncio>=1.0.0",
     "pytest-cov>=6.2.1",

src/haiku/rag/app.py

@@ -53,9 +53,9 @@ class HaikuRAGApp:
             await self.client.delete_document(doc_id)
             self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]")
 
-    async def search(self, query: str, limit: int = 5, k: int = 60):
+    async def search(self, query: str, limit: int = 5):
         async with HaikuRAG(db_path=self.db_path) as self.client:
-            results = await self.client.search(query, limit=limit, k=k)
+            results = await self.client.search(query, limit=limit)
             if not results:
                 self.console.print("[red]No results found.[/red]")
                 return

src/haiku/rag/cli.py

@@ -135,11 +135,6 @@ def search(
         "-l",
         help="Maximum number of results to return",
     ),
-    k: int = typer.Option(
-        60,
-        "--k",
-        help="Reciprocal Rank Fusion k parameter",
-    ),
     db: Path = typer.Option(
         Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
@@ -147,7 +142,7 @@ def search(
     ),
 ):
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.search(query=query, limit=limit, k=k))
+    asyncio.run(app.search(query=query, limit=limit))
 
 
 @cli.command("ask", help="Ask a question using the QA agent")

src/haiku/rag/client.py

@@ -11,9 +11,11 @@ from haiku.rag.config import Config
 from haiku.rag.reader import FileReader
 from haiku.rag.reranking import get_reranker
 from haiku.rag.store.engine import Store
-from haiku.rag.store.factory import create_repositories
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.store.repositories.settings import SettingsRepository
 from haiku.rag.utils import text_to_docling_document
 
 
@@ -34,9 +36,8 @@ class HaikuRAG:
         if not db_path.parent.exists():
             Path.mkdir(db_path.parent, parents=True)
         self.store = Store(db_path, skip_validation=skip_validation)
-        repos = create_repositories(self.store)
+        self.document_repository = DocumentRepository(self.store)
-        self.document_repository = repos["document"]
+        self.chunk_repository = ChunkRepository(self.store)
-        self.chunk_repository = repos["chunk"]
 
     async def __aenter__(self):
         """Async context manager entry."""
@@ -316,14 +317,14 @@ class HaikuRAG:
         return await self.document_repository.list_all(limit=limit, offset=offset)
 
     async def search(
-        self, query: str, limit: int = 5, k: int = 60
+        self, query: str, limit: int = 5, search_type: str = "hybrid"
     ) -> list[tuple[Chunk, float]]:
-        """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking.
+        """Search for relevant chunks using the specified search method with optional reranking.
 
         Args:
             query: The search query string.
             limit: Maximum number of results to return.
-            k: Parameter for Reciprocal Rank Fusion (default: 60).
+            search_type: Type of search - "vector", "fts", or "hybrid" (default).
 
         Returns:
             List of (chunk, score) tuples ordered by relevance.
@@ -332,12 +333,29 @@ class HaikuRAG:
         reranker = get_reranker()
 
         if reranker is None:
-            return await self.chunk_repository.search_chunks_hybrid(query, limit, k)
+            # No reranking - return direct search results
+            if search_type == "vector":
+                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_results = await self.chunk_repository.search_chunks_hybrid(
+        search_limit = limit * 3
-            query, limit * 3, k
+        if search_type == "vector":
-        )
+            search_results = await self.chunk_repository.search_chunks(
+                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]
         reranked_results = await reranker.rerank(query, chunks, top_n=limit)
@@ -507,8 +525,7 @@ class HaikuRAG:
         self.store.recreate_embeddings_table()
 
         # Update settings to current config
-        repos = create_repositories(self.store)
+        settings_repo = SettingsRepository(self.store)
-        settings_repo = repos["settings"]
         settings_repo.save_current_settings()
 
         documents = await self.list_documents()

src/haiku/rag/store/factory.py

@@ -1,23 +0,0 @@
-from pathlib import Path
-
-from haiku.rag.store.engine import Store
-from haiku.rag.store.repositories.chunk import ChunkRepository
-from haiku.rag.store.repositories.document import DocumentRepository
-from haiku.rag.store.repositories.settings import SettingsRepository
-
-
-def create_store(
-    db_path: Path,
-    skip_validation: bool = False,
-) -> Store:
-    """Create a Store instance."""
-    return Store(db_path, skip_validation)
-
-
-def create_repositories(store: Store):
-    """Create repository instances for the store."""
-    return {
-        "chunk": ChunkRepository(store),
-        "document": DocumentRepository(store),
-        "settings": SettingsRepository(store),
-    }

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

@@ -1,8 +1,8 @@
 import json
-import re
 from uuid import uuid4
 
 from docling_core.types.doc.document import DoclingDocument
+from lancedb.rerankers import RRFReranker
 
 from haiku.rag.chunker import chunker
 from haiku.rag.embeddings import get_embedder
@@ -17,6 +17,13 @@ class ChunkRepository:
         self.store = store
         self.embedder = get_embedder()
 
+    def _ensure_fts_index(self) -> None:
+        """Ensure FTS index exists on the content column."""
+        try:
+            self.store.chunks_table.create_fts_index("content", replace=True)
+        except Exception:
+            pass
+
     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"
@@ -175,69 +182,63 @@ class ChunkRepository:
         # Generate embedding for the query
         query_embedding = await self.embedder.embed(query)
 
-        # Perform vector search
+        # Perform vector search with proper query type
         results = (
-            self.store.chunks_table.search(query_embedding)
+            self.store.chunks_table.search(query_embedding, query_type="vector")
             .limit(limit)
             .to_pydantic(self.store.ChunkRecord)
         )
 
-        # Get document info for each chunk
+        return await self._process_search_results(results)
-        chunks_with_scores = []
-        for chunk_record in results:
-            # Get document info
-            doc_results = list(
-                self.store.documents_table.search()
-                .where(f"id = '{chunk_record.document_id}'")
-                .limit(1)
-                .to_pydantic(DocumentRecord)
-            )
-
-            doc_uri = doc_results[0].uri if doc_results else None
-            doc_meta = doc_results[0].metadata if doc_results else "{}"
-
-            chunk = Chunk(
-                id=chunk_record.id,
-                document_id=chunk_record.document_id,
-                content=chunk_record.content,
-                metadata=json.loads(chunk_record.metadata)
-                if chunk_record.metadata
-                else {},
-                document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
-            )
-
-            # LanceDB returns similarity score (higher is better)
-            score = getattr(
-                chunk_record, "_distance", 0.8
-            )  # Default score if not available
-            chunks_with_scores.append(
-                (chunk, 1.0 - score)
-            )  # Convert distance to similarity
-
-        return chunks_with_scores
 
     async def search_chunks_fts(
         self, query: str, limit: int = 5
     ) -> list[tuple[Chunk, float]]:
         """Search for chunks using full-text search."""
-        # Extract keywords for search
+        if not query.strip():
-        words = re.findall(r"\b\w+\b", query.lower())
-
-        if not words:
             return []
 
-        # Search by content similarity (approximate FTS using vector search)
+        # Ensure FTS index exists
-        # This is a fallback since LanceDB doesn't have built-in FTS
+        self._ensure_fts_index()
-        return await self.search_chunks(query, limit)
+
+        # 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(
-        self, query: str, limit: int = 5, k: int = 60
+        self, query: str, limit: int = 5
     ) -> list[tuple[Chunk, float]]:
-        """Hybrid search - for now, just use vector search in LanceDB."""
+        """Hybrid search combining vector and full-text search with native LanceDB RRFReranker."""
-        # For LanceDB, we'll use vector search as the primary method
+        if not query.strip():
-        # In the future, this could be enhanced with additional ranking strategies
+            return []
-        return await self.search_chunks(query, limit)
+
+        # Ensure FTS index exists for hybrid search
+        self._ensure_fts_index()
+
+        # Generate embedding for the query since LanceDB doesn't have embedding function configured
+        query_embedding = await self.embedder.embed(query)
+
+        # Create RRF reranker (k parameter is handled internally)
+        reranker = RRFReranker()
+
+        # Perform native hybrid search with RRF reranking
+        results = (
+            self.store.chunks_table.search(query_type="hybrid")
+            .vector(query_embedding)
+            .text(query)
+            .rerank(reranker)
+            .limit(limit)
+            .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."""
@@ -294,3 +295,41 @@ class ChunkRepository:
                 adjacent_chunks.append(c)
 
         return adjacent_chunks
+
+    async def _process_search_results(self, results) -> list[tuple[Chunk, float]]:
+        """Process search results into chunks with document info and scores."""
+        chunks_with_scores = []
+
+        for chunk_record in results:
+            # Get document info
+            doc_results = list(
+                self.store.documents_table.search()
+                .where(f"id = '{chunk_record.document_id}'")
+                .limit(1)
+                .to_pydantic(DocumentRecord)
+            )
+
+            doc_uri = doc_results[0].uri if doc_results else None
+            doc_meta = doc_results[0].metadata if doc_results else "{}"
+
+            chunk = Chunk(
+                id=chunk_record.id,
+                document_id=chunk_record.document_id,
+                content=chunk_record.content,
+                metadata=json.loads(chunk_record.metadata)
+                if chunk_record.metadata
+                else {},
+                document_uri=doc_uri,
+                document_meta=json.loads(doc_meta) if doc_meta else {},
+            )
+
+            # Get distance score - LanceDB returns _distance (lower is better)
+            distance = getattr(chunk_record, "_distance", 1.0)
+
+            # Convert distance to similarity score (higher is better)
+            # Using exponential decay to convert distance to similarity
+            score = max(0.0, 1.0 / (1.0 + distance))
+
+            chunks_with_scores.append((chunk, score))
+
+        return chunks_with_scores

tests/test_app.py

@@ -150,7 +150,7 @@ async def test_search(app: HaikuRAGApp, monkeypatch):
     with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
         await app.search("query")
 
-    mock_client.search.assert_called_once_with("query", limit=5, k=60)
+    mock_client.search.assert_called_once_with("query", limit=5)
     assert mock_rich_print_search.call_count == len(mock_results)
 
 
@@ -167,7 +167,7 @@ async def test_search_no_results(app: HaikuRAGApp, monkeypatch):
     with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
         await app.search("query")
 
-    mock_client.search.assert_called_once_with("query", limit=5, k=60)
+    mock_client.search.assert_called_once_with("query", limit=5)
     mock_print.assert_called_once_with("[red]No results found.[/red]")
 
 

tests/test_cli.py

@@ -78,7 +78,7 @@ def test_search():
         result = runner.invoke(cli, ["search", "query"])
 
         assert result.exit_code == 0
-        mock_app_instance.search.assert_called_once_with(query="query", limit=5, k=60)
+        mock_app_instance.search.assert_called_once_with(query="query", limit=5)
 
 
 def test_serve():

tests/test_search.py

@@ -14,8 +14,9 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
     store = Store(temp_db_path)
     doc_repo = DocumentRepository(store)
     chunk_repo = ChunkRepository(store)
+
+    # Load first 20 documents with embeddings (reduced for faster testing)
     num_documents = 20
-    # Load first 10 documents with embeddings (reduced for faster testing)
     documents = []
     for i in range(num_documents):
         doc_data = qa_corpus[i]
@@ -41,7 +42,7 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
         )
         documents.append((created_document, doc_data))
 
-    for i in range(5):  # Test with first few documents
+    for i in range(3):  # Test with first few documents
         target_document, doc_data = documents[i]
         question = doc_data["question"]
 
@@ -55,10 +56,6 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
         target_document_ids = {chunk.document_id for chunk, _ in fts_results}
         assert target_document.id in target_document_ids
 
-    for i in range(num_documents):  # Test with first few documents
-        target_document, doc_data = documents[i]
-        question = doc_data["question"]
-
         # Test hybrid search
         hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5)
         target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}

uv.lock

@@ -1033,7 +1033,6 @@ dependencies = [
     { name = "pydantic-ai" },
     { name = "python-dotenv" },
     { name = "rich" },
-    { name = "sqlite-vec" },
     { name = "tiktoken" },
     { name = "typer" },
     { name = "watchfiles" },
@@ -1072,7 +1071,6 @@ requires-dist = [
     { name = "pydantic-ai", specifier = ">=0.7.2" },
     { name = "python-dotenv", specifier = ">=1.1.0" },
     { name = "rich", specifier = ">=14.0.0" },
-    { name = "sqlite-vec", specifier = ">=0.1.6" },
     { name = "tiktoken", specifier = ">=0.9.0" },
     { name = "typer", specifier = ">=0.16.0" },
     { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" },
@@ -1086,7 +1084,7 @@ dev = [
     { name = "mkdocs", specifier = ">=1.6.1" },
     { name = "mkdocs-material", specifier = ">=9.6.14" },
     { name = "pre-commit", specifier = ">=4.2.0" },
-    { name = "pyright", specifier = ">=1.1.403" },
+    { name = "pyright", specifier = ">=1.1.404" },
     { name = "pytest", specifier = ">=8.4.0" },
     { name = "pytest-asyncio", specifier = ">=1.0.0" },
     { name = "pytest-cov", specifier = ">=6.2.1" },
@@ -2775,15 +2773,15 @@ wheels = [
 
 [[package]]
 name = "pyright"
-version = "1.1.403"
+version = "1.1.404"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "nodeenv" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/fe/f6/35f885264ff08c960b23d1542038d8da86971c5d8c955cfab195a4f672d7/pyright-1.1.403.tar.gz", hash = "sha256:3ab69b9f41c67fb5bbb4d7a36243256f0d549ed3608678d381d5f51863921104", size = 3913526, upload-time = "2025-07-09T07:15:52.882Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e2/6e/026be64c43af681d5632722acd100b06d3d39f383ec382ff50a71a6d5bce/pyright-1.1.404.tar.gz", hash = "sha256:455e881a558ca6be9ecca0b30ce08aa78343ecc031d37a198ffa9a7a1abeb63e", size = 4065679, upload-time = "2025-08-20T18:46:14.029Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/49/b6/b04e5c2f41a5ccad74a1a4759da41adb20b4bc9d59a5e08d29ba60084d07/pyright-1.1.403-py3-none-any.whl", hash = "sha256:c0eeca5aa76cbef3fcc271259bbd785753c7ad7bcac99a9162b4c4c7daed23b3", size = 5684504, upload-time = "2025-07-09T07:15:50.958Z" },
+    { url = "https://files.pythonhosted.org/packages/84/30/89aa7f7d7a875bbb9a577d4b1dc5a3e404e3d2ae2657354808e905e358e0/pyright-1.1.404-py3-none-any.whl", hash = "sha256:c7b7ff1fdb7219c643079e4c3e7d4125f0dafcc19d253b47e898d130ea426419", size = 5902951, upload-time = "2025-08-20T18:46:12.096Z" },
 ]
 
 [[package]]
@@ -3501,18 +3499,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e7/9c/0e6afc12c269578be5c0c1c9f4b49a8d32770a080260c333ac04cc1c832d/soupsieve-2.7-py3-none-any.whl", hash = "sha256:6e60cc5c1ffaf1cebcc12e8188320b72071e922c2e897f737cadce79ad5d30c4", size = 36677, upload-time = "2025-04-20T18:50:07.196Z" },
 ]
 
-[[package]]
-name = "sqlite-vec"
-version = "0.1.6"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/88/ed/aabc328f29ee6814033d008ec43e44f2c595447d9cccd5f2aabe60df2933/sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3", size = 164075, upload-time = "2024-11-20T16:40:29.847Z" },
-    { url = "https://files.pythonhosted.org/packages/a7/57/05604e509a129b22e303758bfa062c19afb020557d5e19b008c64016704e/sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361", size = 165242, upload-time = "2024-11-20T16:40:31.206Z" },
-    { url = "https://files.pythonhosted.org/packages/f2/48/dbb2cc4e5bad88c89c7bb296e2d0a8df58aab9edc75853728c361eefc24f/sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3", size = 103704, upload-time = "2024-11-20T16:40:33.729Z" },
-    { url = "https://files.pythonhosted.org/packages/80/76/97f33b1a2446f6ae55e59b33869bed4eafaf59b7f4c662c8d9491b6a714a/sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24", size = 151556, upload-time = "2024-11-20T16:40:35.387Z" },
-    { url = "https://files.pythonhosted.org/packages/6a/98/e8bc58b178266eae2fcf4c9c7a8303a8d41164d781b32d71097924a6bebe/sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c", size = 281540, upload-time = "2024-11-20T16:40:37.296Z" },
-]
-
 [[package]]
 name = "sse-starlette"
 version = "2.3.6"