pruneで用いるトークン計算の改善
This commit is contained in:
@@ -1,12 +1,18 @@
|
||||
//! Conditional Prune algorithm for context window management.
|
||||
//! Prune — context projection for old tool-result content.
|
||||
//!
|
||||
//! Removes `content` from old [`Item::ToolResult`] entries, leaving only
|
||||
//! their `summary`. This reclaims tokens while preserving the "what
|
||||
//! happened" trail.
|
||||
//! LLM 送信時のコンテキストから古い [`Item::ToolResult`] の `content` を
|
||||
//! 省略して、コンテキスト窓のトークンを回収する。`summary` は残すので
|
||||
//! 「何が起きたか」の痕跡は保たれる。
|
||||
//!
|
||||
//! このモジュールは pure な「候補抽出」と「適用」だけを提供する。
|
||||
//! `min_savings` 判定や savings 推定はこの crate には置かず、上位層
|
||||
//! （`pod::prune_hook` など）が usage 履歴ベースのトークン会計と組み合わせて行う。
|
||||
//! # 設計方針
|
||||
//!
|
||||
//! Prune は **コンテキスト射影** であり、history の変換ではない。
|
||||
//! この crate が提供するのは pure な候補抽出 [`prunable_indices`] のみで、
|
||||
//! 射影の適用は上位層（`pod::prune_hook` 等）が LLM に送る一時コンテキスト
|
||||
//! に対してだけ行う。Worker の永続履歴は決して変更されない。
|
||||
//!
|
||||
//! `min_savings` 判定や savings 推定もこの crate には置かず、上位層が
|
||||
//! usage 履歴ベースのトークン会計と組み合わせて行う。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -45,13 +51,6 @@ impl Default for PruneConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
/// Find indices where each "turn" begins.
|
||||
///
|
||||
/// A turn starts at every user message. Returns the indices of those
|
||||
@@ -88,24 +87,6 @@ pub fn prunable_indices(items: &[Item], protected_turns: usize) -> Vec<usize> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 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] {
|
||||
if content.is_some() {
|
||||
*content = None;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
PruneResult {
|
||||
pruned_count: count,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -158,51 +139,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
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 candidates = prunable_indices(&items, 2);
|
||||
let result = apply_prune(&mut items, &candidates);
|
||||
assert_eq!(result.pruned_count, 2);
|
||||
|
||||
for item in &items {
|
||||
if let Item::ToolResult {
|
||||
summary, content, ..
|
||||
} = item
|
||||
{
|
||||
if summary == "s1" || summary == "s2" {
|
||||
assert!(content.is_none(), "old content should be pruned");
|
||||
} else {
|
||||
assert!(content.is_some(), "protected content should remain");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
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 first_indices = prunable_indices(&items, 2);
|
||||
assert_eq!(apply_prune(&mut items, &first_indices).pruned_count, 1);
|
||||
|
||||
// 2 周目: 候補は (まだ) いるかもしれないが、すでに content=None なので
|
||||
// apply_prune は 0 件と数える。
|
||||
let second_indices = prunable_indices(&items, 2);
|
||||
assert!(second_indices.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn already_pruned_items_excluded_from_candidates() {
|
||||
let items = make_history(&[
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
//! PruneHook — applies conditional pruning before each LLM request.
|
||||
//! PruneHook — projects the LLM request context before each call.
|
||||
//!
|
||||
//! Prune は **コンテキスト射影** として実装する。`PreLlmRequest` hook に
|
||||
//! 渡される `context: &mut Vec<Item>` は Worker が毎 turn 冒頭で history を
|
||||
//! clone した一時配列 (`worker.rs:701`)。ここで ToolResult.content を省いても
|
||||
//! Worker の永続履歴には影響しない。`prunable_indices` で候補を抽出し、
|
||||
//! `min_savings` を満たせば content を `None` に射影する。
|
||||
//!
|
||||
//! Wraps the pure `prune` API from `llm-worker` as a [`Hook<PreLlmRequest>`].
|
||||
//! `min_savings` の判定は usage 履歴ベースのトークン会計
|
||||
//! ([`crate::token_counter::savings_for_drop_impl`]) で行う。
|
||||
|
||||
@@ -9,7 +14,7 @@ use std::sync::{Arc, Mutex};
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::interceptor::PreRequestAction;
|
||||
use llm_worker::prune::{PruneConfig, apply_prune, prunable_indices};
|
||||
use llm_worker::prune::{PruneConfig, prunable_indices};
|
||||
use session_store::UsageRecord;
|
||||
use tracing::debug;
|
||||
|
||||
@@ -66,13 +71,23 @@ impl Hook<PreLlmRequest> for PruneHook {
|
||||
return PreRequestAction::Continue;
|
||||
}
|
||||
|
||||
let result = apply_prune(context, &candidates);
|
||||
if result.pruned_count > 0 {
|
||||
// 射影: context (= history の clone) 上の対象 ToolResult だけ content を
|
||||
// drop する。Worker の永続履歴は別インスタンスなので影響を受けない。
|
||||
let mut projected = 0usize;
|
||||
for &i in &candidates {
|
||||
if let Item::ToolResult { content, .. } = &mut context[i] {
|
||||
if content.is_some() {
|
||||
*content = None;
|
||||
projected += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if projected > 0 {
|
||||
debug!(
|
||||
pruned = result.pruned_count,
|
||||
pruned = projected,
|
||||
estimated_savings_tokens = savings.tokens,
|
||||
source = ?savings.source,
|
||||
"Pruned old tool-result content"
|
||||
"Projected old tool-result content out of request context"
|
||||
);
|
||||
}
|
||||
PreRequestAction::Continue
|
||||
|
||||
Reference in New Issue
Block a user