refac: rerank

backend/open_webui/retrieval/utils.py

@@ -168,7 +168,7 @@ def query_doc_with_hybrid_search(
         ):
             log.warning(f"query_doc_with_hybrid_search:no_docs {collection_name}")
             return {"documents": [], "metadatas": [], "distances": []}
-        
+
         # Now safely check the documents content after confirming attributes exist
         if (
             not collection_result.documents
@@ -516,11 +516,13 @@ def get_reranking_function(reranking_engine, reranking_model, reranking_function
     if reranking_function is None:
         return None
     if reranking_engine == "external":
-        return lambda sentences, user=None: reranking_function.predict(
-            sentences, user=user
+        return lambda query, documents, user=None: reranking_function.predict(
+            [(query, doc.page_content) for doc in documents], user=user
         )
     else:
-        return lambda sentences, user=None: reranking_function.predict(sentences)
+        return lambda query, documents, user=None: reranking_function.predict(
+            [(query, doc.page_content) for doc in documents]
+        )
 
 
 def get_sources_from_items(
@@ -1064,9 +1066,7 @@ class RerankCompressor(BaseDocumentCompressor):
 
         scores = None
         if reranking:
-            scores = self.reranking_function(
-                [(query, doc.page_content) for doc in documents]
-            )
+            scores = self.reranking_function(query, documents)
         else:
             from sentence_transformers import util