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();
|
||||
}
|
||||
|
||||
/// 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`.
|
||||
pub fn push_child(&mut self, section: Section, index: usize, child: AstNode) {
|
||||
let node = &mut self.section_mut(section)[index];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue