use super::types::{
ConsolidationResult, MemoryConfig, MemoryEntry, MemoryIndex, MemoryQuery, MemoryTier,
};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
pub struct LongTermMemory {
entries: Arc<RwLock<HashMap<u64, MemoryEntry>>>,
capacity: usize,
index: Arc<MemoryIndex>,
}
impl LongTermMemory {
pub fn new(capacity: usize, index: Arc<MemoryIndex>) -> Self {
Self {
entries: Arc::new(RwLock::new(HashMap::new())),
capacity,
index,
}
}
pub fn with_config(config: &MemoryConfig, index: Arc<MemoryIndex>) -> Self {
Self::new(config.long_term_capacity, index)
}
pub async fn store(&self, mut entry: MemoryEntry) -> anyhow::Result<u64> {
entry.tier = MemoryTier::LongTerm;
let mut entries = self.entries.write().await;
if entries.len() >= self.capacity && !entries.is_empty() {
self.archive_oldest(&mut entries).await?;
}
let id = entry.id;
entries.insert(id, entry.clone());
drop(entries);
self.index.index_entry(&entry).await;
Ok(id)
}
pub async fn retrieve(&self, id: u64) -> Option<MemoryEntry> {
let mut entries = self.entries.write().await;
if let Some(entry) = entries.get_mut(&id) {
entry.accessed();
return Some(entry.clone());
}
None
}
pub async fn query(&self, query: &MemoryQuery) -> Vec<MemoryEntry> {
let entries = self.entries.read().await;
let mut results: Vec<MemoryEntry> = entries
.values()
.filter(|e| super::types::matches_query(e, query))
.cloned()
.collect();
results.sort_by(|a, b| {
b.importance
.partial_cmp(&a.importance)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| b.accessed_at.cmp(&a.accessed_at))
});
if let Some(limit) = query.limit {
results.truncate(limit);
}
results
}
pub async fn update(&self, id: u64, f: impl FnOnce(&mut MemoryEntry)) -> bool {
let mut entries = self.entries.write().await;
if let Some(entry) = entries.get_mut(&id) {
f(entry);
return true;
}
false
}
pub async fn remove(&self, id: u64) -> Option<MemoryEntry> {
let mut entries = self.entries.write().await;
if let Some(entry) = entries.remove(&id) {
self.index.remove_entry(&entry).await;
return Some(entry);
}
None
}
pub async fn count(&self) -> usize {
self.entries.read().await.len()
}
pub async fn entries(&self) -> Vec<MemoryEntry> {
self.entries.read().await.values().cloned().collect()
}
pub async fn consolidate(&self) -> ConsolidationResult {
let mut entries = self.entries.write().await;
let to_consolidate: Vec<u64> = entries
.values()
.filter(|e| e.importance < 0.3)
.map(|e| e.id)
.collect();
let merged = 0;
let mut removed = 0;
for id in &to_consolidate {
if entries.remove(id).is_some() {
removed += 1;
}
}
ConsolidationResult {
entries_merged: merged,
entries_removed: removed,
new_summaries: Vec::new(),
}
}
async fn archive_oldest(
&self,
entries: &mut tokio::sync::RwLockWriteGuard<'_, HashMap<u64, MemoryEntry>>,
) -> anyhow::Result<()> {
let oldest = entries
.values()
.filter(|e| e.importance < 0.5)
.min_by_key(|e| e.accessed_at)
.map(|e| e.id);
if let Some(id) = oldest {
if let Some(entry) = entries.remove(&id) {
self.index.remove_entry(&entry).await;
}
}
Ok(())
}
}
impl Clone for LongTermMemory {
fn clone(&self) -> Self {
Self {
entries: self.entries.clone(),
capacity: self.capacity,
index: self.index.clone(),
}
}
}
pub struct ArchiveMemory {
entries: Arc<RwLock<HashMap<u64, MemoryEntry>>>,
}
impl ArchiveMemory {
pub fn new() -> Self {
Self {
entries: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn store(&self, mut entry: MemoryEntry) -> anyhow::Result<u64> {
entry.tier = MemoryTier::Archive;
let mut entries = self.entries.write().await;
let id = entry.id;
entries.insert(id, entry);
Ok(id)
}
pub async fn retrieve(&self, id: u64) -> Option<MemoryEntry> {
self.entries.read().await.get(&id).cloned()
}
pub async fn count(&self) -> usize {
self.entries.read().await.len()
}
}
impl Default for ArchiveMemory {
fn default() -> Self {
Self::new()
}
}
impl Clone for ArchiveMemory {
fn clone(&self) -> Self {
Self {
entries: self.entries.clone(),
}
}
}