from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.chains import LLMChain, ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
import gradio as gr

# LLM
llm = ChatOpenAI()

# load data
loader = TextLoader('data/codes.txt')
data = loader.load()

# Split and store into vector
from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
all_splits = text_splitter.split_documents(data)
# persist data into SQLite DB
vectorstore = Chroma.from_documents(documents=all_splits, embedding=OpenAIEmbeddings(), persist_directory='persist')

# Prompt 
prompt = ChatPromptTemplate(
    messages=[
        SystemMessagePromptTemplate.from_template(
            "You are a nice chatbot having a conversation with a human about UF student government and its 800 codes which are about the budget and student org funding. Please reference context first when trying to respond to queries and try to cite sources."
        ),
        # The `variable_name` here is what must align with memory
        MessagesPlaceholder(variable_name="chat_history"),
        HumanMessagePromptTemplate.from_template("{question}"),
    ]
)

# Notice that we `return_messages=True` to fit into the MessagesPlaceholder
# Notice that `"chat_history"` aligns with the MessagesPlaceholder name
retriever = vectorstore.as_retriever()
memory = ConversationBufferMemory(memory_key="chat_history",return_messages=True)
qa = ConversationalRetrievalChain.from_llm(llm, retriever=retriever, memory=memory)

def queryAI(query: str):
    return qa({"question": query})["answer"]

website_description = "This is an experimental bot to help students get help with the 800 codes which govern the student activity fees and service budget and student org funding. This is intended to be a helpful resource but its accuracy has not been verfied yet!"

gui = gr.Interface(fn=queryAI, inputs="text", outputs="text", title="800 Codes Bot - EXPERIMENTAL", description=website_description, flagging_options=["Flag Correct/Helpful", "Flag Incorrect"])
gui.launch()