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:
parent
776ac527f1
commit
62996e27d7
10 changed files with 450 additions and 403 deletions
|
|
@ -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());
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue