Skip to content

Commit 0fdb39f

Browse files
ybalbert001Yuanbo Li
and
Yuanbo Li
authored
Fix: The topk parameter doesn't work in sagemaker rerank tool (#12150)
Co-authored-by: Yuanbo Li <[email protected]>
1 parent dae1b5a commit 0fdb39f

File tree

1 file changed

+3
-5
lines changed

1 file changed

+3
-5
lines changed

api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,7 @@
1010

1111
class SageMakerReRankTool(BuiltinTool):
1212
sagemaker_client: Any = None
13-
sagemaker_endpoint: str | None = None
14-
topk: int | None = None
13+
sagemaker_endpoint: str = None
1514

1615
def _sagemaker_rerank(self, query_input: str, docs: list[str], rerank_endpoint: str):
1716
inputs = [query_input] * len(docs)
@@ -47,8 +46,7 @@ def _invoke(
4746
self.sagemaker_endpoint = tool_parameters.get("sagemaker_endpoint")
4847

4948
line = 2
50-
if not self.topk:
51-
self.topk = tool_parameters.get("topk", 5)
49+
topk = tool_parameters.get("topk", 5)
5250

5351
line = 3
5452
query = tool_parameters.get("query", "")
@@ -75,7 +73,7 @@ def _invoke(
7573
sorted_candidate_docs = sorted(candidate_docs, key=operator.itemgetter("score"), reverse=True)
7674

7775
line = 9
78-
return [self.create_json_message(res) for res in sorted_candidate_docs[: self.topk]]
76+
return [self.create_json_message(res) for res in sorted_candidate_docs[:topk]]
7977

8078
except Exception as e:
8179
return self.create_text_message(f"Exception {str(e)}, line : {line}")

0 commit comments

Comments
 (0)