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) {
|
||||
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()
|
||||
.map(|e| count(&e.content))
|
||||
.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
|
||||
.map(|e| count(&e.content)).sum();
|
||||
let mem_tokens: usize = self.context.loaded_nodes.iter()
|
||||
.map(|node| count(&node.render()))
|
||||
.sum();
|
||||
.map(|node| count(&node.render())).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 {
|
||||
identity_tokens: id_tokens,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue