import gradio as gr
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chat_models import ChatOpenAI

from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.retrievers import ContextualCompressionRetriever
from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate



from langchain.vectorstores import Chroma


chat = ChatOpenAI()

embedding_function = HuggingFaceEmbeddings(model_name = "BAAI/bge-large-en-v1.5",model_kwargs={'device': 'cpu'},encode_kwargs={"normalize_embeddings": True})

def add_docs(path):

    loader = PyPDFLoader(file_path=path)
    docs = loader.load_and_split(text_splitter=RecursiveCharacterTextSplitter(chunk_size = 500,
                                                                                chunk_overlap = 100,
                                                                                length_function = len,
                                                                                is_separator_regex=False))
    model_vectorstore = Chroma
    db = model_vectorstore.from_documents(documents=docs,embedding= embedding_function, persist_directory="output/general_knowledge")
    return db


def answer_query(message, chat_history):
    base_compressor = LLMChainExtractor.from_llm(chat)
    db = Chroma(persist_directory = "output/general_knowledge", embedding_function=embedding_function)
    base_retriever = db.as_retriever()
    mq_retriever = MultiQueryRetriever.from_llm(retriever = base_retriever, llm=chat)
    compression_retriever = ContextualCompressionRetriever(base_compressor=base_compressor, base_retriever=mq_retriever)

    matched_docs = compression_retriever.get_relevant_documents(query = message)

    context = ""

    for doc in matched_docs:
        page_content = doc.page_content
        context+=page_content
        context += "\n\n"
    template = """
    Answer the following question only by using the context given below in the triple backticks, do not use any other information to answer the question.
    If you can't answer the given question with the given context, you can return an emtpy string ('')

    Context: ```{context}```
    ----------------------------
    Question: {query}
    ----------------------------
    Answer: """
    
    human_message_prompt = HumanMessagePromptTemplate.from_template(template=template)
    chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])
    prompt = chat_prompt.format_prompt(query = message, context = context)
    response = chat(messages=prompt.to_messages()).content

    chat_history.append((message,response))
    return "", chat_history



with gr.Blocks() as demo:
    gr.HTML("<h1 align = 'center'>Smart Assistant</h1>")
    gr.HTML("<h2 align = 'center'>Upload any PDF and ask your questions.</h2>")

    with gr.Row():

        upload_files = gr.File(label = 'Upload a PDF',file_types=['.pdf'],file_count='single')

    chatbot = gr.Chatbot()
    msg = gr.Textbox(label = "Enter your question here")
    upload_files.upload(add_docs,upload_files)
    msg.submit(answer_query,[msg,chatbot],[msg,chatbot])


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