measure_budget: count from typed sources, not message scanning
Identity tokens from system_prompt + personality vec. Journal from journal entries vec. Memory from loaded_nodes. Conversation is the remainder. No string prefix matching. Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
parent
4bdc7ae112
commit
4580f5dade
1 changed files with 8 additions and 27 deletions
|
|
@ -185,35 +185,16 @@ impl Agent {
|
||||||
fn measure_budget(&mut self) {
|
fn measure_budget(&mut self) {
|
||||||
let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
|
let count = |s: &str| self.tokenizer.encode_with_special_tokens(s).len();
|
||||||
|
|
||||||
let mut id_tokens: usize = 0;
|
let id_tokens = count(&self.context.system_prompt)
|
||||||
|
+ self.context.personality.iter()
|
||||||
|
.map(|(_, content)| count(content)).sum::<usize>();
|
||||||
let jnl_tokens: usize = self.context.journal.iter()
|
let jnl_tokens: usize = self.context.journal.iter()
|
||||||
.map(|e| count(&e.content))
|
.map(|e| count(&e.content)).sum();
|
||||||
.sum();
|
|
||||||
let mut conv_tokens: usize = 0;
|
|
||||||
let mut in_conversation = false;
|
|
||||||
|
|
||||||
for msg in &self.messages {
|
|
||||||
let tokens = crate::agent::context::msg_token_count(&self.tokenizer, msg);
|
|
||||||
|
|
||||||
if in_conversation {
|
|
||||||
conv_tokens += tokens;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
if in_conversation {
|
|
||||||
conv_tokens += tokens;
|
|
||||||
} else if msg.role == Role::System || (!in_conversation && conv_tokens == 0) {
|
|
||||||
id_tokens += tokens;
|
|
||||||
} else {
|
|
||||||
in_conversation = true;
|
|
||||||
conv_tokens += tokens;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Memory = nodes loaded during the session via tool calls
|
|
||||||
let mem_tokens: usize = self.context.loaded_nodes.iter()
|
let mem_tokens: usize = self.context.loaded_nodes.iter()
|
||||||
.map(|node| count(&node.render()))
|
.map(|node| count(&node.render())).sum();
|
||||||
.sum();
|
let total: usize = self.messages.iter()
|
||||||
|
.map(|m| crate::agent::context::msg_token_count(&self.tokenizer, m)).sum();
|
||||||
|
let conv_tokens = total.saturating_sub(id_tokens + jnl_tokens);
|
||||||
|
|
||||||
self.context_budget = ContextBudget {
|
self.context_budget = ContextBudget {
|
||||||
identity_tokens: id_tokens,
|
identity_tokens: id_tokens,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue