crates/assistant/src/assistant.rs 🔗
@@ -3,6 +3,7 @@ pub mod assistant_panel;
pub mod assistant_settings;
mod codegen;
mod completion_provider;
+mod omit_ranges;
mod prompts;
mod saved_conversation;
mod search;
Write a randomized test and fix bugs in the logic for omitting slash commands from completion requests (#12164)
Max Brunsfeld created
Release Notes:
- N/AChange summary
crates/assistant/src/assistant.rs | 1
crates/assistant/src/assistant_panel.rs | 38 ++--------
crates/assistant/src/omit_ranges.rs | 101 +++++++++++++++++++++++++++
3 files changed, 110 insertions(+), 30 deletions(-)
Detailed changes
crates/assistant/src/assistant_panel.rs 🔗
@@ -5,6 +5,7 @@ use crate::{
ambient_context::*,
assistant_settings::{AssistantDockPosition, AssistantSettings, ZedDotDevModel},
codegen::{self, Codegen, CodegenKind},
+ omit_ranges::text_in_range_omitting_ranges,
prompts::prompt::generate_content_prompt,
search::*,
slash_command::{
@@ -3556,38 +3557,15 @@ pub struct Message {
impl Message {
fn to_request_message(&self, buffer: &Buffer) -> LanguageModelRequestMessage {
- let mut slash_command_ranges = self.slash_command_ranges.iter().peekable();
- let mut content = String::with_capacity(self.offset_range.len());
- let mut offset = self.offset_range.start;
- let mut chunks = buffer.text_for_range(self.offset_range.clone());
- while let Some(chunk) = chunks.next() {
- if let Some(slash_command_range) = slash_command_ranges.peek() {
- match offset.cmp(&slash_command_range.start) {
- Ordering::Less => {
- let max_len = slash_command_range.start - offset;
- if chunk.len() < max_len {
- content.push_str(chunk);
- offset += chunk.len();
- } else {
- content.push_str(&chunk[..max_len]);
- offset += max_len;
- chunks.seek(slash_command_range.end);
- slash_command_ranges.next();
- }
- }
- Ordering::Equal | Ordering::Greater => {
- chunks.seek(slash_command_range.end);
- offset = slash_command_range.end;
- slash_command_ranges.next();
- }
- }
- } else {
- content.push_str(chunk);
- }
- }
+ let mut content = text_in_range_omitting_ranges(
+ buffer.as_rope(),
+ self.offset_range.clone(),
+ &self.slash_command_ranges,
+ );
+ content.truncate(content.trim_end().len());
LanguageModelRequestMessage {
role: self.role,
- content: content.trim_end().into(),
+ content,
}
}
}
crates/assistant/src/omit_ranges.rs 🔗
@@ -0,0 +1,101 @@
+use rope::Rope;
+use std::{cmp::Ordering, ops::Range};
+
+pub(crate) fn text_in_range_omitting_ranges(
+ rope: &Rope,
+ range: Range<usize>,
+ omit_ranges: &[Range<usize>],
+) -> String {
+ let mut content = String::with_capacity(range.len());
+ let mut omit_ranges = omit_ranges
+ .iter()
+ .skip_while(|omit_range| omit_range.end <= range.start)
+ .peekable();
+ let mut offset = range.start;
+ let mut chunks = rope.chunks_in_range(range.clone());
+ while let Some(chunk) = chunks.next() {
+ if let Some(omit_range) = omit_ranges.peek() {
+ match offset.cmp(&omit_range.start) {
+ Ordering::Less => {
+ let max_len = omit_range.start - offset;
+ if chunk.len() < max_len {
+ content.push_str(chunk);
+ offset += chunk.len();
+ } else {
+ content.push_str(&chunk[..max_len]);
+ chunks.seek(omit_range.end.min(range.end));
+ offset = omit_range.end;
+ omit_ranges.next();
+ }
+ }
+ Ordering::Equal | Ordering::Greater => {
+ chunks.seek(omit_range.end.min(range.end));
+ offset = omit_range.end;
+ omit_ranges.next();
+ }
+ }
+ } else {
+ content.push_str(chunk);
+ offset += chunk.len();
+ }
+ }
+
+ content
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use rand::{rngs::StdRng, Rng as _};
+ use util::RandomCharIter;
+
+ #[gpui::test(iterations = 100)]
+ fn test_text_in_range_omitting_ranges(mut rng: StdRng) {
+ let text = RandomCharIter::new(&mut rng).take(1024).collect::<String>();
+ let rope = Rope::from(text.as_str());
+
+ let mut start = rng.gen_range(0..=text.len() / 2);
+ let mut end = rng.gen_range(text.len() / 2..=text.len());
+ while !text.is_char_boundary(start) {
+ start -= 1;
+ }
+ while !text.is_char_boundary(end) {
+ end += 1;
+ }
+ let range = start..end;
+
+ let mut ix = 0;
+ let mut omit_ranges = Vec::new();
+ for _ in 0..rng.gen_range(0..10) {
+ let mut start = rng.gen_range(ix..=text.len());
+ while !text.is_char_boundary(start) {
+ start += 1;
+ }
+ let mut end = rng.gen_range(start..=text.len());
+ while !text.is_char_boundary(end) {
+ end += 1;
+ }
+ omit_ranges.push(start..end);
+ ix = end;
+ if ix == text.len() {
+ break;
+ }
+ }
+
+ let mut expected_text = text[range.clone()].to_string();
+ for omit_range in omit_ranges.iter().rev() {
+ let start = omit_range
+ .start
+ .saturating_sub(range.start)
+ .min(range.len());
+ let end = omit_range.end.saturating_sub(range.start).min(range.len());
+ expected_text.replace_range(start..end, "");
+ }
+
+ assert_eq!(
+ text_in_range_omitting_ranges(&rope, range.clone(), &omit_ranges),
+ expected_text,
+ "text: {text:?}\nrange: {range:?}\nomit_ranges: {omit_ranges:?}"
+ );
+ }
+}