thanhtd/NewAppLangchain

1import os2import tempfile3import shutil4import streamlit as st5from langchain.chat_models import ChatOpenAI6from langchain.document_loaders import UnstructuredFileLoader7from langchain.memory import ConversationBufferMemory8from langchain.memory.chat_message_histories import StreamlitChatMessageHistory9from langchain.embeddings import OpenAIEmbeddings10from langchain.callbacks.base import BaseCallbackHandler11from langchain.chains import ConversationalRetrievalChain12from langchain.text_splitter import RecursiveCharacterTextSplitter13from langchain.vectorstores import Chroma14from langchain.prompts import PromptTemplate,ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate15 16from datetime import datetime17 18st.set_page_config(page_title="ヘルプデスクのチャットボット", page_icon="🤖")19st.title("ヘルプデスクのチャットボット")20 21persist_directory = "db"22 23@st.cache_resource(ttl="1h")24def configure_retriever(uploaded_files, persist_directory, api_key):25    # Read documents26    docs = []27    temp_dir = tempfile.TemporaryDirectory()28    for file in uploaded_files:29        print(file)30 31        temp_filepath = os.path.join(temp_dir.name, file.name)32        with open(temp_filepath, "wb") as f:33            f.write(file.getvalue())34        loader = UnstructuredFileLoader(temp_filepath, unstructured_kwargs={'autodetect_encoding': True})35        docs.extend(loader.load())36    # Remove temp file37    if (len(docs) > 0):38        shutil.rmtree(temp_dir.name)39 40 41    # Split documents42    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)43    split_docs = text_splitter.split_documents(docs)44 45    46    # Create embeddings and store in vectordb47    embeddings = OpenAIEmbeddings(openai_api_key=api_key)48    #embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")49    vectordb = Chroma.from_documents(documents=split_docs, 50                                embedding=embeddings,51                                persist_directory=persist_directory)52 53    vectordb.persist()54    vectordb = None55 56    # Define retriever57    # retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={"k": 2, "fetch_k": 4})58 59    # return retriever60 61 62class StreamHandler(BaseCallbackHandler):63    def __init__(self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""):64        self.container = container65        self.text = initial_text66        self.run_id_ignore_token = None67 68    def on_llm_start(self, serialized: dict, prompts: list, **kwargs):69        # Workaround to prevent showing the rephrased question as output70        if prompts[0].startswith("Human"):71            self.run_id_ignore_token = kwargs.get("run_id")72 73    def on_llm_new_token(self, token: str, **kwargs) -> None:74        if self.run_id_ignore_token == kwargs.get("run_id", False):75            return76        77        self.text += token78        self.container.markdown(self.text)79 80 81class PrintRetrievalHandler(BaseCallbackHandler):82    def __init__(self, container):83        self.status = container.status("**コンテキストの取得**")84 85    def on_retriever_start(self, serialized: dict, query: str, **kwargs):86        self.status.write(f"**Question:** {query}")87        self.status.update(label=f"**コンテキストの取得:** {query}")88 89    def on_retriever_end(self, documents, **kwargs):90        for idx, doc in enumerate(documents):91            source = os.path.basename(doc.metadata["source"])92            self.status.write(f"**Document {idx} from {source}**")93            self.status.markdown(doc.page_content)94        self.status.update(state="complete")95 96 97openai_api_key = st.sidebar.text_input("OpenAI API キー", type="password")98if not openai_api_key:99    st.info("OpenAI API キーを追加してください。")100    st.stop()101 102accepted_file_types = ["doc", "docx", "pdf", "txt", "xlsx", "xls"]103 104uploaded_files = st.sidebar.file_uploader(105    label="ファイルをアップロードする", type=accepted_file_types, accept_multiple_files=True106)107# Comment because docuemnt may already uploaded 108# if not uploaded_files:109#     st.info("書類をアップロードしてください。")110#     st.stop()111if uploaded_files:112    configure_retriever(uploaded_files, persist_directory, openai_api_key)113 114# Setup memory for contextual conversation115msgs = StreamlitChatMessageHistory()116memory = ConversationBufferMemory(memory_key="chat_history", 117                                  chat_memory=msgs, 118                                  return_messages=True)119 120print("Message====", msgs.messages)121 122# Setup LLM and QA chain123llm = ChatOpenAI(124    model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, temperature=0, streaming=True125)126 127embedding = OpenAIEmbeddings(openai_api_key=openai_api_key)128 129vectordb = Chroma(persist_directory=persist_directory, 130                embedding_function=embedding,131                   )132 133prompt_template="""134    You are a chatbot having a conversation with a human.135    Given the following extracted parts of a long document and a question, create a final answer in Japanese.136    If you don't know the answer, just say that you don't know, don't try to make up an answer.137    138    {context}        139    140    Question: {question}141    """142PROMPT = ChatPromptTemplate.from_messages(143    [144    SystemMessagePromptTemplate.from_template(145        """146        You are a chatbot having a conversation with a human.147        Given the following extracted parts of a long document and a question, create a final answer in Japanese.148        If you don't know the answer, just say that you don't know, don't try to make up an answer.149        150        {context}  151        """),152    HumanMessagePromptTemplate.from_template('{question}')153    ]154)155#Custom prompt156chain_kwargs = {"prompt":PROMPT}157 158qa_chain = ConversationalRetrievalChain.from_llm(159    llm, retriever=vectordb.as_retriever(search_kwargs={"k": 2}), 160    memory=memory, verbose=True,161    combine_docs_chain_kwargs=chain_kwargs162)163 164 165if len(msgs.messages) == 0 or st.sidebar.button("メッセージ履歴をクリアする"):166    msgs.clear()167    msgs.add_ai_message("ご用件をお伺いしてもよろしいですか？")168 169if st.sidebar.button("DBリセット"):170    # zip db directory171    shutil.make_archive(persist_directory, 'zip', os.getcwd()+"/"+persist_directory)172    # To cleanup, you can delete the collection173    vectordb.delete_collection()174    vectordb.persist()175 176avatars = {"human": "user", "ai": "assistant"}177for msg in msgs.messages:178    st.chat_message(avatars[msg.type]).write(msg.content)179 180if user_query := st.chat_input(placeholder="何でも聞いてください！"):181    st.chat_message("user").write(user_query)182 183    with st.chat_message("assistant"):184        retrieval_handler = PrintRetrievalHandler(st.container())185        stream_handler = StreamHandler(st.empty())186        response = qa_chain.run(user_query, callbacks=[retrieval_handler, stream_handler])