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

Adapt search to lancedb

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-08-29 15:55:46 +03:00
parent c69daecf53
commit 0592a7d122
No known key found for this signature in database
10 changed files with 132 additions and 121 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
pyproject.toml
View file

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

4
src/haiku/rag/app.py
View file

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

7
src/haiku/rag/cli.py
View file

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

43
src/haiku/rag/client.py
View file

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

23
src/haiku/rag/store/factory.py
View file

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

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

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

4
tests/test_app.py
View file

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

2
tests/test_cli.py
View file

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

9
tests/test_search.py
View file

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

22
uv.lock
View file

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