Update app.py

app.py

@@ -10,6 +10,8 @@ import base64
 # Load environment variables
 load_dotenv()
 
+icons = {"assistant": "👽", "user": "👦🏻"}
+
 # Configure the Llama index settings
 Settings.llm = HuggingFaceInferenceAPI(
     model_name="meta-llama/Meta-Llama-3-8B-Instruct",
@@ -58,17 +60,16 @@ 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, streaming=True)
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
     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
@@ -80,7 +81,7 @@ 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']):
+    with st.chat_message(message['role'], avatar=icon[message['role']]):
         st.write(message['content'])
 
 with st.sidebar:
@@ -103,7 +104,7 @@ user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
 
 if user_prompt and uploaded_file:
     st.session_state.messages.append({'role': 'user', "content": user_prompt})
-    with st.chat_message("user", avatar="👽"):
+    with st.chat_message("user", avatar="👦🏻"):
         st.write(user_prompt)
 
 if st.session_state.messages[-1]["role"] != "assistant":