Update app.py

app.py

@@ -61,13 +61,18 @@ def handle_query(query):
     
     query_engine = index.as_query_engine(text_qa_template=text_qa_template)
     answer = query_engine.query(query)
-    
+
+    final_ans = []
     if hasattr(answer, 'response'):
-        return answer.response
+        final_ans.append(answer.response)
     elif isinstance(answer, dict) and 'response' in answer:
-        return answer['response']
+        final_ans.append(answer['response'])
     else:
-        return "Sorry, I couldn't find an answer."
+        final_ans.append("Sorry, I couldn't find an answer.")
+
+    ans = " ".join(final_ans)
+    for i in ans:
+        yiled str(i)
 
 
 # Streamlit app initialization
@@ -77,7 +82,12 @@ st.markdown("chat here👇")
 
 if 'messages' not in st.session_state:
     st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}]
-
+    
+# Display or clear chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+        
 with st.sidebar:
     st.title("Menu:")
     uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button")
@@ -93,9 +103,19 @@ 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})
-    response = handle_query(user_prompt)
-    st.session_state.messages.append({'role': 'assistant', "content": response})
+    with st.chat_message("user", avatar="⛷️"):
+        st.write(prompt)
+    # response = handle_query(user_prompt)
+    # 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'])
+
+# Generate a new response if last message is not from assistant
+if st.session_state.messages[-1]["role"] != "assistant":
+    with st.chat_message("assistant"):
+        response = handle_query(user_prompt)
+        full_response = st.write_stream(response)
+    message = {"role": "assistant", "content": full_response}
+    st.session_state.messages.append(message)