Update app.py

app.py

@@ -112,20 +112,20 @@ def conversation(qa_chain, message, history):
         formatted_chat_history = format_chat_history(history)
         response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
         response_answer = response["answer"]
-        if "Helpful Answer:" in response_answer:
-            response_answer = response_answer.split("Helpful Answer:")[-1]
         response_sources = response["source_documents"]
-        response_source1 = response_sources[0].page_content.strip()
-        response_source1_page = response_sources[0].metadata["page"] + 1
-        response_source2 = response_sources[1].page_content.strip()
-        response_source2_page = response_sources[1].metadata["page"] + 1
-        response_source3 = response_sources[2].page_content.strip()
-        response_source3_page = response_sources[2].metadata["page"] + 1
+
+        sources = []
+        for doc in response_sources:
+            sources.append({
+                "content": doc.page_content.strip(),
+                "page": doc.metadata["page"] + 1
+            })
+
         new_history = history + [(message, response_answer)]
-        return qa_chain, new_history, response_answer, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
+        return qa_chain, new_history, response_answer, sources
     except Exception as e:
         st.error(f"Error in conversation: {e}")
-        return qa_chain, history, "", "", 0, "", 0, "", 0
+        return qa_chain, history, "", []
 
 def main():
     st.sidebar.title("PDF Chatbot")
@@ -146,7 +146,7 @@ def main():
     if 'qa_chain' not in st.session_state:
         st.session_state['qa_chain'] = None
 
-    st.sidebar.markdown("### Select Large Language Model (LLM) and input parameters")
+    st.sidebar.markdown("### Select Large Language Model (LLM)")
     llm_option = st.sidebar.radio("Available LLMs", list_llm_simple)
 
     if st.sidebar.button("Initialize Question Answering Chatbot"):
@@ -163,19 +163,15 @@ def main():
 
         if st.button("Submit"):
             with st.spinner("Generating response..."):
-                qa_chain, history, response_answer, response_source1, source1_page, response_source2, source2_page, response_source3, source3_page = conversation(st.session_state['qa_chain'], message, history)
+                qa_chain, history, response_answer, sources = conversation(st.session_state['qa_chain'], message, history)
                 st.session_state['qa_chain'] = qa_chain
 
                 st.markdown("### Chatbot Response")
                 st.write(response_answer)
 
                 with st.expander("Relevant context from the source document"):
-                    st.text_area("Source 1", value=response_source1, height=100)
-                    st.text(f"Page: {source1_page}")
-                    st.text_area("Source 2", value=response_source2, height=100)
-                    st.text(f"Page: {source2_page}")
-                    st.text_area("Source 3", value=response_source3, height=100)
-                    st.text(f"Page: {source3_page}")
+                    for source in sources:
+                        st.text_area(f"Source - Page {source['page']}", value=source["content"], height=100)
 
 if __name__ == "__main__":
     main()