Update app.py

app.py

@@ -3,37 +3,6 @@ import streamlit as st
 assistant_icon_path = "chat_icon/assistant.ico"
 user_icon_path = "chat_icon/user01.png"
 
-# Custom CSS for alignment
-st.markdown("""
-    <style>
-        .chat-container {
-            display: flex;
-            width: 100%;
-        }
-        .chat-left {
-            justify-content: flex-start;
-            text-align: left;
-        }
-        .chat-right {
-            justify-content: flex-end;
-            text-align: right;
-        }
-        .chat-box {
-            max-width: 60%;
-            padding: 10px;
-            margin: 5px;
-            border-radius: 10px;
-            display: inline-block;
-        }
-        .user-box {
-            background-color: #d1e7fd;
-        }
-        .assistant-box {
-            background-color: #f1f1f1;
-        }
-    </style>
-""", unsafe_allow_html=True)
-
 # 세션 상태에서 approval_state 및 messages 초기화
 if "approval_state" not in st.session_state:
     st.session_state.approval_state = "require"
@@ -42,23 +11,16 @@ if "messages" not in st.session_state:
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
+    col1, col2 = st.columns([1, 1])  # 두 개의 동일한 크기의 열 생성
+    
     if message["role"] == "user":
-        align_class = "chat-left"
-        box_class = "user-box"
+        with col1:  # 왼쪽 열에 사용자 메시지 표시
+            with st.chat_message("user", avatar=user_icon_path):
+                st.markdown(message["content"])
     else:
-        align_class = "chat-right"
-        box_class = "assistant-box"
-
-    st.markdown(
-        f"""
-        <div class="chat-container {align_class}">
-            <div class="chat-box {box_class}">
-                <b>{message["role"].capitalize()}:</b> {message["content"]}
-            </div>
-        </div>
-        """, 
-        unsafe_allow_html=True
-    )
+        with col2:  # 오른쪽 열에 AI 메시지 표시
+            with st.chat_message("assistant", avatar=assistant_icon_path):
+                st.markdown(message["content"])
 
 # Start with an assistant message
 if not st.session_state.messages: