pruneのトークン計算置き換え・Podに接続
This commit is contained in:
@@ -14,10 +14,19 @@
|
||||
//! `min_savings` 判定や savings 推定もこの crate には置かず、上位層が
|
||||
//! usage 履歴ベースのトークン会計と組み合わせて行う。
|
||||
|
||||
use std::ops::Range;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::llm_client::types::Item;
|
||||
|
||||
/// Callback that estimates the token savings for dropping `history[range]`.
|
||||
///
|
||||
/// Injected into [`crate::Worker`] via `set_savings_estimator` so the
|
||||
/// Worker can make `min_savings` decisions without knowing about usage
|
||||
/// measurement sources. Return `0` to signal "no data / refuse to prune".
|
||||
pub type SavingsEstimator = Box<dyn Fn(&[Item], Range<usize>) -> u64 + Send + Sync>;
|
||||
|
||||
/// Configuration for the Prune algorithm.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PruneConfig {
|
||||
@@ -64,6 +73,24 @@ fn find_turn_starts(items: &[Item]) -> Vec<usize> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Set `content = None` on each `Item::ToolResult` at the given indices.
|
||||
///
|
||||
/// Returns the number of items that were actually modified — items that
|
||||
/// are already content-less are counted as 0. Intended for use on a
|
||||
/// request-context clone (never on a persistent history).
|
||||
pub fn project(items: &mut [Item], indices: &[usize]) -> usize {
|
||||
let mut count = 0;
|
||||
for &i in indices {
|
||||
if let Item::ToolResult { content, .. } = &mut items[i] {
|
||||
if content.is_some() {
|
||||
*content = None;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
count
|
||||
}
|
||||
|
||||
/// Indices of `Item::ToolResult { content: Some(_), .. }` that lie outside
|
||||
/// the last `protected_turns` turns. Pure: does not mutate `items`.
|
||||
///
|
||||
@@ -150,6 +177,62 @@ mod tests {
|
||||
assert!(prunable_indices(&items, 2).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_drops_content_and_counts_modifications() {
|
||||
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 count = project(&mut items, &candidates);
|
||||
assert_eq!(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 projected out");
|
||||
} else {
|
||||
assert!(content.is_some(), "protected content should remain");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_skips_already_pruned_items() {
|
||||
// indices points at an item whose content is already None.
|
||||
// project() should count it as 0 modifications.
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", None)]),
|
||||
("turn2", vec![("s2", Some("hello"))]),
|
||||
]);
|
||||
// Manually target s1 (index 3) even though it's already None.
|
||||
let target = items
|
||||
.iter()
|
||||
.position(|it| matches!(it, Item::ToolResult { summary, .. } if summary == "s1"))
|
||||
.unwrap();
|
||||
let count = project(&mut items, &[target]);
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_is_idempotent() {
|
||||
let big = "x".repeat(64);
|
||||
let mut items = make_history(&[
|
||||
("turn1", vec![("s1", Some(&big))]),
|
||||
("turn2", vec![]),
|
||||
("turn3", vec![]),
|
||||
("turn4", vec![]),
|
||||
]);
|
||||
let candidates = prunable_indices(&items, 2);
|
||||
assert_eq!(project(&mut items, &candidates), 1);
|
||||
// 2 周目: 候補は一度の prunable_indices 結果を使い回しても 0 件。
|
||||
assert_eq!(project(&mut items, &candidates), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protected_turns_boundary_exact() {
|
||||
// 3 turns with protected_turns=2: only turn 1 is a candidate.
|
||||
|
||||
@@ -162,6 +162,12 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
|
||||
/// Cancel notification channel (for interrupting execution)
|
||||
cancel_tx: mpsc::Sender<()>,
|
||||
cancel_rx: mpsc::Receiver<()>,
|
||||
/// Prune configuration. `None` disables the prune projection.
|
||||
prune_config: Option<crate::prune::PruneConfig>,
|
||||
/// Callback that estimates token savings for a drop range, injected
|
||||
/// by higher layers that own usage measurements. `None` disables
|
||||
/// the prune projection.
|
||||
savings_estimator: Option<crate::prune::SavingsEstimator>,
|
||||
/// State marker
|
||||
_state: PhantomData<S>,
|
||||
}
|
||||
@@ -303,6 +309,28 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
|
||||
self.interceptor = Box::new(interceptor);
|
||||
}
|
||||
|
||||
/// Configure the prune projection applied to each outgoing request
|
||||
/// context.
|
||||
///
|
||||
/// Both this and [`set_savings_estimator`](Self::set_savings_estimator)
|
||||
/// must be set for the projection to fire; missing either one is a
|
||||
/// no-op. See the crate-level [`prune`](crate::prune) docs for the
|
||||
/// semantics.
|
||||
pub fn set_prune_config(&mut self, config: Option<crate::prune::PruneConfig>) {
|
||||
self.prune_config = config;
|
||||
}
|
||||
|
||||
/// Inject the callback used to estimate token savings for a prune
|
||||
/// candidate range.
|
||||
///
|
||||
/// The callback is invoked with the *request context* (a clone of
|
||||
/// history) and the candidate index range. It must be pure/idempotent
|
||||
/// since it may be called once per LLM request. Return `0` to signal
|
||||
/// "no data" or "refuse to prune".
|
||||
pub fn set_savings_estimator(&mut self, estimator: Option<crate::prune::SavingsEstimator>) {
|
||||
self.savings_estimator = estimator;
|
||||
}
|
||||
|
||||
/// Get a mutable reference to the timeline (for additional handler registration)
|
||||
pub fn timeline_mut(&mut self) -> &mut Timeline {
|
||||
&mut self.timeline
|
||||
@@ -697,8 +725,40 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
|
||||
cb(current_turn);
|
||||
}
|
||||
|
||||
// Interceptor: pre_llm_request
|
||||
// Clone the history into a per-request context. Everything
|
||||
// below (prune projection, interceptor hooks) mutates only
|
||||
// this clone, so the persistent `self.history` stays intact.
|
||||
let mut request_context = self.history.clone();
|
||||
|
||||
// Prune projection: if both the config and the savings
|
||||
// estimator are configured, drop ToolResult.content from
|
||||
// prunable candidates whose estimated savings meet the
|
||||
// threshold. Worker does not own usage history itself; the
|
||||
// estimator is injected by the layer that does.
|
||||
if let (Some(config), Some(estimator)) =
|
||||
(&self.prune_config, &self.savings_estimator)
|
||||
{
|
||||
let candidates =
|
||||
crate::prune::prunable_indices(&request_context, config.protected_turns);
|
||||
if !candidates.is_empty() {
|
||||
let first = *candidates.first().unwrap();
|
||||
let last = *candidates.last().unwrap() + 1;
|
||||
let savings = estimator(&request_context, first..last);
|
||||
if savings >= config.min_savings {
|
||||
let pruned =
|
||||
crate::prune::project(&mut request_context, &candidates);
|
||||
if pruned > 0 {
|
||||
debug!(
|
||||
pruned,
|
||||
estimated_savings_tokens = savings,
|
||||
"Projected old tool-result content out of request context"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Interceptor: pre_llm_request
|
||||
match self.interceptor.pre_llm_request(&mut request_context).await {
|
||||
PreRequestAction::Cancel(reason) => {
|
||||
info!(reason = %reason, "Aborted by interceptor");
|
||||
@@ -899,6 +959,8 @@ impl<C: LlmClient> Worker<C, Mutable> {
|
||||
last_run_interrupted: false,
|
||||
cancel_tx,
|
||||
cancel_rx,
|
||||
prune_config: None,
|
||||
savings_estimator: None,
|
||||
_state: PhantomData,
|
||||
}
|
||||
}
|
||||
@@ -1147,6 +1209,8 @@ impl<C: LlmClient> Worker<C, Mutable> {
|
||||
|
||||
cancel_tx: self.cancel_tx,
|
||||
cancel_rx: self.cancel_rx,
|
||||
prune_config: self.prune_config,
|
||||
savings_estimator: self.savings_estimator,
|
||||
_state: PhantomData,
|
||||
}
|
||||
}
|
||||
@@ -1217,6 +1281,8 @@ impl<C: LlmClient> Worker<C, Locked> {
|
||||
|
||||
cancel_tx: self.cancel_tx,
|
||||
cancel_rx: self.cancel_rx,
|
||||
prune_config: self.prune_config,
|
||||
savings_estimator: self.savings_estimator,
|
||||
_state: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user