After linting.

Author: kalaspuffar

app/models.py

@@ -47,5 +47,4 @@ class QueryMultipleBody(BaseModel):
 class QueryMultipleDocs(BaseModel):
     query: str
     docs: List[str]
-    config: dict
     k: int = 4

app/routes/document_routes.py

@@ -42,6 +42,11 @@
 from app.utils.health import is_health_ok
 
 router = APIRouter()
+reRankerInstance = Reranker(
+    model_name=os.getenv("SIMPLE_RERANKER_MODEL_NAME"),
+    model_type=os.getenv("SIMPLE_RERANKER_MODEL_TYPE"),
+    lang=os.getenv("SIMPLE_RERANKER_LANG"),
+)
 
 
 def get_user_id(request: Request, entity_id: str = None) -> str:
@@ -704,17 +709,10 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody
         )
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @router.post("/rerank")
 async def rerank_documents_by_query(request: Request, body: QueryMultipleDocs):
     try:
-        rk = Reranker(
-            body.config.get("model_name", "flashrank"),
-            model_type=body.config.get("model_type"),
-            lang=body.config.get("lang"),
-            api_provider=body.config.get("api_provider"),
-            api_key=body.config.get("api_key"),
-        )
-
         docs = []
         for i, d in enumerate(body.docs):
             if isinstance(d, str):
@@ -730,10 +728,12 @@ async def rerank_documents_by_query(request: Request, body: QueryMultipleDocs):
 
         top_k = body.k
 
-        results = rk.rank(query=body.query, docs=docs)
+        results = reRankerInstance.rank(query=body.query, docs=docs)
         items = results.top_k(top_k) if top_k else results
 
-        return [[getattr(r.document, "text", None), r.score] for r in items]
+        return [
+            {"text": getattr(r.document, "text", None), "score": r.score} for r in items
+        ]
     except Exception as e:
         logger.error(
             "Error in reranking documents | Query: %s | Error: %s | Traceback: %s",
@@ -743,6 +743,7 @@ async def rerank_documents_by_query(request: Request, body: QueryMultipleDocs):
         )
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @router.post("/text")
 async def extract_text_from_file(
     request: Request,

requirements.txt

@@ -38,4 +38,5 @@ python-pptx==1.0.2
 xlrd==2.0.2
 pydantic==2.9.2
 rerankers[transformers]==0.6.0
+rerankers[flashrank]==0.6.0
 chardet==5.2.0