Update app.py

app.py

@@ -10,9 +10,9 @@ warnings.filterwarnings("ignore")
 st.set_page_config(layout="wide")
 st.title("AI Document Processor with Conversational RAG")
 
-# Initialize conversation history in session state
-if "conversation_history" not in st.session_state:
-    st.session_state.conversation_history = []
+# Initialize chat history in session state
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
 # Sidebar for file upload and settings
 with st.sidebar:
@@ -27,14 +27,16 @@ with st.sidebar:
     llm_option = st.selectbox("Select LLM Model", ["GPT-4o", "GPT-4o-mini"], index=1)
 
     temp_dir = "temp_uploaded_files"
-    
+
+    # Initialize file_paths to prevent undefined variable errors
+    file_paths = []  
+
     # Clear the previous uploads when new files are uploaded
     if uploaded_files:
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir)  # Delete the old directory and its contents
         os.makedirs(temp_dir)  # Create a fresh directory
 
-        file_paths = []  # List to store saved file paths
         for file in uploaded_files:
             file_path = os.path.join(temp_dir, file.name)
             with open(file_path, "wb") as f:
@@ -42,31 +44,37 @@ with st.sidebar:
             file_paths.append(file_path)
             st.write(f"✅ Saved: {file.name}")
 
-# Clear conversation history if the user refreshes or starts a new session
+# Button to clear conversation history
 if st.button("Clear Conversation History"):
-    st.session_state.conversation_history = []  # Reset history
-
-# Chat interface
-st.write("### Chat Interface")
+    st.session_state.messages = []
 
-chat_display = "\n".join(st.session_state.conversation_history)
-# st.text_area("Conversation History", chat_display, height=300, disabled=True)
+# Display previous chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"], unsafe_allow_html=True)  # Allow better formatting
 
-# Input for user question
-user_input = st.text_input("Ask a question:")
-llm_model = llm_option
+# User input
+if user_input := st.chat_input("Ask a question:"):
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(user_input)
+    st.session_state.messages.append({"role": "user", "content": user_input})
 
-if st.button("Retrieve and Answer"):
-    if user_input or uploaded_files:
-        answer = multimodelrag(user_input, file_paths, embeding,  llm_model, conversation)
+    # Get AI response
+    answer = multimodelrag(user_input, file_paths, embeding, llm_option, conversation)
 
-        # Update conversation history
-        st.session_state.conversation_history.append(f"User: {user_input}")
-        st.session_state.conversation_history.append(f"AI: {answer}")
+    # Check if answer is valid
+    if isinstance(answer, dict):
+        chat_display = "\n\n".join(
+            [f"📂 <b>File:</b> <code>{key}</code><br><br>📝 <b>Extracted Text:</b><br><br>{value}<br>{'-'*50}" 
+            for key, value in answer.items()]
+        )
+    else:
+        chat_display = f"❌ **Error:** Unexpected response format from AI model.\n\n{answer}"
 
-        # Refresh chat display
-        chat_display = "\n".join(st.session_state.conversation_history)
-        st.text_area("Conversation History", chat_display, height=400, disabled=True)
+    # Display AI response in chat
+    with st.chat_message("assistant"):
+        st.markdown(chat_display, unsafe_allow_html=True)  # Allow HTML formatting for bold text
 
-        st.write("### Answer:")
-        st.write(answer)
+    # Store assistant response in conversation history
+    st.session_state.messages.append({"role": "assistant", "content": chat_display})