import pandas as pd
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
import gradio as gr

# Carga los datos de entrenamiento
df = pd.read_csv('./botreformasconstrucciones.csv')

# Crea un arreglo con los contextos
context_data = []
for i in range(len(df)):
  context = ""
  for j in range(3):
    context += df.columns[j]
    context += ": "
    context += df.iloc[i, j]  # Cambia esto
    context += " "
  context_data.append(context)

# Importa las bibliotecas necesarias
import os
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma

# Obtiene la clave de API de Groq
groq_key = os.environ.get('groq_api_keys')

# Crea un objeto ChatGroq con el modelo de lenguaje
llm = ChatGroq(model="llama-3.3-70b-versatile", api_key=groq_key)

# Crea un objeto HuggingFaceEmbeddings con el modelo de embeddings
embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

# Crea un objeto Chroma con el nombre de la colección
vectorstore = Chroma(
    collection_name="unipapel_papeleria",
    embedding_function=embed_model,
)

# Agrega los textos a la colección
vectorstore.add_texts(context_data)

# Crea un objeto retriever con la colección
retriever = vectorstore.as_retriever()

# Crea un objeto PromptTemplate con el prompt
template = ("""Eres un asistente experto de unipapel.es, especializado en temas relacionados con la papelería. Tu función es ofrecer orientación a los clientes que realicen preguntas. Si no encuentras una respuesta, invítalos a contactar con nosotros. Puedes proporcionar información sobre nuestras categorías de productos según lo que te soliciten. La URL de la tienda es www.unipapel.es. A continuación, algunas de nuestras categorías:

Accesorios para pizarras

Acuarelas y lápices acuarelables

Archivadores de palanca

Blocs de dibujo

Bolígrafos

Carpetas clasificadoras

Cuadernos de caligrafía

Papelería escolar

Pegamentos y adhesivos

Pizarras blancas

    Context: {context}
    Question: {question}
    Answer:""")

# Crea un objeto rag_prompt con el prompt
rag_prompt = PromptTemplate.from_template(template)

# Crea un objeto StrOutputParser para parsear la salida
from langchain_core.output_parsers import StrOutputParser

# Crea un objeto RunnablePassthrough para ejecutar el modelo
from langchain_core.runnables import RunnablePassthrough

# Crea un objeto rag_chain con el modelo y el prompt
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | StrOutputParser()
)

# Importa la biblioteca Gradio
import gradio as gr

# Crea una función para procesar la entrada del usuario
def rag_memory_stream(message, history):
    partial_text = ""
    for new_text in rag_chain.stream(message):
        partial_text += new_text
        yield partial_text

# Crea un objeto Gradio con la función y el título
examples = [
    "¿Qué tipo de bolígrafos son ideales para tomar apuntes en clase?",
    "¿Cuáles son los mejores tipos de papel para imprimir documentos profesionales?",
    "¿Qué materiales de oficina son esenciales para un espacio de trabajo productivo?"
]
description = "Chatbot con Unipapel"
title = "Experto en papeleria y todo lo referente"
demo = gr.ChatInterface(fn=rag_memory_stream,
                        type="messages",
                        title=title,
                        description=description,
                        fill_height=True,
                        examples=examples,
                        theme="glass",
)

# Lanza la aplicación
if __name__ == "__main__":
    demo.launch()