danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

FTS and hybrid search using Reciprocal Rank Fusion

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-06-16 19:28:45 +02:00
parent 892d1a378d
commit 0b5f89ebad
No known key found for this signature in database
4 changed files with 192 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
src/haiku/rag/store/engine.py
View file

@ -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)"

179
src/haiku/rag/store/repositories/chunk.py
View file

@ -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]:

2
tests/test_chunk.py
View file

@ -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)

18
tests/test_search.py
View file

@ -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()