Restore trim_conversation: dedup memories, evict to budget, snap boundary
Ported the old trim_entries logic to the new AstNode types: - Phase 1: Dedup Memory nodes by key (keep last), drop DMN entries - Phase 2: While over budget, evict lowest-scored memory (if memories > 50% of conv tokens) or oldest conversation entry - Phase 3: Snap to User message boundary at start Called from compact() which runs on startup and on /compact. Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
parent
7237baba11
commit
bbffc2213e
2 changed files with 84 additions and 0 deletions
|
|
@ -795,6 +795,88 @@ impl ContextState {
|
||||||
self.section_mut(section).clear();
|
self.section_mut(section).clear();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Dedup and trim conversation entries to fit within the context budget.
|
||||||
|
///
|
||||||
|
/// Phase 1: Drop duplicate memories (keep last) and DMN entries.
|
||||||
|
/// Phase 2: While over budget, drop lowest-scored memory (if memories
|
||||||
|
/// are > 50% of conversation tokens) or oldest conversation entry.
|
||||||
|
/// Phase 3: Snap to user message boundary at start.
|
||||||
|
pub fn trim_conversation(&mut self) {
|
||||||
|
let max_tokens = context_budget_tokens();
|
||||||
|
let fixed = self.system.iter().map(|n| n.tokens()).sum::<usize>()
|
||||||
|
+ self.identity.iter().map(|n| n.tokens()).sum::<usize>()
|
||||||
|
+ self.journal.iter().map(|n| n.tokens()).sum::<usize>();
|
||||||
|
|
||||||
|
// Phase 1: dedup memories by key (keep last), drop DMN
|
||||||
|
let mut seen_keys: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
|
||||||
|
let mut drop = std::collections::HashSet::new();
|
||||||
|
|
||||||
|
for (i, node) in self.conversation.iter().enumerate() {
|
||||||
|
if let AstNode::Leaf(leaf) = node {
|
||||||
|
match leaf.body() {
|
||||||
|
NodeBody::Dmn(_) => { drop.insert(i); }
|
||||||
|
NodeBody::Memory { key, .. } => {
|
||||||
|
if let Some(prev) = seen_keys.insert(key.clone(), i) {
|
||||||
|
drop.insert(prev);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !drop.is_empty() {
|
||||||
|
let mut i = 0;
|
||||||
|
self.conversation.retain(|_| { let keep = !drop.contains(&i); i += 1; keep });
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 2: while over budget, evict
|
||||||
|
loop {
|
||||||
|
let total: usize = self.conversation.iter().map(|n| n.tokens()).sum();
|
||||||
|
if fixed + total <= max_tokens { break; }
|
||||||
|
let mt: usize = self.conversation.iter()
|
||||||
|
.filter(|n| matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
|
||||||
|
.map(|n| n.tokens()).sum();
|
||||||
|
let ct = total - mt;
|
||||||
|
|
||||||
|
if mt > ct {
|
||||||
|
// Memories > 50% — drop lowest-scored
|
||||||
|
if let Some(i) = self.lowest_scored_memory() {
|
||||||
|
self.conversation.remove(i);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Drop oldest non-memory entry
|
||||||
|
if let Some(i) = self.conversation.iter().position(|n|
|
||||||
|
!matches!(n, AstNode::Leaf(l) if matches!(l.body(), NodeBody::Memory { .. })))
|
||||||
|
{
|
||||||
|
self.conversation.remove(i);
|
||||||
|
} else {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 3: snap to user message boundary
|
||||||
|
while let Some(first) = self.conversation.first() {
|
||||||
|
if matches!(first, AstNode::Branch { role: Role::User, .. }) { break; }
|
||||||
|
self.conversation.remove(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn lowest_scored_memory(&self) -> Option<usize> {
|
||||||
|
self.conversation.iter().enumerate()
|
||||||
|
.filter_map(|(i, n)| {
|
||||||
|
if let AstNode::Leaf(l) = n {
|
||||||
|
if let NodeBody::Memory { score: Some(s), .. } = l.body() {
|
||||||
|
return Some((i, *s));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
})
|
||||||
|
.min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
||||||
|
.map(|(i, _)| i)
|
||||||
|
}
|
||||||
|
|
||||||
/// Push a child node into a branch at `index` in `section`.
|
/// Push a child node into a branch at `index` in `section`.
|
||||||
pub fn push_child(&mut self, section: Section, index: usize, child: AstNode) {
|
pub fn push_child(&mut self, section: Section, index: usize, child: AstNode) {
|
||||||
let node = &mut self.section_mut(section)[index];
|
let node = &mut self.section_mut(section)[index];
|
||||||
|
|
|
||||||
|
|
@ -555,6 +555,8 @@ impl Agent {
|
||||||
|
|
||||||
self.load_startup_journal().await;
|
self.load_startup_journal().await;
|
||||||
|
|
||||||
|
self.context.lock().await.trim_conversation();
|
||||||
|
|
||||||
let mut st = self.state.lock().await;
|
let mut st = self.state.lock().await;
|
||||||
st.generation += 1;
|
st.generation += 1;
|
||||||
st.last_prompt_tokens = 0;
|
st.last_prompt_tokens = 0;
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue