From ca3c5d9fb91d7d815f594caa7a339b0282abf1a9 Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Tue, 3 Dec 2024 13:07:50 +0800 Subject: [PATCH 1/2] fix: better wenxin rerank handler, close #11252 Signed-off-by: yihong0618 --- .../model_providers/wenxin/rerank/rerank.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py b/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py index 9e6a7dd99e10a5..ef63fdeb7b7512 100644 --- a/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py @@ -17,7 +17,13 @@ class WenxinRerank(_CommonWenxin): def rerank(self, model: str, query: str, docs: list[str], top_n: Optional[int] = None): access_token = self._get_access_token() url = f"{self.api_bases[model]}?access_token={access_token}" - + # For issue #11252 + # for wenxin Rerank model top_n length should be equal or less than docs length + if top_n is not None and top_n > len(docs): + top_n = len(docs) + # for wenxin Rerank model, query should not be an empty string + if query == "": + query = " " # FIXME: this is a workaround for wenxin rerank model for better user experience. try: response = httpx.post( url, @@ -25,7 +31,11 @@ def rerank(self, model: str, query: str, docs: list[str], top_n: Optional[int] = headers={"Content-Type": "application/json"}, ) response.raise_for_status() - return response.json() + data = response.json() + # wenxin error handling + if "error_code" in data: + raise InternalServerError(data["error_msg"]) + return data except httpx.HTTPStatusError as e: raise InternalServerError(str(e)) @@ -69,6 +79,7 @@ def _invoke( results = wenxin_rerank.rerank(model, query, docs, top_n) rerank_documents = [] + assert "results" in results, "results key not found in response" for result in results["results"]: index = result["index"] if "document" in result: From be6d28e6a33251c3500eb0696fb3b0e43d30bd5b Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Tue, 3 Dec 2024 13:51:56 +0800 Subject: [PATCH 2/2] fix: adress comments Signed-off-by: yihong0618 --- .../model_runtime/model_providers/wenxin/rerank/rerank.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py b/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py index ef63fdeb7b7512..1c7b9490185337 100644 --- a/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/wenxin/rerank/rerank.py @@ -79,7 +79,9 @@ def _invoke( results = wenxin_rerank.rerank(model, query, docs, top_n) rerank_documents = [] - assert "results" in results, "results key not found in response" + if "results" not in results: + raise ValueError("results key not found in response") + for result in results["results"]: index = result["index"] if "document" in result: