diff --git a/pyproject.toml b/pyproject.toml index aec622e1..74d8e3b2 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 6577b318..c9083cbd 100644 --- a/src/haiku/rag/app.py +++ b/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 diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 4a3318c9..6594c923 100644 --- a/src/haiku/rag/cli.py +++ b/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") diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 4e11af3b..e721af09 100644 --- a/src/haiku/rag/client.py +++ b/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 = repos["document"] - self.chunk_repository = repos["chunk"] + self.document_repository = DocumentRepository(self.store) + self.chunk_repository = ChunkRepository(self.store) 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( - query, limit * 3, k - ) + search_limit = limit * 3 + 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 = repos["settings"] + settings_repo = SettingsRepository(self.store) settings_repo.save_current_settings() documents = await self.list_documents() diff --git a/src/haiku/rag/store/factory.py b/src/haiku/rag/store/factory.py deleted file mode 100644 index 38020f79..00000000 --- a/src/haiku/rag/store/factory.py +++ /dev/null @@ -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), - } diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index c6e7bb32..29216bdc 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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 - 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 + return await self._process_search_results(results) 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 - words = re.findall(r"\b\w+\b", query.lower()) - - if not words: + if not query.strip(): return [] - # Search by content similarity (approximate FTS using vector search) - # This is a fallback since LanceDB doesn't have built-in FTS - return await self.search_chunks(query, limit) + # 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( - 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.""" - # For LanceDB, we'll use vector search as the primary method - # In the future, this could be enhanced with additional ranking strategies - return await self.search_chunks(query, limit) + """Hybrid search combining vector and full-text search with native LanceDB RRFReranker.""" + if not query.strip(): + return [] + + # 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 diff --git a/tests/test_app.py b/tests/test_app.py index c1e40b64..037d0f9c 100644 --- a/tests/test_app.py +++ b/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]") diff --git a/tests/test_cli.py b/tests/test_cli.py index 518fb370..d025a16b 100644 --- a/tests/test_cli.py +++ b/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(): diff --git a/tests/test_search.py b/tests/test_search.py index 9f917076..f3e4c41f 100644 --- a/tests/test_search.py +++ b/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} diff --git a/uv.lock b/uv.lock index 9a80e9f9..9b625d70 100644 --- a/uv.lock +++ b/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"