Update app.py to support reranking using SweRankLLM-Small

app.py

@@ -1,7 +1,44 @@
 from fastapi import FastAPI
+from reranker import RankLLM, RankListwiseOSLLM, Result, RankingExecInfo
+from pydantic import BaseModel
+from typing import Optional, List, Tuple
+
+# load RankListwiseOSLLM
+reranker = RankListwiseOSLLM("Salesforce/SweRankLLM-small")
+
+class RerankRequest(BaseModel):
+    query: str
+    hits: List[Tuple[int, str]]
+
+class RerankResponse(BaseModel):
+    hits: List[Tuple[int, str]]
 
 app = FastAPI()
 
 @app.get("/")
 def hello_world():
-    return {"msg": "Success"}
+    return {"msg": "Success"}
+
+
+@app.get("/rerank")
+def rerank(request: RerankRequest):
+    hits = request.hits
+    sorted_hits = sorted(hits, key=lambda x: x[0]) # sort hits again for safety
+    
+    result = Result(
+        query=request["query"],
+        hits = [{"content": hit} for hit in sorted_hits]
+    )
+    
+    reranked_result = reranker.permutation_pipeline(
+        result,
+        1, 
+        len(hits),
+        logging=True
+    )
+
+    response = [(i, item["content"]) for i, item in enumerate(reranked_result.hits)]
+
+    return {"reranked": response}
+
+