WIP: trim_entries dedup, context_window rename, compact simplification

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-02 15:58:03 -04:00
parent 809679b6ce
commit d419587c1b
3 changed files with 58 additions and 53 deletions

View file

@ -909,23 +909,24 @@ impl Agent {
self.do_compact(); self.do_compact();
} }
/// Internal compaction — rebuilds context window from current messages. /// Dedup memory entries, trim to fit, reload journal for new time range.
fn do_compact(&mut self) { fn trim_and_reload(&mut self, entries: &[ConversationEntry]) {
let conversation: Vec<Message> = self.context.entries.iter() self.context.entries = crate::thought::context::trim_entries(
.map(|e| e.api_message().clone()).collect();
let messages = crate::thought::context::trim_conversation(
&self.context, &self.context,
&conversation, entries,
&self.client.model,
&self.tokenizer, &self.tokenizer,
); );
self.context.entries = messages.into_iter() self.load_startup_journal();
.map(ConversationEntry::Message).collect();
self.last_prompt_tokens = 0; self.last_prompt_tokens = 0;
self.publish_context_state(); self.publish_context_state();
} }
/// Internal compaction — dedup memory entries and trim to fit.
fn do_compact(&mut self) {
let entries = self.context.entries.clone();
self.trim_and_reload(&entries);
}
/// Emergency compaction using stored config — called on context overflow. /// Emergency compaction using stored config — called on context overflow.
fn emergency_compact(&mut self) { fn emergency_compact(&mut self) {
self.do_compact(); self.do_compact();
@ -964,32 +965,14 @@ impl Agent {
} }
}; };
// Filter out system messages, keep everything else (including Memory entries) // Filter out system messages, dedup memory, trim to fit
let entries: Vec<ConversationEntry> = entries let entries: Vec<ConversationEntry> = entries
.into_iter() .into_iter()
.filter(|e| e.message().role != Role::System) .filter(|e| e.message().role != Role::System)
.collect(); .collect();
self.trim_and_reload(&entries);
// Trim to fit context budget
let n = entries.len();
let conversation: Vec<Message> = entries.iter()
.map(|e| e.api_message().clone()).collect();
let trimmed = crate::thought::context::trim_conversation(
&self.context,
&conversation,
&self.client.model,
&self.tokenizer,
);
// Keep only the entries that survived trimming (by count from the end)
let keep = trimmed.len();
self.context.entries = entries.into_iter()
.skip(n.saturating_sub(keep))
.collect();
dbglog!("[restore] {} entries, journal: {} entries", dbglog!("[restore] {} entries, journal: {} entries",
self.context.entries.len(), self.context.journal.len()); self.context.entries.len(), self.context.journal.len());
self.last_prompt_tokens = 0;
self.publish_context_state();
true true
} }

View file

