use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectedItem {
pub source_id: String,
pub source_type: SourceType,
pub content: String,
pub timestamp: DateTime<Utc>,
pub importance: u8,
pub tags: Vec<String>,
pub metadata: HashMap<String, String>,
pub related_ids: Vec<String>,
pub session_id: Option<String>,
pub file_refs: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SourceType {
Episode,
MemoryEntry,
SessionEvent,
InteractionTrace,
ToolResult,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectedBatch {
pub items: Vec<CollectedItem>,
pub time_range: (DateTime<Utc>, DateTime<Utc>),
pub source_counts: HashMap<SourceType, usize>,
}
pub struct ShortTermCollector {
max_age_secs: u64,
min_importance: u8,
}
impl ShortTermCollector {
pub fn new(max_age_secs: u64, min_importance: u8) -> Self {
Self {
max_age_secs,
min_importance,
}
}
pub fn collect_episodes(&self, episodes: &[EpisodeData]) -> Vec<CollectedItem> {
let cutoff = Utc::now() - chrono::Duration::seconds(self.max_age_secs as i64);
episodes
.iter()
.filter(|ep| ep.timestamp >= cutoff && ep.importance >= self.min_importance)
.map(|ep| CollectedItem {
source_id: ep.id.clone(),
source_type: SourceType::Episode,
content: ep.content.clone(),
timestamp: ep.timestamp,
importance: ep.importance,
tags: ep.tags.clone(),
metadata: ep.context.clone(),
related_ids: ep.related_ids.clone(),
session_id: Some(ep.session_id.clone()),
file_refs: Vec::new(),
})
.collect()
}
pub fn assemble_batch(&self, items: Vec<CollectedItem>) -> CollectedBatch {
let mut source_counts: HashMap<SourceType, usize> = HashMap::new();
for item in &items {
*source_counts.entry(item.source_type).or_insert(0) += 1;
}
let time_range = if items.is_empty() {
let now = Utc::now();
(now, now)
} else {
let min = items.iter().map(|i| i.timestamp).min().unwrap();
let max = items.iter().map(|i| i.timestamp).max().unwrap();
(min, max)
};
CollectedBatch {
items,
time_range,
source_counts,
}
}
}
#[derive(Debug, Clone)]
pub struct EpisodeData {
pub id: String,
pub content: String,
pub timestamp: DateTime<Utc>,
pub importance: u8,
pub tags: Vec<String>,
pub context: HashMap<String, String>,
pub related_ids: Vec<String>,
pub session_id: String,
}
#[derive(Debug, Clone)]
pub struct MemoryEntryData {
pub content: String,
pub role: String,
pub timestamp: DateTime<Utc>,
}
#[cfg(test)]
#[path = "../../tests/unit/consolidation/collector/collector_test.rs"]
mod tests;