fix(rag): Fix rerank bug of EmbeddingRetriever (#1499)

Co-authored-by: dong <[email protected]>
eosphoros-ai · May 9, 2024 · 87a67cc · 87a67cc
1 parent b38305b
commit 87a67cc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dbgpt/rag/retriever/embedding.py b/dbgpt/rag/retriever/embedding.py
@@ -120,7 +120,7 @@ def _retrieve_with_score(
         new_candidates_with_score = cast(
             List[Chunk], reduce(lambda x, y: x + y, candidates_with_score)
         )
-        new_candidates_with_score = self._rerank.rank(new_candidates_with_score)
+        new_candidates_with_score = self._rerank.rank(new_candidates_with_score, query)
         return new_candidates_with_score
 
     async def _aretrieve(
@@ -207,7 +207,9 @@ async def _aretrieve_with_score(
                 "rerank_cls": self._rerank.__class__.__name__,
             },
         ):
-            new_candidates_with_score = self._rerank.rank(new_candidates_with_score)
+            new_candidates_with_score = self._rerank.rank(
+                new_candidates_with_score, query
+            )
             return new_candidates_with_score
 
     async def _similarity_search(