Use parameter for reranker

comps/cores/proto/docarray.py

@@ -50,6 +50,7 @@ class EmbedDoc1024(BaseDoc):
 class SearchedDoc(BaseDoc):
     retrieved_docs: DocList[TextDoc]
     initial_query: str
+    top_n: int = 1
 
     class Config:
         json_encoders = {np.ndarray: lambda x: x.tolist()}

comps/reranks/README.md

@@ -100,3 +100,12 @@ curl http://localhost:8000/v1/reranking \
   -d '{"initial_query":"What is Deep Learning?", "retrieved_docs": [{"text":"Deep Learning is not..."}, {"text":"Deep learning is..."}]}' \
   -H 'Content-Type: application/json'
 ```
+
+You can add the parameter `top_n` to specify the return number of the reranker model, default value is 1.
+
+```bash
+curl http://localhost:8000/v1/reranking \
+  -X POST \
+  -d '{"initial_query":"What is Deep Learning?", "retrieved_docs": [{"text":"Deep Learning is not..."}, {"text":"Deep learning is..."}], "top_n":2}' \
+  -H 'Content-Type: application/json'
+```

comps/reranks/langchain/reranking_tei_xeon.py

@@ -1,6 +1,7 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+import heapq
 import json
 import os
 import time
@@ -39,14 +40,16 @@ def reranking(input: SearchedDoc) -> LLMParamsDoc:
     headers = {"Content-Type": "application/json"}
     response = requests.post(url, data=json.dumps(data), headers=headers)
     response_data = response.json()
-    best_response = max(response_data, key=lambda response: response["score"])
+    best_response_list = heapq.nlargest(input.top_n, response_data, key=lambda x: x["score"])
     template = """Answer the question based only on the following context:
     {context}
     Question: {question}
     """
     prompt = ChatPromptTemplate.from_template(template)
-    doc = input.retrieved_docs[best_response["index"]]
-    final_prompt = prompt.format(context=doc.text, question=input.initial_query)
+    context_str = ""
+    for best_response in best_response_list:
+        context_str = context_str + " " + input.retrieved_docs[best_response["index"]].text
+    final_prompt = prompt.format(context=context_str, question=input.initial_query)
     statistics_dict["opea_service@reranking_tgi_gaudi"].append_latency(time.time() - start, None)
     return LLMParamsDoc(query=final_prompt.strip())