Optionally perform reranking when searching

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,8 +290,23 @@ class HaikuRAG:
         Returns:
             List of (chunk, score) tuples ordered by relevance.
         """
+
+        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.
 

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"