ronn_memory/
working.rs

1//! Working Memory - Short-term storage with attention weighting
2
3use crate::{MemoryError, MemoryId, Result, current_timestamp};
4use ronn_core::tensor::Tensor;
5use std::collections::{HashMap, VecDeque};
6
7/// Configuration for working memory
8#[derive(Debug, Clone)]
9pub struct WorkingMemoryConfig {
10    /// Maximum number of items in working memory
11    pub capacity: usize,
12
13    /// Time-to-live for items in milliseconds
14    pub ttl_ms: u64,
15
16    /// Whether to use attention weighting
17    pub attention_enabled: bool,
18}
19
20impl Default for WorkingMemoryConfig {
21    fn default() -> Self {
22        Self {
23            capacity: 100,
24            ttl_ms: 60_000, // 1 minute
25            attention_enabled: true,
26        }
27    }
28}
29
30/// An item in working memory
31#[derive(Clone)]
32pub struct WorkingMemoryItem {
33    pub id: MemoryId,
34    pub data: Tensor,
35    pub importance: f64,
36    pub timestamp: u64,
37    pub access_count: u64,
38}
39
40/// Working memory with attention-weighted storage
41pub struct WorkingMemory {
42    config: WorkingMemoryConfig,
43    items: HashMap<MemoryId, WorkingMemoryItem>,
44    lru_order: VecDeque<MemoryId>,
45    next_id: MemoryId,
46}
47
48impl WorkingMemory {
49    /// Create new working memory with configuration
50    pub fn new(config: WorkingMemoryConfig) -> Self {
51        Self {
52            config,
53            items: HashMap::new(),
54            lru_order: VecDeque::new(),
55            next_id: 1,
56        }
57    }
58
59    /// Store data with importance score
60    pub fn store(&mut self, data: Tensor, importance: f64) -> Result<MemoryId> {
61        // Evict old items if at capacity
62        while self.items.len() >= self.config.capacity {
63            self.evict_lru()?;
64        }
65
66        let id = self.next_id;
67        self.next_id += 1;
68
69        let item = WorkingMemoryItem {
70            id,
71            data,
72            importance,
73            timestamp: current_timestamp(),
74            access_count: 0,
75        };
76
77        self.items.insert(id, item);
78        self.lru_order.push_back(id);
79
80        Ok(id)
81    }
82
83    /// Retrieve data by ID
84    pub fn get(&self, id: MemoryId) -> Result<Tensor> {
85        self.items
86            .get(&id)
87            .map(|item| item.data.clone())
88            .ok_or_else(|| MemoryError::WorkingMemory(format!("Item {} not found", id)))
89    }
90
91    /// Search for similar items
92    pub fn search_similar(&self, _query: &Tensor, limit: usize) -> Result<Vec<MemoryId>> {
93        // For MVP: return most recent items
94        Ok(self.lru_order.iter().rev().take(limit).copied().collect())
95    }
96
97    /// Drain old items for consolidation
98    pub fn drain_old_items(&mut self) -> Result<Vec<(MemoryId, Tensor, f64)>> {
99        let current_time = current_timestamp();
100        let mut drained = Vec::new();
101
102        let expired_ids: Vec<MemoryId> = self
103            .items
104            .iter()
105            .filter(|(_, item)| current_time - item.timestamp > self.config.ttl_ms)
106            .map(|(id, _)| *id)
107            .collect();
108
109        for id in expired_ids {
110            if let Some(item) = self.items.remove(&id) {
111                drained.push((item.id, item.data, item.importance));
112                self.lru_order.retain(|&x| x != id);
113            }
114        }
115
116        Ok(drained)
117    }
118
119    /// Evict least recently used item
120    fn evict_lru(&mut self) -> Result<()> {
121        if let Some(id) = self.lru_order.pop_front() {
122            self.items.remove(&id);
123        }
124        Ok(())
125    }
126
127    /// Get number of items
128    pub fn len(&self) -> usize {
129        self.items.len()
130    }
131
132    /// Check if empty
133    pub fn is_empty(&self) -> bool {
134        self.items.is_empty()
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
142    use super::*;
143    use ronn_core::types::{DataType, TensorLayout};
144
145    #[test]
146    fn test_store_and_retrieve() -> Result<()> {
147        let mut wm = WorkingMemory::new(WorkingMemoryConfig::default());
148
149        let data = vec![1.0f32, 2.0, 3.0];
150        let tensor = Tensor::from_data(data, vec![1, 3], DataType::F32, TensorLayout::RowMajor)?;
151
152        let id = wm.store(tensor.clone(), 0.5)?;
153        let retrieved = wm.get(id)?;
154
155        assert_eq!(retrieved.shape(), tensor.shape());
156        assert_eq!(wm.len(), 1);
157
158        Ok(())
159    }
160
161    #[test]
162    fn test_capacity_eviction() -> Result<()> {
163        let config = WorkingMemoryConfig {
164            capacity: 3,
165            ..Default::default()
166        };
167        let mut wm = WorkingMemory::new(config);
168
169        // Add 4 items (should evict oldest)
170        for i in 0..4 {
171            let data = vec![i as f32; 2];
172            let tensor =
173                Tensor::from_data(data, vec![1, 2], DataType::F32, TensorLayout::RowMajor)?;
174            wm.store(tensor, 0.5)?;
175        }
176
177        assert_eq!(wm.len(), 3);
178
179        Ok(())
180    }
181}
ronn_memory/working.rs

ronn_memory/
working.rs
