Update app.py

app.py

@@ -117,18 +117,24 @@ if prompt := st.chat_input("Message..."):
     try:
         co = cohere.ClientV2(api_key)
         with st.chat_message("assistant", avatar=AI_PFP):
-            response = co.chat(
+            message_placeholder = st.empty()
+            full_response = ""
+            
+            for event in co.chat_stream(
                 model=selected_model,
                 messages=st.session_state.messages
-            )
-            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)
-
+            ):
+                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 + "▌")
+            
+            message_placeholder.markdown(full_response)
+            reply = full_response
+        
         st.session_state.messages.append({"role": "assistant", "content": reply})
 
     except Exception as e: