import streamlit as st
# from streamlit_extras.add_vertical_space import add_vertical_space
import os
import pickle
from PyPDF2 import PdfReader
from langchain.document_loaders import UnstructuredPDFLoader, OnlinePDFLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import OpenAI,AzureOpenAI
from langchain.chains.question_answering import load_qa_chain
from langchain.callbacks import get_openai_callback


# Sidebar contents
with st.sidebar:
    st.title('🤗💬 LLM Chat App')
    st.markdown('''
    ## About
    This app is an LLM-powered chatbot built using:
    - [Streamlit](https://streamlit.io/)
    - [LangChain](https://python.langchain.com/)
    - [OpenAI](https://platform.openai.com/docs/models) LLM model

    ''')
    # add_vertical_space(5)
    st.write('Made by Nick')


def main():
    st.header("智能点餐机器人 💬")

    # # embeddings
    store_name = "coffee"

    if os.path.exists(f"{store_name}.pkl"):
        with open(f"{store_name}.pkl", "rb") as f:
            VectorStore = pickle.load(f)
        st.write('Embeddings Loaded from the Disk')
    else:
        st.write('Reading from prompt ...')
        loader = PyPDFLoader("./咖啡语料.pdf")
        data = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=512,
            chunk_overlap=128,
            length_function=len
        )
        texts = text_splitter.split_documents(data)
        embeddings = OpenAIEmbeddings(chunk_size = 1)
        VectorStore = FAISS.from_texts([t.page_content for t in texts], embedding=embeddings)
        with open(f"{store_name}.pkl", "wb") as f:
            pickle.dump(VectorStore, f)
            

    
    query = st.text_input("Ask questions about Starbucks coffee:")
    

    if query:
        docs = VectorStore.similarity_search(query=query, k=3)

        llm = AzureOpenAI(
            engine="text-davinci-003",
            model_name="text-davinci-003", 
        )
        chain = load_qa_chain(llm=llm, chain_type="stuff")
        with get_openai_callback() as cb:
            response = chain.run(input_documents=docs, question=query)
            print(cb)
        st.write(response)


if __name__ == '__main__':
    main()