Skip to content

Commit

Permalink
assistant2: More improvement to prompt building efficiency
Browse files Browse the repository at this point in the history
  • Loading branch information
mgsloan committed Jan 10, 2025
1 parent ec4c674 commit 8ccb33d
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 35 deletions.
76 changes: 48 additions & 28 deletions crates/assistant2/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ pub struct ContextSnapshot {
pub tooltip: Option<SharedString>,
pub icon_path: Option<SharedString>,
pub kind: ContextKind,
/// Concatenating these strings yields text to send to the model. Not refreshed by `snapshot`.
/// Joining these strings separated by \n yields text for model. Not refreshed by `snapshot`.
pub text: Box<[SharedString]>,
}

Expand Down Expand Up @@ -210,62 +210,82 @@ pub fn attach_context_to_message(
let mut fetch_context = Vec::new();
let mut thread_context = Vec::new();

let mut capacity = 0;
for context in contexts {
capacity += context.text.len();
match context.kind {
ContextKind::File => file_context.push(context),
ContextKind::Directory => directory_context.push(context),
ContextKind::FetchedUrl => fetch_context.push(context),
ContextKind::Thread => thread_context.push(context),
}
}
if !file_context.is_empty() {
capacity += 1;
}
if !directory_context.is_empty() {
capacity += 1;
}
if !fetch_context.is_empty() {
capacity += 1 + fetch_context.len();
}
if !thread_context.is_empty() {
capacity += 1 + thread_context.len();
}
if capacity == 0 {
return;
}

let mut context_text = String::new();
let mut context_chunks = Vec::with_capacity(capacity);

if !file_context.is_empty() {
context_text.push_str("The following files are available:\n");
for context in file_context {
for chunk in context.text {
context_text.push_str(&chunk);
context_chunks.push("The following files are available:\n");
for context in &file_context {
for chunk in &context.text {
context_chunks.push(&chunk);
}
context_text.push('\n');
}
}

if !directory_context.is_empty() {
context_text.push_str("The following directories are available:\n");
for context in directory_context {
for chunk in context.text {
context_text.push_str(&chunk);
context_chunks.push("The following directories are available:\n");
for context in &directory_context {
for chunk in &context.text {
context_chunks.push(&chunk);
}
context_text.push('\n');
}
}

if !fetch_context.is_empty() {
context_text.push_str("The following fetched results are available\n");
for context in fetch_context {
context_text.push_str(&context.name);
context_text.push('\n');
for chunk in context.text {
context_text.push_str(&chunk);
context_chunks.push("The following fetched results are available:\n");
for context in &fetch_context {
context_chunks.push(&context.name);
for chunk in &context.text {
context_chunks.push(&chunk);
}
context_text.push('\n');
}
}

if !thread_context.is_empty() {
context_text.push_str("The following previous conversation threads are available\n");
for context in thread_context {
context_text.push_str(&context.name);
context_text.push('\n');
for chunk in context.text {
context_text.push_str(&chunk);
context_chunks.push("The following previous conversation threads are available:\n");
for context in &thread_context {
context_chunks.push(&context.name);
for chunk in &context.text {
context_chunks.push(&chunk);
}
context_text.push('\n');
}
}

if !context_text.is_empty() {
message.content.push(MessageContent::Text(context_text));
debug_assert!(
context_chunks.len() == capacity,
"attach_context_message calculated capacity of {}, but length was {}",
capacity,
context_chunks.len()
);

if !context_chunks.is_empty() {
message
.content
.push(MessageContent::Text(context_chunks.join("\n")));
}
}
13 changes: 6 additions & 7 deletions crates/assistant2/src/context_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -462,13 +462,12 @@ fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {

buffer.push_str("```\n");

if buffer.len() > capacity {
log::error!(
"to_fenced_codeblock calculated capacity {} but length was {}",
capacity,
buffer.len()
);
}
debug_assert!(
buffer.len() == capacity - 1 || buffer.len() == capacity,
"to_fenced_codeblock calculated capacity of {}, but length was {}",
capacity,
buffer.len(),
);

buffer.into()
}
Expand Down

0 comments on commit 8ccb33d

Please sign in to comment.