from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer, util
import numpy as np
import json
import uvicorn
import os

os.environ["HF_HOME"] = "/tmp"
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp"

app = FastAPI()

# Cargar el model
modelo = SentenceTransformer("Francogv/Modelo-Embedding")

# Cargar la base de embeddings
with open("base_semantica.json", "r", encoding="utf-8") as f:
    base = json.load(f)

# Pydantic schema para entrada
class PreguntaInput(BaseModel):
    pregunta: str
    top_k: int = 3

# Función de búsqueda
def buscar_semanticamente(pregunta, top_k=3):
    emb_pregunta = modelo.encode(pregunta)
    resultados = []
    for item in base:
        emb_item = np.array(item["embedding"], dtype=np.float32)
        score = util.cos_sim(emb_pregunta, emb_item).item()
        resultados.append((score, item))
    resultados.sort(reverse=True, key=lambda x: x[0])
    return resultados[:top_k]

# Endpoint principal
@app.post("/buscar")
async def buscar(input: PreguntaInput):
    resultados = buscar_semanticamente(input.pregunta, input.top_k)
    return {
        "pregunta": input.pregunta,
        "resultados": [
            {
                "score": round(score, 4),
                "titulo": item["titulo"],
                "url": item["url"],
                "texto": item["texto"]
            } for score, item in resultados
        ]
    }

# Para correr en Hugging Face
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)