Update app.py

app.py

@@ -1,34 +1,38 @@
-import streamlit as st
-from document_chat import ingest_pdf, process_query_with_memory
-
-#configure streamlit app
-st.set_page_config(page_title="AI Document Q&A Chatbot", layout="wide")
-st.title("📄 AI-Powered Document Chatbot")
-st.write("Upload a document and ask questions!")
-
-#upload document
-uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
-if uploaded_file:
-    file_path = "uploaded_doc.pdf"
-    with open(file_path, "wb") as f:
-        f.write(uploaded_file.getbuffer())
-
-    st.success("File uploaded! Processing...")
-    ingest_pdf(file_path)
-
-if "chat_history" not in st.session_state:
-    st.session_state["chat_history"] = []
-
-query = st.text_input("Ask a question:")
-if query:
-    with st.spinner("Thinking..."):
-        response = process_query_with_memory(query, st.session_state["chat_history"])
-        st.session_state["chat_history"].append((query, response))
-        st.write(response)
-
-# Show chat history
-if st.session_state["chat_history"]:
-    st.subheader("Chat History")
-    for q, a in st.session_state["chat_history"]:
-        st.write(f"**User:** {q}")
-        st.write(f"**Bot:** {a}")
+import streamlit as st
+from document_chat import ingest_pdf, process_query_with_memory
+from langchain.memory import ConversationBufferMemory
+
+# Configure Streamlit app
+st.set_page_config(page_title="AI Document Q&A Chatbot", layout="wide")
+st.title("📄 AI-Powered Document Chatbot")
+st.write("Upload a document and ask questions!")
+
+# Upload document
+uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
+if uploaded_file:
+    file_path = "uploaded_doc.pdf"
+    with open(file_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+
+    st.success("File uploaded! Processing...")
+    ingest_pdf(file_path)
+
+# Initialize memory if not exists
+if "memory" not in st.session_state:
+    st.session_state["memory"] = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+
+query = st.text_input("Ask a question:")
+if query:
+    with st.spinner("Thinking..."):
+        response = process_query_with_memory(query, st.session_state["memory"])
+        st.session_state["memory"].save_context({"input": query}, {"output": response})
+        st.write(response)
+
+# Show chat history
+if st.session_state["memory"].chat_memory.messages:
+    st.subheader("Chat History")
+    for i in range(0, len(st.session_state["memory"].chat_memory.messages), 2):
+        user_message = st.session_state["memory"].chat_memory.messages[i].content
+        bot_response = st.session_state["memory"].chat_memory.messages[i + 1].content if i + 1 < len(st.session_state["memory"].chat_memory.messages) else "..."
+        st.write(f"**User:** {user_message}")
+        st.write(f"**Bot:** {bot_response}")