fix: xinference reranker return_documents (#6888)

langgenius · Aug 1, 2024 · cc4785f · cc4785f
1 parent 093f902
commit cc4785f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/api/core/model_runtime/model_providers/xinference/rerank/rerank.py b/api/core/model_runtime/model_providers/xinference/rerank/rerank.py
@@ -57,6 +57,7 @@ def _invoke(self, model: str, credentials: dict,
                 documents=docs,
                 query=query,
                 top_n=top_n,
+                return_documents=True
             )
         except RuntimeError as e:
             raise InvokeServerUnavailableError(str(e))
@@ -66,7 +67,7 @@ def _invoke(self, model: str, credentials: dict,
         for idx, result in enumerate(response['results']):
             # format document
             index = result['index']
-            page_content = result['document']
+            page_content = result['document'] if isinstance(result['document'], str) else result['document']['text']
             rerank_document = RerankDocument(
                 index=index,
                 text=page_content,

diff --git a/api/tests/integration_tests/model_runtime/__mock/xinference.py b/api/tests/integration_tests/model_runtime/__mock/xinference.py
@@ -106,7 +106,7 @@ def get(self: Session, url: str, **kwargs):
     def _check_cluster_authenticated(self):
         self._cluster_authed = True
 
-    def rerank(self: RESTfulRerankModelHandle, documents: list[str], query: str, top_n: int) -> dict:
+    def rerank(self: RESTfulRerankModelHandle, documents: list[str], query: str, top_n: int, return_documents: bool) -> dict:
         # check if self._model_uid is a valid uuid
         if not re.match(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', self._model_uid) and \
             self._model_uid != 'rerank':