import os
import uvicorn
from fastapi import FastAPI, Request
from sentence_transformers import SentenceTransformer, CrossEncoder

app = FastAPI()
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
reranker_model = CrossEncoder("BAAI/bge-reranker-large")

@app.get("/")
def root():
    return {"message": "Embedding & Reranker API is live!"}

@app.post("/embed")
async def embed(request: Request):
    data = await request.json()
    text = data.get("text", "")
    if not text:
        return {"error": "Missing text"}
    embedding = embedding_model.encode(text).tolist()
    return {"embedding": embedding}

@app.post("/rerank")
async def rerank(request: Request):
    data = await request.json()
    query = data.get("query")
    docs = data.get("documents", [])
    if not query or not docs:
        return {"error": "Missing query or documents"}

    inputs = [[query, doc["content"]] for doc in docs]
    scores = reranker_model.predict(inputs)
    for doc, score in zip(docs, scores):
        doc["rerank_score"] = float(score)
    return {"documents": docs}

@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 10000))
    uvicorn.run(app, host="0.0.0.0", port=port)