WIP: trim_entries dedup, context_window rename, compact simplification
Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
parent
809679b6ce
commit
d419587c1b
3 changed files with 58 additions and 53 deletions
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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(),
|
||||||
|
|
|
||||||
|
|
@ -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.
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue