Update app.py

app.py

@@ -1,14 +1,37 @@
 import streamlit as st
 import os
 import json
+import base64
+from langchain.memory import ConversationBufferWindowMemory
+from langchain_community.chat_message_histories import StreamlitChatMessageHistory
 from utils.ingestion import DocumentProcessor
 from utils.llm import LLMProcessor
 from utils.qa import QAEngine
 
+# Configure Streamlit page
 st.set_page_config(page_title="AI-Powered Document QA", layout="wide")
-st.title("📄 AI-Powered Document QA")
 
-# Initialize processors
+# Background Image
+def add_bg_from_local(image_file):
+    with open(image_file, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read())
+    st.markdown(
+        f"""
+        <style>
+            .stApp {{
+                background-image: url(data:image/png;base64,{encoded_string.decode()});
+                background-size: cover;
+            }}
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
+# Path to background image
+image_bg = "./image/background.jpeg"  # Change this path accordingly
+add_bg_from_local(image_bg)
+
+# Initialize document processing & AI components
 document_processor = DocumentProcessor()
 llm_processor = LLMProcessor()
 qa_engine = QAEngine()
@@ -17,49 +40,59 @@ qa_engine = QAEngine()
 os.makedirs("temp", exist_ok=True)
 
 # Sidebar for file upload
-st.sidebar.header("📂 Upload a PDF")
+st.sidebar.header("Upload a PDF")
 uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type=["pdf"])
 
+# Initialize chat memory
+memory_storage = StreamlitChatMessageHistory(key="chat_messages")
+memory = ConversationBufferWindowMemory(
+    memory_key="chat_history", human_prefix="User", chat_memory=memory_storage, k=5
+)
+
 # Document upload & processing
-if uploaded_file:
+if uploaded_file and "document_uploaded" not in st.session_state:
     pdf_path = os.path.join("temp", uploaded_file.name)
-    
+
     with open(pdf_path, "wb") as f:
         f.write(uploaded_file.read())
 
-    st.sidebar.success("✅ File uploaded successfully!")
+    st.sidebar.success("File uploaded successfully!")
 
-    with st.spinner("🔄 Processing document..."):
+    with st.spinner("Processing document..."):
         document_processor.process_document(pdf_path)
 
-    st.sidebar.success("✅ Document processed successfully!")
+    st.sidebar.success("Document processed successfully!")
     st.session_state["document_uploaded"] = True
-else:
-    st.session_state["document_uploaded"] = False
 
-# Divider between sections
+# Chat interface layout
+st.markdown("<h2 style='text-align: center;'>AI Chat Assistant</h2>", unsafe_allow_html=True)
 st.markdown("---")
 
-# Q&A Section
-st.header("🔍 Ask a Question")
+# Display chat history
+for message in memory_storage.messages:
+    role = "user" if message.type == "human" else "assistant"
+    with st.chat_message(role):
+        st.markdown(message.content)
 
-question = st.text_input("Ask a question:", placeholder="What are the key insights?")
+# User input at the bottom
+user_input = st.chat_input("Ask me anything...")
 
-if st.button("💡 Get Answer"):
-    if question:
-        with st.spinner("🧠 Generating response..."):
-            if st.session_state["document_uploaded"]:
-                # Use document-based QA if a file is uploaded
-                answer = qa_engine.query(question)
-            else:
-                # Use AI-based response if no document is uploaded
-                answer = llm_processor.generate_answer("", question)
-                st.warning("⚠️ No document uploaded. This response is generated from general AI knowledge and may not be document-specific.")
+if user_input:
+    # Store user message in memory
+    memory_storage.add_user_message(user_input)
 
-        st.subheader("📝 Answer:")
-        st.write(answer.content)
-    else:
-        st.warning("⚠️ Please enter a question.")
+    with st.chat_message("user"):
+        st.markdown(user_input)
 
-st.markdown("---")
-st.caption("🤖 Powered by ChromaDB + Groq LLM | Built with ❤️ using Streamlit")
+    with st.spinner("Generating response..."):
+        if st.session_state.get("document_uploaded", False):
+            answer = qa_engine.query(user_input)
+        else:
+            answer = llm_processor.generate_answer("", user_input)
+            st.warning("No document uploaded. This response is generated from general AI knowledge and may not be document-specific.")
+
+    # Store AI response in memory
+    memory_storage.add_ai_message(answer)
+
+    with st.chat_message("assistant"):
+        st.markdown(answer)