Update app.py

app.py

@@ -1,3 +1,4 @@
+import random
 import streamlit as st
 
 assistant_icon_path = "chat_icon/assistant.ico"
@@ -9,34 +10,76 @@ if "approval_state" not in st.session_state:
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat messages from history on app rerun
+# Custom CSS to align messages
+st.markdown("""
+    <style>
+    .chat-container {
+        display: flex;
+        align-items: center;
+        margin-bottom: 10px;
+    }
+    .user-container {
+        justify-content: flex-start;
+    }
+    .assistant-container {
+        justify-content: flex-end;
+    }
+    .chat-bubble {
+        padding: 10px;
+        border-radius: 10px;
+        max-width: 70%;
+    }
+    .user-bubble {
+        background-color: #d1e7dd;
+        text-align: left;
+    }
+    .assistant-bubble {
+        background-color: #f8d7da;
+        text-align: right;
+    }
+    .avatar {
+        width: 40px;
+        height: 40px;
+        border-radius: 50%;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# Display chat messages from history
 for message in st.session_state.messages:
-    col1, col2 = st.columns([1, 1])  # 두 개의 동일한 크기의 열 생성
-    
     if message["role"] == "user":
-        with col1:  # 왼쪽 열에 사용자 메시지 표시
-            with st.chat_message("user", avatar=user_icon_path):
-                st.markdown(message["content"])
+        st.markdown(
+            f"""
+            <div class="chat-container user-container">
+                <img src="{user_icon_path}" class="avatar">
+                <div class="chat-bubble user-bubble">{message["content"]}</div>
+            </div>
+            """, unsafe_allow_html=True)
     else:
-        with col2:  # 오른쪽 열에 AI 메시지 표시
-            with st.chat_message("assistant", avatar=assistant_icon_path):
-                st.markdown(message["content"])
+        st.markdown(
+            f"""
+            <div class="chat-container assistant-container">
+                <div class="chat-bubble assistant-bubble">{message["content"]}</div>
+                <img src="{assistant_icon_path}" class="avatar">
+            </div>
+            """, unsafe_allow_html=True)
 
 # Start with an assistant message
 if not st.session_state.messages:
-    st.session_state.messages.append({"role": "assistant", "content": "type yes to start"})
+    st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": "Type yes to start"})
+    st.rerun()
 
 if prompt := st.chat_input():
     if st.session_state.approval_state == "require":
-        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.session_state.messages.append({"role": "user", "avatar": user_icon_path, "content": prompt})
         if prompt.lower() == "yes":
             st.session_state.approval_state = "approved"
-            st.session_state.messages.append({"role": "assistant", "content": "token approved now you can use ai-tom chatbot"})
+            st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": "Token approved. Now you can use AI-Tom chatbot."})
         else:
-            st.session_state.messages.append({"role": "assistant", "content": "wrong token try again"})
+            st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": "Wrong token, try again."})
     elif st.session_state.approval_state == "approved":
-        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.session_state.messages.append({"role": "user", "avatar": user_icon_path, "content": prompt})
         response = f"Echo: {prompt}"
-        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": response})
 
-    st.rerun()  # 화면 갱신
+    st.rerun()