Update app.py

app.py

@@ -11,6 +11,7 @@ class RerankRequest(BaseModel):
     hits: List[Tuple[int, str]]
 
 class RerankResponse(BaseModel):
+    query: str
     hits: List[Tuple[int, str]]
 
 app = FastAPI()
@@ -21,13 +22,13 @@ def hello_world():
 
 
 @app.get("/rerank")
-def rerank(request: RerankRequest):
+def rerank(request: RerankRequest) -> RerankResponse:
     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]
+        query=request.query,
+        hits = [{"content": hit[1]} for hit in sorted_hits]
     )
     
     reranked_result = reranker.permutation_pipeline(
@@ -37,8 +38,11 @@ def rerank(request: RerankRequest):
         logging=True
     )
 
-    response = [(i, item["content"]) for i, item in enumerate(reranked_result.hits)]
+    reranked_hits = [(i + 1, item["content"]) for i, item in enumerate(reranked_result.hits)]
 
-    return {"reranked": response}
+    return {
+        "query": request.query,
+        "reranked_hits": reranked_hits
+    }