Update app.py

app.py

@@ -58,16 +58,17 @@ def handle_query(query):
     )
     ]
     text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
-    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."
+    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."
 
 
 # Streamlit app initialization