Update crs_arena/arena.py

crs_arena/arena.py

@@ -37,6 +37,7 @@ from battle_manager import (
     get_unique_user_id,
 )
 from crs_fighter import CRSFighter
+from streamlit_lottie import st_lottie_spinner
 from utils import (
     download_and_extract_item_embeddings,
     download_and_extract_models,
@@ -68,6 +69,8 @@ if not os.path.exists("data/embed_items"):
 CONVERSATION_LOG_DIR = "data/arena/conversation_logs/"
 os.makedirs(CONVERSATION_LOG_DIR, exist_ok=True)
 
+TYPING_PLACEHOLDER_JSON = json.load(open("asset/Typing_animation.json", "r"))
+
 
 # Callbacks
 def record_vote(vote: str) -> None:
@@ -233,15 +236,22 @@ def chat_col(crs_id: int, color: str):
             st.session_state[f"messages_{crs_id}"].append(
                 {"role": "user", "message": prompt}
             )
-            # Get the CRS response
-            response_crs = messages_crs.chat_message("assistant").write_stream(
-                get_crs_response(st.session_state[f"crs{crs_id}"], prompt)
-            )
 
-            # Add CRS response to chat history
-            st.session_state[f"messages_{crs_id}"].append(
-                {"role": "assistant", "message": response_crs}
-            )
+            crs_message = messages_crs.chat_message("assistant").empty()
+            with crs_message:
+                # Placeholder for the CRS response
+                with st_lottie_spinner(
+                    TYPING_PLACEHOLDER_JSON, height=40, width=40
+                ):
+                    response_crs = get_crs_response(
+                        st.session_state[f"crs{crs_id}"], prompt
+                    )
+                    # Add CRS response to chat history
+                    st.session_state[f"messages_{crs_id}"].append(
+                        {"role": "assistant", "message": response_crs}
+                    )
+                # Display the CRS response
+                crs_message.write(response_crs)
 
         frustrated_col, satisfied_col = st.columns(2)
         if frustrated_col.button(