Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Conversation::messages_for_offsets with empty message at the end

Browse source
This commit is contained in:
Antonio Scandurra 2023-06-23 08:58:30 +02:00
parent 1707652643
commit ff07d0c2ed
1 changed files with 27 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

34
crates/ai/src/assistant.rs
View file

@ -1300,21 +1300,21 @@ impl Conversation {
) -> Vec<Message> { ) -> Vec<Message> {
let mut result = Vec::new(); let mut result = Vec::new();
let buffer_len = self.buffer.read(cx).len();
let mut messages = self.messages(cx).peekable(); let mut messages = self.messages(cx).peekable();
let mut offsets = offsets.into_iter().peekable(); let mut offsets = offsets.into_iter().peekable();
let mut current_message = messages.next();
while let Some(offset) = offsets.next() { while let Some(offset) = offsets.next() {
// Skip messages that start after the offset. // Locate the message that contains the offset.
while messages.peek().map_or(false, |message| { while current_message.as_ref().map_or(false, |message| {
message.range.end < offset || (message.range.end == offset && offset < buffer_len) !message.range.contains(&offset) && messages.peek().is_some()
}) { }) {
messages.next(); current_message = messages.next();
} }
let Some(message) = messages.peek() else { continue }; let Some(message) = current_message.as_ref() else { break };
// Skip offsets that are in the same message. // Skip offsets that are in the same message.
while offsets.peek().map_or(false, |offset| { while offsets.peek().map_or(false, |offset| {
message.range.contains(offset) || message.range.end == buffer_len message.range.contains(offset) || messages.peek().is_none()
}) { }) {
offsets.next(); offsets.next();
} }
@ -2360,6 +2360,26 @@ mod tests {
[message_1.id, message_3.id] [message_1.id, message_3.id]
); );
let message_4 = conversation
.update(cx, |conversation, cx| {
conversation.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)
})
.unwrap();
assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_2.id, Role::User, 4..8),
(message_3.id, Role::User, 8..12),
(message_4.id, Role::User, 12..12)
]
);
assert_eq!(
message_ids_for_offsets(&conversation, &[0, 4, 8, 12], cx),
[message_1.id, message_2.id, message_3.id, message_4.id]
);
fn message_ids_for_offsets( fn message_ids_for_offsets(
conversation: &ModelHandle<Conversation>, conversation: &ModelHandle<Conversation>,
offsets: &[usize], offsets: &[usize],