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,6 +709,40 @@ 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:
+        docs = []
+        for i, d in enumerate(body.docs):
+            if isinstance(d, str):
+                docs.append(ReRankDocument(text=d, doc_id=i))
+            else:
+                docs.append(
+                    ReRankDocument(
+                        text=d.get("text", ""),
+                        doc_id=d.get("doc_id", i),
+                        metadata=d.get("metadata", {}) or {},
+                    )
+                )
+
+        top_k = body.k
+
+        results = reRankerInstance.rank(query=body.query, docs=docs)
+        items = results.top_k(top_k) if top_k else results
+
+        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",
+            body.query,
+            str(e),
+            traceback.format_exc(),
+        )
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.post("/text")
 async def extract_text_from_file(
     request: Request,

requirements.txt

@@ -37,5 +37,6 @@ python-magic==0.4.27
 python-pptx==1.0.2
 xlrd==2.0.2
 pydantic==2.9.2
-chardet==5.2.0
 rerankers[transformers]==0.6.0
+rerankers[flashrank]==0.6.0
+chardet==5.2.0