Return tuple (Chunk, score,) in rerank()

src/haiku/rag/reranking/base.py

@@ -7,7 +7,7 @@ class RerankerBase:
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
-    ) -> list[Chunk]:
+    ) -> list[tuple[Chunk, float]]:
         raise NotImplementedError(
             "Reranker is an abstract class. Please implement the rerank method in a subclass."
         )

src/haiku/rag/reranking/cohere.py

@@ -16,7 +16,7 @@ class CohereReranker(RerankerBase):
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
-    ) -> list[Chunk]:
+    ) -> list[tuple[Chunk, float]]:
         if not chunks:
             return []
 
@@ -29,6 +29,6 @@ class CohereReranker(RerankerBase):
         reranked_chunks = []
         for result in response.results:
             original_chunk = chunks[result.index]
-            reranked_chunks.append(original_chunk)
+            reranked_chunks.append((original_chunk, result.relevance_score))
 
         return reranked_chunks

src/haiku/rag/reranking/mxbai.py

@@ -13,7 +13,7 @@ class MxBAIReranker(RerankerBase):
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
-    ) -> list[Chunk]:
+    ) -> list[tuple[Chunk, float]]:
         if not chunks:
             return []
 
@@ -23,6 +23,6 @@ class MxBAIReranker(RerankerBase):
         reranked_chunks = []
         for result in results:
             original_chunk = chunks[result.index]
-            reranked_chunks.append(original_chunk)
+            reranked_chunks.append((original_chunk, result.score))
 
         return reranked_chunks

tests/test_reranker.py

@@ -34,7 +34,8 @@ async def test_mxbai_reranker():
     reranked = await reranker.rerank(
         "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
     )
-    assert [r.document_id for r in reranked] == [0, 2]
+    assert [chunk.document_id for chunk, score in reranked] == [0, 2]
+    assert all(isinstance(score, float) for chunk, score in reranked)
 
 
 @pytest.mark.asyncio
@@ -43,12 +44,13 @@ async def test_cohere_reranker():
         from haiku.rag.reranking.cohere import CohereReranker
 
         reranker = CohereReranker()
-        assert reranker._model == "rerank-v3.5"
+        reranker._model = "rerank-v3.5"
 
         reranked = await reranker.rerank(
             "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
         )
-        assert [r.document_id for r in reranked] == [0, 2]
+        assert [chunk.document_id for chunk, score in reranked] == [0, 2]
+        assert all(isinstance(score, float) for chunk, score in reranked)
 
     except ImportError:
         pytest.skip("Cohere package not installed")