@ -40,14 +40,14 @@ use poc_memory::agent::ui_channel::{ContextInfo, StatusInfo, StreamTarget, UiMes
/// Hard compaction threshold — context is rebuilt immediately. /// Hard compaction threshold — context is rebuilt immediately.
/// Uses config percentage of model context window. /// Uses config percentage of model context window.
fn compaction_threshold(model: &str, app: &AppConfig) -> u32 { fn compaction_threshold(app: &AppConfig) -> u32 {
(poc_memory::thought::context::model_context_window(model) as u32) * app.compaction.hard_threshold_pct / 100 (poc_memory::thought::context::context_window() as u32) * app.compaction.hard_threshold_pct / 100
} }
/// Soft threshold — nudge the model to journal before compaction. /// Soft threshold — nudge the model to journal before compaction.
/// Fires once; the hard threshold handles the actual rebuild. /// Fires once; the hard threshold handles the actual rebuild.
fn pre_compaction_threshold(model: &str, app: &AppConfig) -> u32 { fn pre_compaction_threshold(app: &AppConfig) -> u32 {
(poc_memory::thought::context::model_context_window(model) as u32) * app.compaction.soft_threshold_pct / 100 (poc_memory::thought::context::context_window() as u32) * app.compaction.soft_threshold_pct / 100
} }
#[tokio::main] #[tokio::main]
@ -280,8 +280,8 @@ impl Session {
async fn check_compaction(&mut self) { async fn check_compaction(&mut self) {
let mut agent_guard = self.agent.lock().await; let mut agent_guard = self.agent.lock().await;
let tokens = agent_guard.last_prompt_tokens(); let tokens = agent_guard.last_prompt_tokens();
let hard = compaction_threshold(agent_guard.model(), &self.config.app); let hard = compaction_threshold(&self.config.app);
let soft = pre_compaction_threshold(agent_guard.model(), &self.config.app); let soft = pre_compaction_threshold(&self.config.app);
if tokens > hard { if tokens > hard {
let _ = self.ui_tx.send(UiMessage::Info(format!( let _ = self.ui_tx.send(UiMessage::Info(format!(
@ -452,7 +452,7 @@ impl Session {
let total_chars: usize = let total_chars: usize =
msgs.iter().map(|e| e.message().content_text().len()).sum(); msgs.iter().map(|e| e.message().content_text().len()).sum();
let prompt_tokens = agent.last_prompt_tokens(); let prompt_tokens = agent.last_prompt_tokens();
let threshold = compaction_threshold(agent.model(), &self.config.app); let threshold = compaction_threshold(&self.config.app);
let _ = self.ui_tx.send(UiMessage::Info(format!( let _ = self.ui_tx.send(UiMessage::Info(format!(
" {} messages, ~{} chars", " {} messages, ~{} chars",
msgs.len(), msgs.len(),

View file

@ -15,27 +15,49 @@ pub struct JournalEntry {
pub content: String, pub content: String,
} }
/// Look up a model's context window size in tokens. /// Context window size in tokens (from config).
pub fn model_context_window(_model: &str) -> usize { pub fn context_window() -> usize {
crate::config::get().api_context_window crate::config::get().api_context_window
} }
/// Context budget in tokens: 60% of the model's context window. /// Context budget in tokens: 60% of the model's context window.
fn context_budget_tokens(model: &str) -> usize { fn context_budget_tokens() -> usize {
model_context_window(model) * 60 / 100 context_window() * 60 / 100
} }
/// Trim conversation to fit within the context budget. /// Dedup and trim conversation entries to fit within the context budget.
/// Returns the trimmed conversation messages (oldest dropped first). ///
pub fn trim_conversation( /// 1. Dedup: if the same memory key appears multiple times, keep only
/// the latest render (drop the earlier Memory entry and its
/// corresponding assistant tool_call message).
/// 2. Trim: drop oldest entries until the conversation fits, snapping
/// to user message boundaries.
pub fn trim_entries(
context: &ContextState, context: &ContextState,
conversation: &[Message], entries: &[ConversationEntry],
model: &str,
tokenizer: &CoreBPE, tokenizer: &CoreBPE,
) -> Vec<Message> { ) -> Vec<ConversationEntry> {
let count = |s: &str| tokenizer.encode_with_special_tokens(s).len(); let count = |s: &str| tokenizer.encode_with_special_tokens(s).len();
let max_tokens = context_budget_tokens(model);
// --- Phase 1: dedup memory entries by key (keep last) ---
let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
for (i, entry) in entries.iter().enumerate() {
if let ConversationEntry::Memory { key, .. } = entry {
if let Some(prev) = seen_keys.insert(key.as_str(), i) {
drop_indices.insert(prev);
}
}
}
let deduped: Vec<ConversationEntry> = entries.iter().enumerate()
.filter(|(i, _)| !drop_indices.contains(i))
.map(|(_, e)| e.clone())
.collect();
// --- Phase 2: trim to fit context budget ---
let max_tokens = context_budget_tokens();
let identity_cost = count(&context.system_prompt) let identity_cost = count(&context.system_prompt)
+ context.personality.iter().map(|(_, c)| count(c)).sum::<usize>(); + context.personality.iter().map(|(_, c)| count(c)).sum::<usize>();
let journal_cost: usize = context.journal.iter().map(|e| count(&e.content)).sum(); let journal_cost: usize = context.journal.iter().map(|e| count(&e.content)).sum();
@ -45,23 +67,23 @@ pub fn trim_conversation(
.saturating_sub(journal_cost) .saturating_sub(journal_cost)
.saturating_sub(reserve); .saturating_sub(reserve);
let msg_costs: Vec<usize> = conversation.iter() let msg_costs: Vec<usize> = deduped.iter()
.map(|m| msg_token_count(tokenizer, m)).collect(); .map(|e| msg_token_count(tokenizer, e.message())).collect();
let total: usize = msg_costs.iter().sum(); let total: usize = msg_costs.iter().sum();
let mut skip = 0; let mut skip = 0;
let mut trimmed = total; let mut trimmed = total;
while trimmed > available && skip < conversation.len() { while trimmed > available && skip < deduped.len() {
trimmed -= msg_costs[skip]; trimmed -= msg_costs[skip];
skip += 1; skip += 1;
} }
// Walk forward to user message boundary // Walk forward to user message boundary
while skip < conversation.len() && conversation[skip].role != Role::User { while skip < deduped.len() && deduped[skip].message().role != Role::User {
skip += 1; skip += 1;
} }
conversation[skip..].to_vec() deduped[skip..].to_vec()
} }
/// Count the token footprint of a message using BPE tokenization. /// Count the token footprint of a message using BPE tokenization.