Update app.py

app.py

@@ -135,12 +135,19 @@ def main():
             background-color: #2980b9;
         }
         .chat-message {
-            background-color: #ecf0f1;
+            display: flex;
+            align-items: center;
             padding: 15px;
             border-radius: 10px;
             margin-bottom: 10px;
             box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
         }
+        .chat-message img {
+            width: 40px;
+            height: 40px;
+            border-radius: 50%;
+            margin-right: 10px;
+        }
         .user-message {
             background-color: #d6eaf8;
         }
@@ -183,12 +190,13 @@ def main():
     # Display chat messages
     for message in st.session_state.messages:
         css_class = "user-message" if message["role"] == "user" else "assistant-message"
-        st.markdown(f'<div class="chat-message {css_class}">{message["content"]}</div>', unsafe_allow_html=True)
+        icon = "user_icon.png" if message["role"] == "user" else "assistant_icon.png"
+        st.markdown(f'<div class="chat-message {css_class}"><img src="{icon}" alt="icon"><div>{message["content"]}</div></div>', unsafe_allow_html=True)
 
     # Chat input
     if prompt := st.chat_input():
         st.session_state.messages.append({"role": "user", "content": prompt})
-        st.markdown(f'<div class="chat-message user-message">{prompt}</div>', unsafe_allow_html=True)
+        st.markdown(f'<div class="chat-message user-message"><img src="user_icon.png" alt="User"><div>{prompt}</div></div>', unsafe_allow_html=True)
 
         # Generate bot response
         if st.session_state.messages[-1]["role"] != "assistant":
@@ -197,7 +205,7 @@ def main():
                     response = user_input(prompt)
                     if response:
                         full_response = ''.join(response['output_text'])
-                        st.markdown(f'<div class="chat-message assistant-message">{full_response}</div>', unsafe_allow_html=True)
+                        st.markdown(f'<div class="chat-message assistant-message"><img src="assistant_icon.png" alt="Assistant"><div>{full_response}</div></div>', unsafe_allow_html=True)
                         message = {"role": "assistant", "content": full_response}
                         st.session_state.messages.append(message)
             except RuntimeError as e: