app.py

import gradio as gr
from langchain_community.llms import HuggingFaceHub
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.output_parsers import StrOutputParser
from langchain_huggingface import HuggingFaceEmbeddings
from langchain import hub
from rerankers import Reranker
import os

# Configuración del token de acceso a Hugging Face (si usas modelo privado)
os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Cargar PDF
loader = PyPDFLoader("80dias.pdf")
documents = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20)
splits = splitter.split_documents(documents)

# Crear embeddings
embedding_model = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
vectordb = Chroma.from_documents(splits, embedding=embeddings)

# Modelo LLM desde HuggingFace (usa uno disponible en Spaces)
llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.1", model_kwargs={"temperature": 0.5, "max_new_tokens": 500})
chain = llm | StrOutputParser()

# Reranker
ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type="colbert")

# Función RAG
def rag_chat(message, history):
    # Solo usamos el mensaje del usuario
    query = message

    results = vectordb.similarity_search_with_score(query)
    context = []
    for doc, score in results:
        if score < 7:
            context.append(doc.page_content)
    if not context:
        return "No tengo información suficiente para responder a esa pregunta."

    ranking = ranker.rank(query=query, docs=context)
    best_context = ranking[0].text

    prompt = f"""Contesta a la siguiente pregunta usando solo el contexto que se proporciona:

    Contexto:
    {best_context}

    Pregunta: {query}
    Respuesta:"""

    return llm.invoke(prompt)
    
# Interfaz Gradio
iface = gr.ChatInterface(
    fn=rag_chat,
    title="Chat Julio Verne - RAG",
    description="Pregunta lo que quieras sobre *La vuelta al mundo en 80 días* de Julio Verne.",
    chatbot=gr.Chatbot(type="messages")  # 👈 Esto elimina el warning de formato obsoleto
)
iface.launch()