Update app.py

app.py

@@ -7,6 +7,7 @@ from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationalRetrievalChain, ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.document_loaders import PyPDFLoader
+import time
 
 # Initialize session state variables
 if "messages" not in st.session_state:
@@ -15,6 +16,8 @@ if "chain" not in st.session_state:
     st.session_state.chain = None
 if "processed_pdfs" not in st.session_state:
     st.session_state.processed_pdfs = False
+if "waiting_for_answer" not in st.session_state:
+    st.session_state.waiting_for_answer = False
 
 def create_sidebar():
     with st.sidebar:
@@ -89,20 +92,15 @@ def process_pdfs(papers, api_key):
             return texts
     return []
 
-def handle_chat(prompt, texts):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    
+def get_assistant_response(prompt, texts):
     try:
         if texts or st.session_state.processed_pdfs:
             result = st.session_state.chain({"question": prompt})
-            response = result["answer"]
+            return result["answer"]
         else:
-            response = "Please upload a PDF first."
-        
-        st.session_state.messages.append({"role": "assistant", "content": response})
-        
+            return "Please upload a PDF first."
     except Exception as e:
-        st.error(f"Error: {str(e)}")
+        return f"Error: {str(e)}"
 
 def main():
     st.set_page_config(page_title="PDF Chat", layout="wide")
@@ -121,20 +119,29 @@ def main():
     # Process PDFs
     texts = process_pdfs(papers, api_key)
     
-    # Chat interface container
+    # Chat interface
     chat_container = st.container()
     
     with chat_container:
-        # Display chat messages
+        # Display existing chat messages
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
-    
-        # Chat input
+        
+        # Get user input
         if prompt := st.chat_input("Ask about your PDFs"):
-            handle_chat(prompt, texts)
-            # Force a rerun to display the new message
-            st.rerun()
+            # Add user message immediately
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            st.chat_message("user").markdown(prompt)
+            
+            # Get assistant response with a loading indicator
+            with st.chat_message("assistant"):
+                with st.spinner("Thinking..."):
+                    response = get_assistant_response(prompt, texts)
+                st.markdown(response)
+            
+            # Add assistant response to messages
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
 if __name__ == "__main__":
     main()