Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import cohere
 import os
+import time
 
 st.set_page_config(page_title="Cohere Chat", layout="wide")
 
@@ -118,7 +119,12 @@ if prompt := st.chat_input("Message..."):
         co = cohere.ClientV2(api_key)
         with st.chat_message("assistant", avatar=AI_PFP):
             message_placeholder = st.empty()
+            
+            progress_bar = st.progress(0)
+            
             full_response = ""
+            chunk_count = 0
+            max_chunks = 20  
             
             for event in co.chat_stream(
                 model=selected_model,
@@ -131,8 +137,18 @@ if prompt := st.chat_input("Message..."):
                                 delta_text = getattr(content_item, "text", "")
                                 full_response += delta_text
                                 message_placeholder.markdown(full_response + "▌")
+                                
+                                chunk_count += 1
+                                progress_value = min(chunk_count / max_chunks, 0.95) 
+                                progress_bar.progress(progress_value)
             
+            progress_bar.progress(1.0)
             message_placeholder.markdown(full_response)
+            
+            time.sleep(0.5)
+            
+            progress_bar.empty()
+            
             reply = full_response
         
         st.session_state.messages.append({"role": "assistant", "content": reply})