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

Co-authored-by: dong <dongzhancai@iie2.com>
Hopshine · Sep 10, 2024 · c4d47ae · c4d47ae
1 parent 75c3eb1
commit c4d47ae
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(