token-counter実装
This commit is contained in:
+88
-136
@@ -4,9 +4,9 @@
|
||||
//! their `summary`. This reclaims tokens while preserving the "what
|
||||
//! happened" trail.
|
||||
//!
|
||||
//! Pruning is **conditional**: it only fires when the estimated token
|
||||
//! savings exceed [`PruneConfig::min_savings`], avoiding unnecessary
|
||||
//! KV-cache invalidation.
|
||||
//! このモジュールは pure な「候補抽出」と「適用」だけを提供する。
|
||||
//! `min_savings` 判定や savings 推定はこの crate には置かず、上位層
|
||||
//! （`pod::prune_hook` など）が usage 履歴ベースのトークン会計と組み合わせて行う。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -20,17 +20,19 @@ pub struct PruneConfig {
|
||||
#[serde(default = "default_protected_turns")]
|
||||
pub protected_turns: usize,
|
||||
|
||||
/// Minimum estimated token savings required to actually prune.
|
||||
/// If the prunable content is smaller than this, we skip to
|
||||
/// avoid pointless KV-cache invalidation.
|
||||
/// Minimum token savings required to actually prune. If the prunable
|
||||
/// content is smaller than this, the caller should skip to avoid
|
||||
/// pointless KV-cache invalidation. The unit is tokens; the caller
|
||||
/// is responsible for measuring savings via a usage-history-aware
|
||||
/// estimator and comparing against this threshold.
|
||||
#[serde(default = "default_min_savings")]
|
||||
pub min_savings: usize,
|
||||
pub min_savings: u64,
|
||||
}
|
||||
|
||||
fn default_protected_turns() -> usize {
|
||||
3
|
||||
}
|
||||
fn default_min_savings() -> usize {
|
||||
fn default_min_savings() -> u64 {
|
||||
4096
|
||||
}
|
||||
|
||||
@@ -43,18 +45,11 @@ impl Default for PruneConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a prune operation.
|
||||
/// Result of [`apply_prune`].
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PruneResult {
|
||||
/// Number of items whose `content` was set to `None`.
|
||||
pub pruned_count: usize,
|
||||
/// Estimated tokens reclaimed.
|
||||
pub estimated_savings: usize,
|
||||
}
|
||||
|
||||
/// Estimate the token count of a string (rough: chars / 4).
|
||||
fn estimate_tokens(s: &str) -> usize {
|
||||
s.len() / 4
|
||||
}
|
||||
|
||||
/// Find indices where each "turn" begins.
|
||||
@@ -70,59 +65,45 @@ fn find_turn_starts(items: &[Item]) -> Vec<usize> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Conditionally prune old tool-result content from `items`.
|
||||
/// Indices of `Item::ToolResult { content: Some(_), .. }` that lie outside
|
||||
/// the last `protected_turns` turns. Pure: does not mutate `items`.
|
||||
///
|
||||
/// Returns `None` if pruning was skipped (not enough savings or not
|
||||
/// enough turns). Returns `Some(PruneResult)` if items were modified.
|
||||
///
|
||||
/// # Algorithm
|
||||
///
|
||||
/// 1. Identify turn boundaries (user-message positions).
|
||||
/// 2. Compute the protection boundary: items before the last
|
||||
/// `protected_turns` turns are candidates.
|
||||
/// 3. Sum the estimated token savings from prunable `content` fields.
|
||||
/// 4. If savings < `min_savings`, skip.
|
||||
/// 5. Otherwise, set `content = None` on each candidate.
|
||||
pub fn prune(items: &mut [Item], config: &PruneConfig) -> Option<PruneResult> {
|
||||
/// Returns an empty vector when there are too few turns or no prunable
|
||||
/// candidates.
|
||||
pub fn prunable_indices(items: &[Item], protected_turns: usize) -> Vec<usize> {
|
||||
let turn_starts = find_turn_starts(items);
|
||||
|
||||
// Not enough turns to have anything outside the protected window.
|
||||
if turn_starts.len() <= config.protected_turns {
|
||||
return None;
|
||||
if turn_starts.len() <= protected_turns {
|
||||
return Vec::new();
|
||||
}
|
||||
let boundary = turn_starts[turn_starts.len() - protected_turns];
|
||||
items[..boundary]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, item)| match item {
|
||||
Item::ToolResult {
|
||||
content: Some(_), ..
|
||||
} => Some(i),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// Everything before this index is a prune candidate.
|
||||
let boundary = turn_starts[turn_starts.len() - config.protected_turns];
|
||||
|
||||
// Collect prunable indices and total savings.
|
||||
let mut total_savings: usize = 0;
|
||||
let mut prunable: Vec<usize> = Vec::new();
|
||||
|
||||
for (i, item) in items[..boundary].iter().enumerate() {
|
||||
if let Item::ToolResult {
|
||||
content: Some(c), ..
|
||||
} = item
|
||||
{
|
||||
total_savings += estimate_tokens(c);
|
||||
prunable.push(i);
|
||||
}
|
||||
}
|
||||
|
||||
if prunable.is_empty() || total_savings < config.min_savings {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Apply: drop content, keep summary.
|
||||
for &i in &prunable {
|
||||
/// Set `content = None` on each item at `indices`. Returns the number
|
||||
/// of items that were actually modified (already-pruned items are
|
||||
/// counted as 0).
|
||||
pub fn apply_prune(items: &mut [Item], indices: &[usize]) -> PruneResult {
|
||||
let mut count = 0;
|
||||
for &i in indices {
|
||||
if let Item::ToolResult { content, .. } = &mut items[i] {
|
||||
*content = None;
|
||||
if content.is_some() {
|
||||
*content = None;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(PruneResult {
|
||||
pruned_count: prunable.len(),
|
||||
estimated_savings: total_savings,
|
||||
})
|
||||
PruneResult {
|
||||
pruned_count: count,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -148,53 +129,48 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_too_few_turns() {
|
||||
let mut items = make_history(&[
|
||||
fn no_candidates_when_too_few_turns() {
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("summary1", Some("big content here"))]),
|
||||
("turn2", vec![("summary2", Some("more content"))]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 3,
|
||||
min_savings: 0,
|
||||
};
|
||||
assert!(prune(&mut items, &config).is_none());
|
||||
assert!(prunable_indices(&items, 3).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_savings_below_threshold() {
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s", Some("tiny"))]), // ~1 token
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
fn candidates_in_unprotected_turns() {
|
||||
let big = "x".repeat(4096 * 4);
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some(&big))]),
|
||||
("turn3", vec![("s3", Some("keep me"))]),
|
||||
("turn4", vec![("s4", Some("keep me too"))]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 2,
|
||||
min_savings: 9999,
|
||||
};
|
||||
assert!(prune(&mut items, &config).is_none());
|
||||
let candidates = prunable_indices(&items, 2);
|
||||
assert_eq!(candidates.len(), 2);
|
||||
// 候補は turn1 と turn2 の ToolResult のみ
|
||||
for &i in &candidates {
|
||||
if let Item::ToolResult { summary, .. } = &items[i] {
|
||||
assert!(summary == "s1" || summary == "s2");
|
||||
} else {
|
||||
panic!("non tool-result selected");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_old_content() {
|
||||
// 4 turns. protected_turns=2 → turns 1-2 are candidates.
|
||||
let big = "x".repeat(4096 * 4); // ~4096 tokens
|
||||
fn apply_drops_content_only() {
|
||||
let big = "x".repeat(64);
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some(&big))]),
|
||||
("turn3", vec![("s3", Some("keep me"))]),
|
||||
("turn4", vec![("s4", Some("keep me too"))]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 2,
|
||||
min_savings: 1000,
|
||||
};
|
||||
|
||||
let result = prune(&mut items, &config).expect("should prune");
|
||||
let candidates = prunable_indices(&items, 2);
|
||||
let result = apply_prune(&mut items, &candidates);
|
||||
assert_eq!(result.pruned_count, 2);
|
||||
assert!(result.estimated_savings >= 8000);
|
||||
|
||||
// Verify: pruned items have content=None, protected items keep content.
|
||||
for item in &items {
|
||||
if let Item::ToolResult {
|
||||
summary, content, ..
|
||||
@@ -210,73 +186,49 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idempotent() {
|
||||
let big = "x".repeat(4096 * 4);
|
||||
fn apply_is_idempotent() {
|
||||
let big = "x".repeat(64);
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 2,
|
||||
min_savings: 100,
|
||||
};
|
||||
let first_indices = prunable_indices(&items, 2);
|
||||
assert_eq!(apply_prune(&mut items, &first_indices).pruned_count, 1);
|
||||
|
||||
let first = prune(&mut items, &config).expect("first prune");
|
||||
assert_eq!(first.pruned_count, 1);
|
||||
|
||||
// Second call: nothing left to prune.
|
||||
assert!(prune(&mut items, &config).is_none());
|
||||
// 2 周目: 候補は (まだ) いるかもしれないが、すでに content=None なので
|
||||
// apply_prune は 0 件と数える。
|
||||
let second_indices = prunable_indices(&items, 2);
|
||||
assert!(second_indices.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn already_pruned_items_skipped() {
|
||||
// Items that already have content=None are not counted as savings.
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", None)]), // already pruned
|
||||
fn already_pruned_items_excluded_from_candidates() {
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", None)]), // already pruned (content=None)
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 2,
|
||||
min_savings: 0, // Even with threshold 0, no savings means no prune
|
||||
};
|
||||
|
||||
assert!(prune(&mut items, &config).is_none());
|
||||
assert!(prunable_indices(&items, 2).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protected_turns_boundary_exact() {
|
||||
// 3 turns with protected_turns=2:
|
||||
// Turn 1 content should be pruned, turns 2-3 protected.
|
||||
let big = "x".repeat(4096 * 4);
|
||||
let mut items = make_history(&[
|
||||
// 3 turns with protected_turns=2: only turn 1 is a candidate.
|
||||
let big = "x".repeat(64);
|
||||
let items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![("s2", Some("protected"))]),
|
||||
("turn3", vec![("s3", Some("also protected"))]),
|
||||
]);
|
||||
let config = PruneConfig {
|
||||
protected_turns: 2,
|
||||
min_savings: 100,
|
||||
};
|
||||
|
||||
let result = prune(&mut items, &config).expect("should prune turn1");
|
||||
assert_eq!(result.pruned_count, 1);
|
||||
|
||||
// Verify s1 pruned, s2 and s3 intact.
|
||||
for item in &items {
|
||||
if let Item::ToolResult {
|
||||
summary, content, ..
|
||||
} = item
|
||||
{
|
||||
match summary.as_str() {
|
||||
"s1" => assert!(content.is_none()),
|
||||
"s2" | "s3" => assert!(content.is_some()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
let candidates = prunable_indices(&items, 2);
|
||||
assert_eq!(candidates.len(), 1);
|
||||
if let Item::ToolResult { summary, .. } = &items[candidates[0]] {
|
||||
assert_eq!(summary, "s1");
|
||||
} else {
|
||||
panic!("expected ToolResult at candidate index");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user