import gradio as gr
import os

os.environ["OPENAI_API_KEY"] = "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"

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain.llms import OpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import DirectoryLoader




txt_loader = DirectoryLoader('.\', glob="**/*.txt")
pdf_loader = DirectoryLoader('.\', glob="**/*.pdf")
doc_loader = DirectoryLoader('.\', glob="**/*.docx")
loaders = [pdf_loader, txt_loader, doc_loader]
documents = []

for loader in loaders:
    documents.extend(loader.load())
    
print(f"Total # of documents: {len(documents)}")

text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
documents = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(documents, embeddings)

from langchain.memory import ConversationBufferMemory
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), vectorstore.as_retriever(), memory=memory)

chat_history = []

def submit_callback(user_message):
    default_prompt = " Please format your response in the following way: Each statement should be in a newline . "
    prompt = default_prompt + user_message
    
    # Process user input and generate chatbot response
    response = qa({"question": prompt, "chat_history": chat_history})
    chat_history.append((prompt, response["answer"]))
    return response["answer"]

iface = gr.Interface(
    fn=submit_callback,
    inputs=gr.inputs.Textbox(lines=2, label="Enter your query"),
    outputs=gr.outputs.Textbox(label="Chatbot Response"),
    #outputs=gr.outputs.HTML(label="Chatbot Response"),
    title="LVE Torpedoes Chatbot",
    layout="vertical",
    description="Enter your query to chat with the LVET chatbot",
    examples=[
        ["What are the practice times for each age group ?"],
        ["What are the eligibility criteria for the Mini Torpedoes program?"],
        ["What is the eligibility to participate in the LVET Swim Team?"],
        ["How many volunteer hours are required per family during the swim season?"],
        ["What strokes can swimmers participate in at swim meets?"],
        ["How are swimmers grouped for practice?"],
        ["When do evaluations take place for new swimmers?"],
        ["Who are LVET's Board Members"],
        ["How can I read swim meet results ?"],
        ["How can I contact LVET's Board Members?"],
        ["What is the penalty for not meeting the required volunteer hours?"],
        ["Volunteer Hours?"],
        ["Registration info?"],
        ["How do I sign up for volunteer jobs to fulfill my volunteer hours?"],
        ["Volunteer jobs that do not require certification or prior experience"],
        ["What are the responsibilities of an Age Group Coordinator?"],
        ["How do I commit my swimmer for meets/events?"],
        ["What age groups and races does the LVET Swim Team participate in?"]
    ],
    theme="default"
    
)
iface.launch(share=True)
while True:
    pass