use crate::context::scoring::{calculate_priority_score, compare_by_priority};
use crate::context::types::{AddResult, BundleConfig, BundleStats, ContextItem};
use chrono::Utc;
use tracing::{debug, trace};
use uuid::Uuid;
#[derive(Debug)]
pub struct BundleAccumulator {
config: BundleConfig,
items: Vec<ContextItem>,
stats: BundleStats,
reference_time: chrono::DateTime<Utc>,
}
impl BundleAccumulator {
#[must_use]
pub fn new(config: BundleConfig) -> Self {
Self {
config,
items: Vec::new(),
stats: BundleStats::default(),
reference_time: Utc::now(),
}
}
#[must_use]
pub fn default_config() -> Self {
Self::new(BundleConfig::default())
}
#[must_use]
pub fn token_efficient() -> Self {
Self::new(BundleConfig::token_efficient())
}
#[must_use]
pub fn comprehensive() -> Self {
Self::new(BundleConfig::comprehensive())
}
pub fn set_reference_time(&mut self, time: chrono::DateTime<Utc>) {
self.reference_time = time;
}
#[must_use]
pub fn config(&self) -> &BundleConfig {
&self.config
}
#[must_use]
pub fn stats(&self) -> &BundleStats {
&self.stats
}
#[must_use]
pub fn size(&self) -> usize {
self.items.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
#[must_use]
pub fn is_full(&self) -> bool {
self.items.len() >= self.config.max_items
}
pub fn add(&mut self, mut item: ContextItem) -> AddResult {
if item.salience() < self.config.min_salience_threshold {
self.stats.total_rejected += 1;
debug!(
item_id = %item.id(),
salience = item.salience(),
threshold = self.config.min_salience_threshold,
"Rejected item below salience threshold"
);
return AddResult::rejected(
self.items.len(),
format!(
"Salience {} below threshold {}",
item.salience(),
self.config.min_salience_threshold
),
);
}
let priority = calculate_priority_score(&item, &self.config, self.reference_time);
item.set_priority(priority);
let evicted_id = if self.items.len() >= self.config.max_items {
self.evict_lowest_priority()
} else {
None
};
let added_id = item.id();
self.items.push(item);
self.stats.total_added += 1;
if let Some(id) = evicted_id {
self.stats.total_evicted += 1;
debug!(
added_id = %added_id,
evicted_id = %id,
bundle_size = self.items.len(),
"Added item, evicted lowest priority"
);
AddResult::accepted_with_eviction(self.items.len(), id)
} else {
trace!(
added_id = %added_id,
bundle_size = self.items.len(),
"Added item to bundle"
);
AddResult::accepted(self.items.len())
}
}
pub fn add_batch(&mut self, items: impl IntoIterator<Item = ContextItem>) -> Vec<AddResult> {
items.into_iter().map(|item| self.add(item)).collect()
}
fn evict_lowest_priority(&mut self) -> Option<Uuid> {
if self.items.is_empty() {
return None;
}
let min_idx = self
.items
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| {
a.priority()
.partial_cmp(&b.priority())
.unwrap_or(std::cmp::Ordering::Equal)
})
.map(|(idx, _)| idx);
if let Some(idx) = min_idx {
let evicted = self.items.remove(idx);
Some(evicted.id())
} else {
None
}
}
pub fn remove(&mut self, id: Uuid) -> bool {
let idx = self.items.iter().position(|item| item.id() == id);
if let Some(idx) = idx {
self.items.remove(idx);
true
} else {
false
}
}
#[must_use]
pub fn contains(&self, id: Uuid) -> bool {
self.items.iter().any(|item| item.id() == id)
}
#[must_use]
pub fn get(&self, id: Uuid) -> Option<&ContextItem> {
self.items.iter().find(|item| item.id() == id)
}
#[must_use]
pub fn to_bundle(&mut self) -> Vec<ContextItem> {
self.recompute_priorities();
self.items.sort_by(compare_by_priority);
self.items.truncate(self.config.max_items);
self.update_stats();
self.items.clone()
}
#[must_use]
pub fn peek_items(&self) -> &[ContextItem] {
&self.items
}
#[must_use]
pub fn episodes_only(&self) -> Vec<&ContextItem> {
self.items
.iter()
.filter(|item| item.item_type() == crate::context::types::ContextItemType::Episode)
.collect()
}
#[must_use]
pub fn patterns_only(&self) -> Vec<&ContextItem> {
self.items
.iter()
.filter(|item| item.item_type() == crate::context::types::ContextItemType::Pattern)
.collect()
}
pub fn clear(&mut self) {
self.items.clear();
self.stats = BundleStats::default();
}
fn recompute_priorities(&mut self) {
for item in &mut self.items {
let priority = calculate_priority_score(item, &self.config, self.reference_time);
item.set_priority(priority);
}
}
fn update_stats(&mut self) {
if self.items.is_empty() {
self.stats.average_salience = 0.0;
self.stats.average_priority = 0.0;
self.stats.oldest_timestamp = None;
self.stats.newest_timestamp = None;
return;
}
self.stats.current_size = self.items.len();
let total_salience: f32 = self.items.iter().map(|i| i.salience()).sum();
let total_priority: f32 = self.items.iter().map(|i| i.priority()).sum();
self.stats.average_salience = total_salience / self.items.len() as f32;
self.stats.average_priority = total_priority / self.items.len() as f32;
self.stats.oldest_timestamp = self.items.iter().map(|i| i.timestamp()).min();
self.stats.newest_timestamp = self.items.iter().map(|i| i.timestamp()).max();
}
pub fn from_episodes(
episodes: Vec<std::sync::Arc<crate::episode::Episode>>,
salience_fn: impl Fn(&crate::episode::Episode) -> f32,
) -> Vec<ContextItem> {
let mut accumulator = Self::default_config();
for episode in episodes {
let salience = salience_fn(&episode);
let item = ContextItem::from_episode(episode, salience);
accumulator.add(item);
}
accumulator.to_bundle()
}
pub fn from_episodes_with_config(
episodes: Vec<std::sync::Arc<crate::episode::Episode>>,
config: BundleConfig,
salience_fn: impl Fn(&crate::episode::Episode) -> f32,
) -> Vec<ContextItem> {
let mut accumulator = Self::new(config);
for episode in episodes {
let salience = salience_fn(&episode);
let item = ContextItem::from_episode(episode, salience);
accumulator.add(item);
}
accumulator.to_bundle()
}
}
impl Default for BundleAccumulator {
fn default() -> Self {
Self::default_config()
}
}