Update app.py

app.py

@@ -58,15 +58,17 @@ def handle_query(query):
     )
     ]
     text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
-    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template, streaming=True)
     answer = query_engine.query(query)
+    yield answer.print_response_stream()
     
-    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."
+    
+    # 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
@@ -77,6 +79,10 @@ 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.'}]
 
+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")
@@ -90,11 +96,19 @@ with st.sidebar:
             st.success("Done")
 
 user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
-if user_prompt:
+# 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})
+
+if user_prompt and uploaded_file:
     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(user_prompt)
 
-for message in st.session_state.messages:
-    with st.chat_message(message['role']):
-        st.write(message['content'])
+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)