add conversation history (MULTI-turn conversation support)

app.py

@@ -6,6 +6,10 @@ import fitz
 from langchain_community.vectorstores import Chroma
 from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.chains import create_history_aware_retriever, create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import HumanMessage, AIMessage
 from langchain.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain_core.output_parsers import StrOutputParser
@@ -159,30 +163,40 @@ async def chat_with_pdf(message, history, state: SessionState):
     retriever = state.db.as_retriever()
     llm = ChatGoogleGenerativeAI(model=LLM_MODEL, temperature=0.7, google_api_key=google_api_key)
 
-    prompt_template = PromptTemplate(
-        template="""
-        You are a helpful assistant for a PDF document.
-        Answer the user's question based on the following context.
-        If you don't know the answer, just say that you don't know, don't try to make up an answer.
-        ----------------
-        Context: {context}
-        Question: {question}
-        """,
-        input_variables=["context", "question"],
+    condenser_prompt = ChatPromptTemplate.from_messages([
+        ("system", "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."),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+    ])
+    
+    history_aware_retriever = create_history_aware_retriever(
+        llm, retriever, condenser_prompt
     )
 
-    rag_chain = (
-        {"context": retriever, "question": RunnablePassthrough()}
-        | prompt_template
-        | llm
-        | StrOutputParser()
-    )
+    qa_prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a helpful assistant for a PDF document. Answer the user's question based on the following context. If you don't know the answer, just say that you don't know, don't try to make up an answer.\n\n{context}"),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+    ])
+
+    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
+
+    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
 
-    response = await rag_chain.ainvoke(message)
-    yield response
+    chat_history_for_chain = []
+    for user_msg, ai_msg in history:
+        chat_history_for_chain.append(HumanMessage(content=user_msg))
+        chat_history_for_chain.append(AIMessage(content=ai_msg))
+
+    response = await rag_chain.ainvoke({
+        "chat_history": chat_history_for_chain,
+        "input": message
+    })
+
+    yield response["answer"]
 
 with gr.Blocks(title="PDF Chatbot") as demo:
-    state = gr.State(value=SessionState())
+    state = gr.State()
 
     gr.Markdown(
         """
@@ -190,27 +204,42 @@ with gr.Blocks(title="PDF Chatbot") as demo:
         Upload a PDF to start a conversation with your document.
         """
     )
-    
-    file_upload_input = gr.File(
-        file_types=[".pdf"],
-        label="Upload your PDF document",
-        interactive=True
-    )
 
-    chat_interface = gr.ChatInterface(
-        fn=chat_with_pdf,
-        additional_inputs=[state],
-        chatbot=gr.Chatbot(type="messages"),
-        textbox=gr.Textbox(placeholder="Type your question here...", scale=7, interactive=False),
-        examples=[["What is the main topic of the document?"], ["Summarize the key findings."], ["Who are the authors?"]],
-        title="Chat Interface",
-        theme="soft"
-    )
+    with gr.Row():
+        file_upload_input = gr.File(
+            file_types=[".pdf"],
+            label="Upload your PDF document",
+            interactive=True
+        )
+
+    with gr.Row(visible=False) as chat_row:
+        chat_interface = gr.ChatInterface(
+            fn=chat_with_pdf,
+            additional_inputs=[state],
+            chatbot=gr.Chatbot(type="messages"),
+            textbox=gr.Textbox(placeholder="Type your question here...", scale=7),
+            examples=[["What is the main topic of the document?"], ["Summarize the key findings."], ["Who are the authors?"]],
+            title="Chat Interface",
+            theme="soft",
+            type="messages"
+        )
+
+    async def process_and_show_chat(file, state):
+        gr.Info("Processing your PDF, please wait...")
+        try:
+            new_state = SessionState()
+            await process_pdf(file, new_state)
+            gr.Info("PDF processed successfully! You can now chat with it.")
+            return gr.update(visible=True), gr.update(interactive=False), new_state
+        except Exception:
+            # The exception is already a gr.Error, so it will be displayed in the UI.
+            # We just need to return the correct UI updates.
+            return gr.update(visible=False), gr.update(interactive=True), state
 
     file_upload_input.upload(
-        fn=process_pdf,
+        fn=process_and_show_chat,
         inputs=[file_upload_input, state],
-        outputs=[file_upload_input, chat_interface.textbox, state]
+        outputs=[chat_row, file_upload_input, state]
     )
 
-demo.launch()
+demo.launch()