Update app.py

app.py

@@ -126,7 +126,9 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     progress(0.9, desc="Done!")
     return qa_chain
 
-def start(llm_model, temperature, max_tokens, top_k, vector_db, list_file_obj, chunk_size, chunk_overlap):
+def start(llm_model, temperature, max_tokens, top_k, 
+          vector_db, list_file_obj, chunk_size, chunk_overlap,
+         qa_chain, message, history):
     # HuggingFaceHub uses HF inference endpoints
     # Use of trust_remote_code as model_kwargs
     # Warning: langchain issue
@@ -174,7 +176,71 @@ def start(llm_model, temperature, max_tokens, top_k, vector_db, list_file_obj, c
     # Create or load vector database
     vector_db = create_db(doc_splits, collection_name)
 
+    formatted_chat_history = format_chat_history(message, history)
+    #print("formatted_chat_history",formatted_chat_history)
+   
+    # Generate response using QA chain
+    response = qa_chain({"question": message, "chat_history": formatted_chat_history})
+    response_answer = response["answer"]
+    if response_answer.find("Helpful Answer:") != -1:
+        response_answer = response_answer.split("Helpful Answer:")[-1]
+    response_sources = response["source_documents"]
+    response_source1 = response_sources[0].page_content.strip()
+    response_source2 = response_sources[1].page_content.strip()
+    response_source3 = response_sources[2].page_content.strip()
+    # Langchain sources are zero-based
+    response_source1_page = response_sources[0].metadata["page"] + 1
+    response_source2_page = response_sources[1].metadata["page"] + 1
+    response_source3_page = response_sources[2].metadata["page"] + 1
+    # print ('chat response: ', response_answer)
+    # print('DB source', response_sources)
     
-    return qa_chain, vector_db, collection_name
-
+    # Append user message and response to chat history
+    new_history = history + [(message, response_answer)]
+    
+    return qa_chain, vector_db, collection_name, new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
     
+def demo():
+    with gr.Blocks(theme="base") as demo:
+        vector_db = gr.State()
+        qa_chain = gr.State()
+        collection_name = gr.State()
+
+        chatbot = gr.Chatbot(height=300)
+        with gr.Accordion("Advanced - Document references", open=False):
+            with gr.Row():
+                doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
+                source1_page = gr.Number(label="Page", scale=1)
+            with gr.Row():
+                doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
+                source2_page = gr.Number(label="Page", scale=1)
+            with gr.Row():
+                doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
+                source3_page = gr.Number(label="Page", scale=1)
+        with gr.Row():
+            msg = gr.Textbox(placeholder="Type message", container=True)
+        with gr.Row():
+            submit_btn = gr.Button("Submit")
+            clear_btn = gr.ClearButton([msg, chatbot])
+
+        msg.submit(start, 
+                   inputs=[llm_model, temperature, max_tokens, top_k, 
+                           vector_db, list_file_obj, chunk_size, chunk_overlap,
+                           qa_chain, message, history], 
+                   outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, 
+                            doc_source2, source2_page, 
+                            doc_source3, source3_page], 
+                   queue=False)
+        submit_btn.click(conversation, \
+            inputs=[qa_chain, msg, chatbot], \
+            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
+            queue=False)
+        clear_btn.click(lambda:[None,"",0,"",0,"",0], \
+            inputs=None, \
+            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
+            queue=False)
+    
+    demo.queue().launch(debug=True)
+
+if __name__ == "__main__":
+    demo()