add chatbot

app.py

@@ -1,11 +1,15 @@
+from typing import List
+
 import gradio
 import gradio as gr
 import spacy
-from langchain.chains import RetrievalQA
+from langchain.chains import RetrievalQA, ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import SpacyTextSplitter
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 from langchain_community.vectorstores import Chroma
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 from langchain_core.prompts import PromptTemplate
 from langchain_google_genai import ChatGoogleGenerativeAI
 
@@ -19,6 +23,14 @@ Helpful Answer:"""
 QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
 
 
+def convert_chat_history_to_messages(chat_history) -> List[BaseMessage]:
+    result = []
+    for human, ai in chat_history:
+        result.append(HumanMessage(content=human))
+        result.append(AIMessage(content=ai))
+    return result
+
+
 class RAGDemo(object):
     def __init__(self):
         self.embedding = None
@@ -79,8 +91,24 @@ class RAGDemo(object):
         resp = basic_qa.invoke(input_text)
         return resp['result']
 
+    def _chat_qa(self, message, chat_history):
+        memory = ConversationBufferMemory(
+            chat_memory=convert_chat_history_to_messages(chat_history),
+            memory_key="chat_history",
+            return_messages=True,
+        )
+        qa = ConversationalRetrievalChain.from_llm(
+            self.chat_model,
+            retriever=self.vector_db.as_retriever(),
+            memory=memory,
+            verbose=True,
+        )
+        resp = qa.invoke(message)
+        chat_history.append((message, resp['answer']))
+        return "", chat_history
+
     def __call__(self):
-        with gr.Blocks() as demo:
+        with gr.Blocks(title="🔥 RAG Demo") as demo:
             gr.Markdown("# RAG Demo\n\nbase on the [RAG learning note](https://www.jianshu.com/p/9792f1e6c3f9) and "
                         "[rag-practice](https://github.com/hiwei93/rag-practice/tree/main)")
             with gr.Tab("Settings"):
@@ -109,6 +137,10 @@ class RAGDemo(object):
                         submit_btn = gr.Button("submit")
                     with gr.Column():
                         output = gr.TextArea(label="answer")
+            with gr.Tab("Chat RAG"):
+                chatbot = gr.Chatbot(label="chat with pdf")
+                input_msg = gr.Textbox(placeholder="input your question...", label="input")
+                clear_btn = gr.ClearButton([chatbot, input_msg])
             initial_btn.click(
                 self._init_settings,
                 inputs=[model_name, api_key, embedding_model, embedding_api_key, data_file]
@@ -119,6 +151,12 @@ class RAGDemo(object):
                 inputs=input_text,
                 outputs=output,
             )
+
+            input_msg.submit(
+                self._chat_qa,
+                inputs=[input_msg, chatbot],
+                outputs=[input_msg, chatbot]
+            )
         return demo