Update app.py

app.py

@@ -58,17 +58,16 @@ def handle_query(query):
     )
     ]
     text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            query_engine = index.as_query_engine(text_qa_template=text_qa_template)
-            answer = query_engine.query(query)
-            
-            if hasattr(answer, 'response'):
-                return answer.response
-            elif isinstance(answer, dict) and 'response' in answer:
-                return answer['response']
-            else:
-                return "Sorry, I couldn't find an answer."
+    
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
+    answer = query_engine.query(query)
+    
+    if hasattr(answer, 'response'):
+        return answer.response
+    elif isinstance(answer, dict) and 'response' in answer:
+        return answer['response']
+    else:
+        return "Sorry, I couldn't find an answer."
 
 
 # Streamlit app initialization
@@ -94,9 +93,16 @@ with st.sidebar:
 user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
 if user_prompt:
     st.session_state.messages.append({'role': 'user', "content": user_prompt})
+    for message in st.session_state.messages:
+        with st.chat_message(message['role']):
+            st.write(message['content'])
     response = handle_query(user_prompt)
+    if st.session_state.messages[-1]["role"] != "assistant":
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+            st.write(response)
     st.session_state.messages.append({'role': 'assistant', "content": response})
 
-for message in st.session_state.messages:
-    with st.chat_message(message['role']):
-        st.write(message['content'])
+# for message in st.session_state.messages:
+#     with st.chat_message(message['role']):
+#         st.write(message['content'])