diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 7f375155..11a60089 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -10,6 +10,7 @@ import httpx 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.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -277,9 +278,9 @@ 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 = 3, k: int = 60, rerank=Config.RERANK ) -> list[tuple[Chunk, float]]: - """Search for relevant chunks using hybrid search (vector similarity + full-text search). + """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking. Args: query: The search query string. @@ -289,7 +290,22 @@ class HaikuRAG: Returns: List of (chunk, score) tuples ordered by relevance. """ - return await self.chunk_repository.search_chunks_hybrid(query, limit, k) + + if not rerank: + return await self.chunk_repository.search_chunks_hybrid(query, limit, k) + + # Get more initial results (3X) for reranking + search_results = await self.chunk_repository.search_chunks_hybrid( + query, limit * 3, k + ) + + # Apply reranking + reranker = get_reranker() + chunks = [chunk for chunk, _ in search_results] + reranked_results = await reranker.rerank(query, chunks, top_n=limit) + + # Return reranked results with scores from reranker + return reranked_results async def ask(self, question: str) -> str: """Ask a question using the configured QA agent. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index d4fd7421..f174a81c 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,6 +19,7 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 + RERANK: bool = False RERANK_PROVIDER: str = "mxbai" RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2"