Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import cohere
 import os
-import time
 
 st.set_page_config(page_title="Cohere Chat", layout="wide")
 
@@ -118,39 +117,18 @@ if prompt := st.chat_input("Message..."):
     try:
         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(
+            response = co.chat(
                 model=selected_model,
                 messages=st.session_state.messages
-            ):
-                if event.type == "content-delta":
-                    if hasattr(event.delta.message, "content") and event.delta.message.content:
-                        for content_item in event.delta.message.content:
-                            if hasattr(content_item, "text"):
-                                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
-        
+            )
+            if hasattr(response, "message") and hasattr(response.message, "content"):
+                content_items = response.message.content
+                reply = "".join(getattr(item, 'text', '') for item in content_items)
+            else:
+                st.write(response)
+                reply = "❗️Couldn't extract reply from the Cohere response."
+            st.markdown(reply)
+
         st.session_state.messages.append({"role": "assistant", "content": reply})
 
     except Exception as e: