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

agent: Include context with first message that introduced it (#27925)

Browse source 
We were including the context at the end which meant it never got
cached. We'll now include it with the first message that introduced it
so it's cached as long as it doesn't change.

This is an improvement, but we probably still need to think of ways to
optimize caching for cases where files in context change.

Release Notes:

- N/A
This commit is contained in:
Agus Zubiaga 2025-04-02 12:14:48 -03:00 • committed by GitHub
parent 7a54dd7190
commit 2eed94ff23
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 36 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

38
crates/agent/src/thread.rs
View file

@ -881,13 +881,9 @@ impl Thread {
log::error!("system_prompt_context not set.")
}
let mut referenced_context_ids = HashSet::default();
let mut added_context_ids = HashSet::<ContextId>::default();
for message in &self.messages {
if let Some(context_ids) = self.context_by_message.get(&message.id) {
referenced_context_ids.extend(context_ids);
}
let mut request_message = LanguageModelRequestMessage {
role: message.role,
content: Vec::new(),
@ -907,6 +903,23 @@ impl Thread {
}
}
// Attach context to this message if it's the first to reference it
if let Some(context_ids) = self.context_by_message.get(&message.id) {
let new_context_ids: Vec<_> = context_ids
.iter()
.filter(|id| !added_context_ids.contains(id))
.collect();
if !new_context_ids.is_empty() {
let referenced_context = new_context_ids
.iter()
.filter_map(|context_id| self.context.get(*context_id));
attach_context_to_message(&mut request_message, referenced_context, cx);
added_context_ids.extend(context_ids.iter());
}
}
if !message.segments.is_empty() {
request_message
.content
@ -933,21 +946,6 @@ impl Thread {
message.cache = index == breakpoint_index;
}
if !referenced_context_ids.is_empty() {
let mut context_message = LanguageModelRequestMessage {
role: Role::User,
content: Vec::new(),
cache: false,
};
let referenced_context = referenced_context_ids
.into_iter()
.filter_map(|context_id| self.context.get(context_id));
attach_context_to_message(&mut context_message, referenced_context, cx);
request.messages.push(context_message);
}
self.attached_tracked_files_state(&mut request.messages, cx);
request