from .utils import catch_error_and_return
from dotenv import load_dotenv
from fastmcp import FastMCP
from typing import List
import requests
import os
load_dotenv()
reranker_mcp = FastMCP(name="Reranker server")
@reranker_mcp.tool
@catch_error_and_return
def rerank_documents(query: str, documents: List[str]):
"""
Rerank texts
:param query: user query
:param documents: documents for rerank
:return: reranked list of documents
"""
url = os.getenv("RERANK_URL")
rerank_model = os.getenv("RERANK_MODEL")
data = {"model": rerank_model, "query": query, "documents": documents}
result = requests.post(url, json=data).json()
return result
# TODO: rerank_with_scores
@reranker_mcp.tool
def health_check():
"""
health check for server
:return: server status
"""
url = os.getenv("RERANK_HEALTH_URL")
result = requests.get(url)
if "unix" in result.text:
return {"status": "ok"}
else:
return {"status": result.text}
if __name__ == "__main__":
reranker_mcp.run(transport="stdio")
