import pandas as pd


df = pd.read_csv('./medical_data.csv')

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)


import os

# Get the secret key from the environment
groq_key = os.environ.get('groq_api_keys')

## LLM used for RAG
from langchain_groq import ChatGroq

llm = ChatGroq(model="llama-3.1-70b-versatile",api_key=groq_key)

## Embedding model!
from langchain_huggingface import HuggingFaceEmbeddings
embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

# create vector store!
from langchain_chroma import Chroma

vectorstore = Chroma(
    collection_name="medical_dataset_store",
    embedding_function=embed_model,
)

# add data to vector nstore
vectorstore.add_texts(context_data)

retriever = vectorstore.as_retriever()

from langchain_core.prompts import PromptTemplate

template = ("""tu eres un experto en mecanica automotriz, puedes hablar de mas cosas, cuando te pregunten por algo relacionado a los vehiculos o motores
debes responder pidiendo la marva y modelo de auto, luego pediras la fecha, y pediras que te digan los sintomas, tu les daras soluciones.

    Context: {context}

    Question: {question}

    Answer:""")

rag_prompt = PromptTemplate.from_template(template)

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | StrOutputParser()
)

import gradio as gr

def rag_memory_stream(message, history):
    partial_text = ""
    for new_text in rag_chain.stream(message):
        partial_text += new_text
        yield partial_text

examples = [
    "I feel dizzy", 
    "what is the possible sickness for fatigue?"
]

description = "Real-time AI App with Groq API and LangChain to Answer medical questions"


title = "Medical Expert :) Try me!"
demo = gr.ChatInterface(fn=rag_memory_stream,
                        type="messages",
                        title=title,
                        description=description,
                        fill_height=True,
                        examples=examples,
                        theme="glass",
)


if __name__ == "__main__":
    demo.launch()