Update app.py

app.py

@@ -28,8 +28,8 @@ qa_engine = QAEngine()
 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", "docx", "html", "pptx"])
+st.sidebar.header("Upload a Document")
+uploaded_file = st.sidebar.file_uploader("Choose a file", type=["pdf", "docx", "html", "pptx", "txt"])
 
 # Initialize chat memory
 memory_storage = StreamlitChatMessageHistory(key="chat_messages")
@@ -39,15 +39,15 @@ memory = ConversationBufferWindowMemory(
 
 # Document upload & processing
 if uploaded_file and "document_uploaded" not in st.session_state:
-    pdf_path = os.path.join("temp", uploaded_file.name)
+    file_path = os.path.join("temp", uploaded_file.name)
 
-    with open(pdf_path, "wb") as f:
+    with open(file_path, "wb") as f:
         f.write(uploaded_file.read())
 
     st.sidebar.success("File uploaded successfully!")
 
     with st.spinner("Processing document..."):
-        document_processor.process_document(pdf_path)
+        document_processor.process_document(file_path)
 
     st.sidebar.success("Document processed successfully!")
     st.session_state["document_uploaded"] = True
@@ -61,16 +61,15 @@ 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(
-            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
-        )
+    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...")
@@ -78,16 +77,16 @@ user_input = st.chat_input("Ask me anything...")
 if user_input:
     memory_storage.add_user_message(user_input)
 
-    with st.chat_message("user"):
-        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.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
+    ):
+        pass
 
     with st.spinner("Generating response..."):
         if st.session_state.get("document_uploaded", False):
@@ -98,13 +97,13 @@ if user_input:
 
     memory_storage.add_ai_message(answer)
 
-    with st.chat_message("assistant"):
-        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
-        )
+    with 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
+    ):
+        pass