Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import os
-import json
 import base64
 from langchain.memory import ConversationBufferWindowMemory
 from langchain_community.chat_message_histories import StreamlitChatMessageHistory
@@ -11,25 +10,14 @@ from utils.qa import QAEngine
 # Configure Streamlit page
 st.set_page_config(page_title="AI-Powered Document QA", layout="wide")
 
-# # 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)
+# Function to encode image in Base64 for avatars
+def encode_image(image_path):
+    with open(image_path, "rb") as file:
+        return base64.b64encode(file.read()).decode()
+
+# Load avatar images
+user_avatar = encode_image("./icons/user.jpg")  # Change path if needed
+ai_avatar = encode_image("./icons/ai.jpg")
 
 # Initialize document processing & AI components
 document_processor = DocumentProcessor()
@@ -41,7 +29,7 @@ os.makedirs("temp", exist_ok=True)
 
 # Sidebar for file upload
 st.sidebar.header("Upload a PDF")
-uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type=["pdf"])
+uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type=["pdf", "docx", "html", "pptx"])
 
 # Initialize chat memory
 memory_storage = StreamlitChatMessageHistory(key="chat_messages")
@@ -64,25 +52,42 @@ if uploaded_file and "document_uploaded" not in st.session_state:
     st.sidebar.success("Document processed successfully!")
     st.session_state["document_uploaded"] = True
 
-# Chat interface layout
+# Chat UI Header
 st.markdown("<h2 style='text-align: center;'>AI Chat Assistant</h2>", unsafe_allow_html=True)
 st.markdown("---")
 
-# Display chat history
+# Display chat history with avatars
 for message in memory_storage.messages:
     role = "user" if message.type == "human" else "assistant"
+    avatar = user_avatar if role == "user" else ai_avatar  # Assign appropriate avatar
+
     with st.chat_message(role):
-        st.markdown(message.content)
+        st.markdown(
+            f"""
+            <div style="display: flex; align-items: center; margin-bottom: 10px;">
+                <img src="data:image/jpeg;base64,{avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
+                <div style="background-color: #f1f1f1; padding: 10px; border-radius: 10px; max-width: 80%;">{message.content}</div>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
 
 # User input at the bottom
 user_input = st.chat_input("Ask me anything...")
 
 if user_input:
-    # Store user message in memory
     memory_storage.add_user_message(user_input)
 
     with st.chat_message("user"):
-        st.markdown(user_input)
+        st.markdown(
+            f"""
+            <div style="display: flex; align-items: center; margin-bottom: 10px;">
+                <img src="data:image/jpeg;base64,{user_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
+                <div style="background-color: #d9f7be; padding: 10px; border-radius: 10px; max-width: 80%;">{user_input}</div>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
 
     with st.spinner("Generating response..."):
         if st.session_state.get("document_uploaded", False):
@@ -91,8 +96,15 @@ if user_input:
             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.content)
+        st.markdown(
+            f"""
+            <div style="display: flex; align-items: center; margin-bottom: 10px;">
+                <img src="data:image/jpeg;base64,{ai_avatar}" width="40" height="40" style="border-radius: 50%; margin-right: 10px;">
+                <div style="background-color: #e6f7ff; padding: 10px; border-radius: 10px; max-width: 80%;">{answer.content}</div>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )