Update app.py

app.py

@@ -1,9 +1,39 @@
-import random
 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"
@@ -12,34 +42,39 @@ if "messages" not in st.session_state:
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
-    with st.chat_message(message["role"], avatar=message["avatar"]):
-        st.markdown(message["content"])
+    if message["role"] == "user":
+        align_class = "chat-left"
+        box_class = "user-box"
+    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
+    )
 
 # Start with an assistant message
 if not st.session_state.messages:
-    st.chat_message("assistant", avatar=assistant_icon_path).markdown("type yes to start")
-    st.session_state.messages.append({"role": "assistant", "avatar":assistant_icon_path, "content": "type yes to start"})
+    st.session_state.messages.append({"role": "assistant", "content": "type yes to start"})
 
 if prompt := st.chat_input():
     if st.session_state.approval_state == "require":
-        st.chat_message("user", avatar=user_icon_path).markdown(prompt)
-        st.session_state.messages.append({"role": "user", "avatar":user_icon_path,"content": prompt})
-        if prompt == "yes":
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        if prompt.lower() == "yes":
             st.session_state.approval_state = "approved"
-            st.chat_message("assistant", avatar=assistant_icon_path).markdown("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"})
+            st.session_state.messages.append({"role": "assistant", "content": "token approved now you can use ai-tom chatbot"})
         else:
-            st.chat_message("assistant", avatar=assistant_icon_path).markdown("wrong token try again")
-            st.session_state.messages.append({"role": "assistant", "avatar":assistant_icon_path, "content": "wrong token try again"})
+            st.session_state.messages.append({"role": "assistant", "content": "wrong token try again"})
     elif st.session_state.approval_state == "approved":
-        # Display user message in chat message container
-        st.chat_message("user", avatar=user_icon_path).markdown(prompt)
-        st.session_state.messages.append({"role": "user", "avatar":user_icon_path, "content": prompt})
-        
+        st.session_state.messages.append({"role": "user", "content": prompt})
         response = f"Echo: {prompt}"
-        # Display assistant response in chat message container
-        with st.chat_message("assistant", avatar=assistant_icon_path):
-            st.markdown(response)
-        # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "avatar":assistant_icon_path, "content": response})
+        st.session_state.messages.append({"role": "assistant", "content": response})
 
+    st.rerun()  # 화면 갱신