Insert empty user message when assisting with the current last message

Antonio Scandurra and Nathan Sobo created 2 years ago

Co-Authored-By: Nathan Sobo <[e-mail]>

Change summary

crates/ai/src/assistant.rs | 21 +++++++++++++++++++++
1 file changed, 21 insertions(+)

Detailed changes

crates/ai/src/assistant.rs

@@ -1023,6 +1023,14 @@ impl Conversation {
     ) -> Vec<MessageAnchor> {
         let mut user_messages = Vec::new();
         let mut tasks = Vec::new();
+
+        let last_message_id = self.message_anchors.iter().rev().find_map(|message| {
+            message
+                .start
+                .is_valid(self.buffer.read(cx))
+                .then_some(message.id)
+        });
+
         for selected_message_id in selected_messages {
             let selected_message_role =
                 if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
@@ -1085,6 +1093,19 @@ impl Conversation {
                     )
                     .unwrap();
 
+                // Queue up the user's next reply
+                if Some(selected_message_id) == last_message_id {
+                    let user_message = self
+                        .insert_message_after(
+                            assistant_message.id,
+                            Role::User,
+                            MessageStatus::Done,
+                            cx,
+                        )
+                        .unwrap();
+                    user_messages.push(user_message);
+                }
+
                 tasks.push(cx.spawn_weak({
                     |this, mut cx| async move {
                         let assistant_message_id = assistant_message.id;