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

# Configuration : Dossier de cache pour éviter les re-téléchargements inutiles
CACHE_DIR = "/tmp/model_cache"
os.makedirs(CACHE_DIR, exist_ok=True)

app = FastAPI()

# Chargement du modèle au démarrage du serveur
print("Chargement du modèle en cours...")
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', cache_folder=CACHE_DIR)
print("Modèle chargé avec succès.")

@app.post("/embed")
async def get_embedding(request: Request):
    try:
        data = await request.json()
        text = data.get("text", "")
        
        if not text:
            return {"error": "Texte manquant"}
        
        # Calcul de l'embedding
        embedding = model.encode(text).tolist()
        return {"embedding": embedding}
        
    except Exception as e:
        return {"error": str(e)}

@app.get("/health")
def health_check():
    return {"status": "ok", "message": "Service embedding opérationnel"}

if __name__ == "__main__":
    # Hugging Face Spaces utilise le port 7860 par défaut
    uvicorn.run(app, host="0.0.0.0", port=7860)