Langchain과 Gradio를 이용해 수면 상담 Chatbot 만들기 (진행 중)

def build_db():
    global cbti_manual_db

    # Load the text file and build a vector store with FAISS
    loader = PyPDFLoader("CBT-I.pdf")
    document = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0, 
                                                    separators=["\n\n",
                                                                "\n",
                                                                " ",
                                                                ".",
                                                                ",",
                                                                "",])
    docs = text_splitter.split_documents(document)

    embeddings = OpenAIEmbeddings()

    cbti_manual_db = FAISS.from_documents(docs, embeddings)

def build_db():
    global cbti_manual_db
    cbti_db_file_name = "cbti_faiss_index"
    
    embeddings = OpenAIEmbeddings()

    # if there's a saved faiss db load it. or make one and save it.
    try:
        cbti_manual_db = FAISS.load_local(cbti_db_file_name, embeddings=embeddings, allow_dangerous_deserialization=True)
    except Exception as e:
        print(f"Failed to load FAISS index: {e}.")

        # Load the text file and build a vector store with FAISS
        loader = PyPDFLoader("Cognitive Behavioral Treatment of Insomnia A Session-by-Session Guide.pdf")
        document = loader.load()

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size = 1000, 
            chunk_overlap = 200,
            separators=["\n\n", "\n", " ", ". ", ", ", "?", "! "]
        )
        docs = text_splitter.split_documents(document)                

        cbti_manual_db = FAISS.from_documents(docs, embeddings)
        cbti_manual_db.save_local(cbti_db_file_name)

def predict(message, history):
    global cbti_manual_db

    history_langchain_format = []

    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))   

    history_langchain_format.append(HumanMessage(content=message))

    # Translate the message into English to enhance RAG
    en_prompt = ChatPromptTemplate.from_messages([
        ("system", "Translate the user message into English"),
        ("user", "{message}")
    ])

    en_chain = en_prompt | llm | parser
    en_message = en_chain.invoke({"message": message})

    # define retriever with similarity search
    retriever = cbti_manual_db.as_retriever(search_type="similarity", search_kwargs={"k": 6})
    prompt = hub.pull("rlm/rag-prompt")

    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
    )

    gpt_response = rag_chain.invoke(en_message)
    
    return gpt_response.content


build_db()
gr.ChatInterface(predict).launch()

def translate_to_eng(org_msg):
     # Detect the language of original message
    detect_lang_prompt = ChatPromptTemplate.from_messages([
        ("system", "Detect the language used for the message. Answer with the language name in English"),
        ("user", "{message}")
    ])

    detect_lang_chain = detect_lang_prompt | llm | parser
    org_lang = detect_lang_chain.invoke({"message": org_msg})

    # Translate the message into English to enhance RAG
    if org_lang != "English":        
        en_prompt = ChatPromptTemplate.from_messages([
            ("system", "Translate the user message into English"),
            ("user", "{message}")
        ])

        en_chain = en_prompt | llm | parser
        en_message = en_chain.invoke({"message": org_msg})
    else:
        en_message = org_msg   
    
    return en_message, org_lang

def translate_to_org(msg, org_lang):
    if org_lang == "English":
        return msg

    # Translate the message into org_lang
    translate_prompt = ChatPromptTemplate.from_messages([
        ("system", "Translate the user message into {org_lang} with your expertise. Use a natural and easy words and expressions"),
        ("user", "{message}")
    ])

    translate_chain = translate_prompt | llm | parser
    result_message = translate_chain.invoke({"org_lang": org_lang, "message": msg})
    
    return result_message

def predict(message, history):
    global cbti_manual_db

    history_langchain_format = []

    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))       

    en_message, org_lang = translate_to_eng(message)
    #org_lang = "English"

    # define retriever with similarity search
    retriever = cbti_manual_db.as_retriever(search_type="similarity", search_kwargs={"k": 6})


    # question + history => new question
    contextualize_q_system_prompt = (
        "Given a chat history and the latest user question "
        "which might reference context in the chat history, "
        "formulate a standalone question which can be understood "
        "without the chat history. Do NOT answer the question, "
        "just reformulate it if needed and otherwise return it as is."
    )

    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", contextualize_q_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )
    history_aware_retriever = create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )

    # Q&A assistant which takes the reformulated question and answer for it
    system_prompt = (
        "You are an assistant for question-answering tasks in the field of sleep medicine. "
        "You talk and behave as a professional sleep consultant and an expert CBT-I practitioner"
        "Use easy words to help the user to improve sleep quality."
        "Use the following pieces of retrieved context to answer "
        "the question. If you don't know the answer, say that you "
        "don't know. Use three sentences maximum and keep the "
        "answer concise."
        "\n\n"
        "{context}"
    )

    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )

    question_answer_chain = qa_prompt | llm

    # build RAG chain
    rag_chain = create_retrieval_chain(history_aware_retriever | format_docs, question_answer_chain)

    #rag_chain = create_retrieval_chain()
    # rag_chain = create_
    '''
    rag_chain = (
        {"context": retriever | format_docs, "input": RunnablePassthrough()}
        | prompt
        | llm
    )
    '''

    try:
        gpt_response = rag_chain.invoke({"input": en_message, "chat_history": history_langchain_format})
    except Exception as e:
        print("error during RAG chain invocation:", e)
        raise e
    
    # add message to history_langchain_format after RAG gives the answer
    history_langchain_format.append(HumanMessage(content=message))    

    # Translate the answer into the original language
    response = translate_to_org(gpt_response['answer'].content, org_lang)
    
    return response