WIP: ContextEntry/ContextSection data structures for incremental token counting

New types — not yet wired to callers:

- ContextEntry: wraps ConversationEntry with cached token count and
  timestamp
- ContextSection: named group of entries with cached token total.
  Private entries/tokens, read via entries()/tokens().
  Mutation via push(entry), set(index, entry), del(index).
- ContextState: system/identity/journal/conversation sections + working_stack
- ConversationEntry::System variant for system prompt entries

Token counting happens once at push time. Sections maintain their
totals incrementally via push/set/del. No more recomputing from
scratch on every budget check.

Does not compile — callers need updating.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-07 20:15:31 -04:00
parent 776ac527f1
commit 62996e27d7
10 changed files with 450 additions and 403 deletions

View file

@ -16,7 +16,7 @@
use crate::agent::api::ApiClient;
use crate::agent::api::*;
use crate::agent::context::{ConversationEntry, ContextState};
use crate::agent::context::{ConversationEntry, ContextEntry, ContextState};
const SCORE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
@ -39,19 +39,22 @@ fn build_messages(
range: std::ops::Range<usize>,
filter: Filter,
) -> Vec<serde_json::Value> {
let mut msgs = vec![
serde_json::json!({"role": "system", "content": &context.system_prompt}),
];
let mut msgs = Vec::new();
for e in context.system.entries() {
msgs.push(serde_json::json!({"role": "system", "content": e.entry.message().content_text()}));
}
let ctx = context.render_context_message();
if !ctx.is_empty() {
msgs.push(serde_json::json!({"role": "user", "content": ctx}));
}
let entries = context.conversation.entries();
for i in range {
let entry = &context.entries[i];
let ce = &entries[i];
let entry = &ce.entry;
let skip = match &filter {
Filter::None => false,
Filter::SkipIndex(idx) => i == *idx,
Filter::SkipKey(key) => matches!(entry, ConversationEntry::Memory { key: k, .. } if k == key),
Filter::SkipKey(key) => matches!(entry, ConversationEntry::Memory { key: k, .. } if k == *key),
Filter::SkipAllMemories => entry.is_memory(),
};
if skip { continue; }
@ -175,16 +178,16 @@ pub async fn score_memories(
context: &ContextState,
client: &ApiClient,
) -> anyhow::Result<MemoryScore> {
let mut memory_keys: Vec<String> = context.entries.iter()
.filter_map(|e| match e {
let mut memory_keys: Vec<String> = context.conversation.entries().iter()
.filter_map(|ce| match &ce.entry {
ConversationEntry::Memory { key, .. } => Some(key.clone()),
_ => None,
})
.collect();
memory_keys.dedup();
let response_indices: Vec<usize> = context.entries.iter().enumerate()
.filter(|(_, e)| e.message().role == Role::Assistant)
let response_indices: Vec<usize> = context.conversation.entries().iter().enumerate()
.filter(|(_, ce)| ce.entry.message().role == Role::Assistant)
.map(|(i, _)| i)
.collect();
@ -198,7 +201,7 @@ pub async fn score_memories(
let http = http_client();
let range = 0..context.entries.len();
let range = 0..context.conversation.entries().len();
let baseline = call_score(&http, client, &build_messages(context, range.clone(), Filter::None)).await?;
@ -242,10 +245,10 @@ pub async fn score_memories(
/// Find the entry index after `start` that contains the Nth assistant response.
/// Returns (end_index, true) if N responses were found, (entries.len(), false) if not.
fn nth_response_end(entries: &[ConversationEntry], start: usize, n: usize) -> (usize, bool) {
fn nth_response_end(entries: &[ContextEntry], start: usize, n: usize) -> (usize, bool) {
let mut count = 0;
for i in start..entries.len() {
if entries[i].message().role == Role::Assistant {
if entries[i].entry.message().role == Role::Assistant {
count += 1;
if count >= n { return (i + 1, true); }
}
@ -267,16 +270,17 @@ pub async fn score_memory(
) -> anyhow::Result<f64> {
const RESPONSE_WINDOW: usize = 50;
let first_pos = match context.entries.iter().position(|e| {
matches!(e, ConversationEntry::Memory { key: k, .. } if k == key)
let entries = context.conversation.entries();
let first_pos = match entries.iter().position(|ce| {
matches!(&ce.entry, ConversationEntry::Memory { key: k, .. } if k == key)
}) {
Some(p) => p,
None => return Ok(0.0),
};
let (end, _) = nth_response_end(&context.entries, first_pos, RESPONSE_WINDOW);
let (end, _) = nth_response_end(entries, first_pos, RESPONSE_WINDOW);
let range = first_pos..end;
if !context.entries[range.clone()].iter().any(|e| e.message().role == Role::Assistant) {
if !entries[range.clone()].iter().any(|ce| ce.entry.message().role == Role::Assistant) {
return Ok(0.0);
}
@ -310,8 +314,8 @@ pub async fn score_memories_incremental(
let store = crate::hippocampus::store::Store::load().unwrap_or_default();
for (i, entry) in context.entries.iter().enumerate() {
if let ConversationEntry::Memory { key, .. } = entry {
for (i, ce) in context.conversation.entries().iter().enumerate() {
if let ConversationEntry::Memory { key, .. } = &ce.entry {
if !seen.insert(key.clone()) { continue; }
let last_scored = store.nodes.get(key.as_str())
.map(|n| n.last_scored)
@ -328,18 +332,18 @@ pub async fn score_memories_incremental(
let http = http_client();
let mut results = Vec::new();
let total_entries = context.entries.len();
let total_entries = context.conversation.entries().len();
let first_quarter = total_entries / 4;
for (pos, key, _) in &candidates {
let (end, full_window) = nth_response_end(&context.entries, *pos, response_window);
let (end, full_window) = nth_response_end(context.conversation.entries(), *pos, response_window);
// Skip memories without a full window, unless they're in the
// first quarter of the conversation (always score those).
if !full_window && *pos >= first_quarter {
continue;
}
let range = *pos..end;
if !context.entries[range.clone()].iter().any(|e| e.message().role == Role::Assistant) {
if !context.conversation.entries()[range.clone()].iter().any(|ce| ce.entry.message().role == Role::Assistant) {
continue;
}
@ -378,10 +382,10 @@ pub async fn score_finetune(
count: usize,
client: &ApiClient,
) -> anyhow::Result<Vec<(usize, f64)>> {
let range = context.entries.len().saturating_sub(count)..context.entries.len();
let range = context.conversation.entries().len().saturating_sub(count)..context.conversation.entries().len();
let response_positions: Vec<usize> = range.clone()
.filter(|&i| context.entries[i].message().role == Role::Assistant)
.filter(|&i| context.conversation.entries()[i].entry.message().role == Role::Assistant)
.collect();
if response_positions.is_empty() {
return Ok(Vec::new());