From 0b5f89ebadc355cf14fac82c5a44bf9c27ce4d26 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 16 Jun 2025 19:28:45 +0200 Subject: [PATCH] FTS and hybrid search using Reciprocal Rank Fusion --- src/haiku/rag/store/engine.py | 9 ++ src/haiku/rag/store/repositories/chunk.py | 179 ++++++++++++++++++++-- tests/test_chunk.py | 2 +- tests/test_search.py | 18 ++- 4 files changed, 192 insertions(+), 16 deletions(-) diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index b0f53893..180eaa19 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -50,6 +50,15 @@ class Store: ) """) + # Create FTS5 table for full-text search + db.execute(""" + CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5( + content, + content='chunks', + content_rowid='id' + ) + """) + # Create indexes for better performance db.execute( "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)" diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 5cb0d3e5..70026238 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -1,4 +1,5 @@ import json +import re from haiku.rag.chunker import chunker from haiku.rag.embeddings.ollama import Embedder @@ -40,6 +41,15 @@ class ChunkRepository(BaseRepository[Chunk]): (entity.id, serialized_embedding), ) + # Insert into FTS5 table for full-text search + cursor.execute( + """ + INSERT INTO chunks_fts(rowid, content) + VALUES (?, ?) + """, + (entity.id, entity.content), + ) + if commit: self.store._connection.commit() return entity @@ -103,6 +113,16 @@ class ChunkRepository(BaseRepository[Chunk]): (serialized_embedding, entity.id), ) + # Update FTS5 table + cursor.execute( + """ + UPDATE chunks_fts + SET content = ? + WHERE rowid = ? + """, + (entity.content, entity.id), + ) + self.store._connection.commit() return entity @@ -113,7 +133,10 @@ class ChunkRepository(BaseRepository[Chunk]): cursor = self.store._connection.cursor() - # Delete the embedding first + # Delete from FTS5 table first + cursor.execute("DELETE FROM chunks_fts WHERE rowid = ?", (entity_id,)) + + # Delete the embedding cursor.execute("DELETE FROM chunk_embeddings WHERE chunk_id = ?", (entity_id,)) # Delete the chunk @@ -197,7 +220,9 @@ class ChunkRepository(BaseRepository[Chunk]): self.store._connection.commit() return deleted_any - async def search_chunks(self, query: str, limit: int = 5) -> list[Chunk]: + async def search_chunks( + self, query: str, limit: int = 5 + ) -> list[tuple[Chunk, float]]: """Search for relevant chunks using vector similarity.""" if self.store._connection is None: raise ValueError("Store connection is not available") @@ -224,17 +249,151 @@ class ChunkRepository(BaseRepository[Chunk]): chunks = [] for row in results: - chunk_id, document_id, content, metadata_json, _ = row + chunk_id, document_id, content, metadata_json, distance = row metadata = json.loads(metadata_json) if metadata_json else {} - chunks.append( - Chunk( - id=chunk_id, - document_id=document_id, - content=content, - metadata=metadata, - ) + chunk = Chunk( + id=chunk_id, + document_id=document_id, + content=content, + metadata=metadata, ) + similarity_score = 1.0 / (1.0 + distance) + chunks.append((chunk, similarity_score)) + + return chunks + + async def search_chunks_fts( + self, query: str, limit: int = 5 + ) -> list[tuple[Chunk, float]]: + """Search for chunks using FTS5 full-text search.""" + if self.store._connection is None: + raise ValueError("Store connection is not available") + + cursor = self.store._connection.cursor() + + # Clean the query for FTS5 - extract keywords for better matching + # Remove special characters and split into words + words = re.findall(r"\b\w+\b", query.lower()) + # Join with OR to find chunks containing any of the keywords + fts_query = " OR ".join(words) if words else query + + # Search using FTS5 + cursor.execute( + """ + SELECT c.id, c.document_id, c.content, c.metadata, rank + FROM chunks_fts + JOIN chunks c ON c.id = chunks_fts.rowid + WHERE chunks_fts MATCH ? + ORDER BY rank + LIMIT ? + """, + (fts_query, limit), + ) + + results = cursor.fetchall() + chunks = [] + + for row in results: + chunk_id, document_id, content, metadata_json, rank = row + metadata = json.loads(metadata_json) if metadata_json else {} + chunk = Chunk( + id=chunk_id, + document_id=document_id, + content=content, + metadata=metadata, + ) + # Convert rank to a score - FTS5 rank is negative BM25 score + # More negative = better match, so we negate it to get positive scores + fts_score = -rank + chunks.append((chunk, fts_score)) + + return chunks + + async def search_chunks_hybrid( + self, query: str, limit: int = 5, k: int = 60 + ) -> list[tuple[Chunk, float]]: + """Hybrid search using Reciprocal Rank Fusion (RRF) combining vector similarity and FTS5 full-text search.""" + if self.store._connection is None: + raise ValueError("Store connection is not available") + + cursor = self.store._connection.cursor() + + # Generate embedding for the query + query_embedding = await self.embedder.embed(query) + serialized_query_embedding = self.store.serialize_embedding(query_embedding) + + # Clean the query for FTS5 - extract keywords for better matching + # Remove special characters and split into words + words = re.findall(r"\b\w+\b", query.lower()) + # Join with OR to find chunks containing any of the keywords + fts_query = " OR ".join(words) if words else query + + # Perform hybrid search using RRF (Reciprocal Rank Fusion) + cursor.execute( + """ + WITH vector_search AS ( + SELECT + c.id, + c.document_id, + c.content, + c.metadata, + ROW_NUMBER() OVER (ORDER BY ce.distance) as vector_rank + FROM chunk_embeddings ce + JOIN chunks c ON c.id = ce.chunk_id + WHERE ce.embedding MATCH ? AND k = ? + ORDER BY ce.distance + ), + fts_search AS ( + SELECT + c.id, + c.document_id, + c.content, + c.metadata, + ROW_NUMBER() OVER (ORDER BY chunks_fts.rank) as fts_rank + FROM chunks_fts + JOIN chunks c ON c.id = chunks_fts.rowid + WHERE chunks_fts MATCH ? + ORDER BY chunks_fts.rank + ), + all_chunks AS ( + SELECT id, document_id, content, metadata FROM vector_search + UNION + SELECT id, document_id, content, metadata FROM fts_search + ), + rrf_scores AS ( + SELECT + a.id, + a.document_id, + a.content, + a.metadata, + COALESCE(1.0 / (? + v.vector_rank), 0) + COALESCE(1.0 / (? + f.fts_rank), 0) as rrf_score + FROM all_chunks a + LEFT JOIN vector_search v ON a.id = v.id + LEFT JOIN fts_search f ON a.id = f.id + ) + SELECT id, document_id, content, metadata, rrf_score + FROM rrf_scores + ORDER BY rrf_score DESC + LIMIT ? + """, + (serialized_query_embedding, limit * 3, fts_query, k, k, limit), + ) + + results = cursor.fetchall() + chunks = [] + + for row in results: + chunk_id, document_id, content, metadata_json, rrf_score = row + metadata = json.loads(metadata_json) if metadata_json else {} + chunk = Chunk( + id=chunk_id, + document_id=document_id, + content=content, + metadata=metadata, + ) + chunks.append((chunk, rrf_score)) + return chunks async def get_by_document_id(self, document_id: int) -> list[Chunk]: diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 0fbd3eac..da0216a4 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -33,7 +33,7 @@ async def test_chunk_repository_operations(qa_corpus: Dataset): # Test chunk search results = await chunk_repo.search_chunks("election", limit=2) assert len(results) <= 2 - assert all(hasattr(chunk, "content") for chunk in results) + assert all(hasattr(chunk, "content") for chunk, _ in results) # Test deleting chunks by document ID deleted = await chunk_repo.delete_by_document_id(created_document.id) diff --git a/tests/test_search.py b/tests/test_search.py index 1f23717e..340a2525 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -14,7 +14,7 @@ async def test_search_qa_corpus(qa_corpus: Dataset): store = Store(":memory:") doc_repo = DocumentRepository(store) chunk_repo = ChunkRepository(store) - num_documents = 10 + num_documents = 20 # Load first 10 documents with embeddings (reduced for faster testing) documents = [] for i in range(num_documents): @@ -40,11 +40,19 @@ async def test_search_qa_corpus(qa_corpus: Dataset): target_document, doc_data = documents[i] question = doc_data["question"] - # Search for chunks using the question - search_results = await chunk_repo.search_chunks(question, limit=5) + # Test vector search + vector_results = await chunk_repo.search_chunks(question, limit=5) + target_document_ids = {chunk.document_id for chunk, _ in vector_results} + assert target_document.id in target_document_ids - # Check if target document is in results - target_document_ids = {chunk.document_id for chunk in search_results} + # Test FTS search + fts_results = await chunk_repo.search_chunks_fts(question, limit=5) + 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) + target_document_ids = {chunk.document_id for chunk, _ in hybrid_results} assert target_document.id in target_document_ids store.close()