import pandas as pd

# Carga los datos de entrenamiento
df = pd.read_csv('./medical_data.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]
    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.1-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="mecanica_automotriz",
    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 = ("""Tu eres un experto en mecánica automotriz, puedes responder preguntas sobre coches y motores.
    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 = [
    "Mi coche no arranca, ¿qué puedo hacer?",
    "¿Cómo puedo cambiar el aceite de mi coche?"
]
description = "Aplicación de IA en tiempo real para responder preguntas sobre mecánica automotriz"
title = "Experto en Mecánica Automotriz :)"
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()