Update app.py

app.py

@@ -6,6 +6,8 @@ import time
 os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
 genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
 
+icons = {"assistant": "robot.png", "user": "man-kddi.png"}
+
 model = genai.GenerativeModel('gemini-1.5-flash-latest')
 prompt = """You are a programming teaching assistant named GenXAI(Generative eXpert AI), created by Pachaiappan. Answer only the programming, error-fixing and code-related question that being asked. 
 Important note, If Question non-related to coding or programming means, you have to say: 'Please ask only coding-related questions.' except those kind of questions "who are you", "who created you".
@@ -40,22 +42,22 @@ if 'messages' not in st.session_state:
 
 if 'history' not in st.session_state:
     st.session_state.history = [] 
-    
+
 for message in st.session_state.messages:
-    with st.chat_message(message['role']):
+    with st.chat_message(message['role'], avatar=icons[message['role']]):
         st.write(message['content'])
         
 user_input = st.chat_input("Ask Your Questions 👉..")
 if user_input:
     st.session_state.messages.append({'role': 'user', 'content': user_input})
-    with st.chat_message("user"):
+    with st.chat_message("user", avatar="man-kddi.png"):
         st.write(user_input)
         
     with st.spinner("Thinking..."):
         response = get_response(user_input)
         
-    with st.chat_message("assistant"):
-        # st.write(response)
+    with st.chat_message("user", avatar="robot.png"):
         st.write_stream(response_streaming(response))
-        message = {"role": "assistant", "content": response}
-        st.session_state.messages.append(message) 
+        
+    message = {"role": "assistant", "content": response}
+    st.session_state.messages.append(message)