Skip to main content

ronn_memory/
episodic.rs

1//! Episodic Memory - Experience storage with temporal indexing
2
3use crate::{MemoryId, Result};
4use ronn_core::tensor::Tensor;
5use std::collections::HashMap;
6
7/// An episode (experience) in memory
8#[derive(Clone)]
9pub struct Episode {
10    pub id: MemoryId,
11    pub data: Tensor,
12    pub timestamp: u64,
13    pub importance: f64,
14}
15
16/// Query parameters for episodic memory
17#[derive(Debug, Clone)]
18pub struct EpisodeQuery {
19    pub start_time: Option<u64>,
20    pub end_time: Option<u64>,
21    pub min_importance: Option<f64>,
22    pub limit: usize,
23}
24
25impl Default for EpisodeQuery {
26    fn default() -> Self {
27        Self {
28            start_time: None,
29            end_time: None,
30            min_importance: None,
31            limit: 10,
32        }
33    }
34}
35
36/// Episodic memory storage
37pub struct EpisodicMemory {
38    episodes: HashMap<MemoryId, Episode>,
39    temporal_index: Vec<(u64, MemoryId)>, // (timestamp, id) for temporal queries
40}
41
42impl EpisodicMemory {
43    /// Create new episodic memory
44    pub fn new() -> Self {
45        Self {
46            episodes: HashMap::new(),
47            temporal_index: Vec::new(),
48        }
49    }
50
51    /// Store an episode
52    pub fn store_episode(&mut self, episode: Episode) -> Result<()> {
53        let id = episode.id;
54        let timestamp = episode.timestamp;
55
56        self.episodes.insert(id, episode);
57
58        // Add to temporal index
59        self.temporal_index.push((timestamp, id));
60
61        // Keep temporal index sorted
62        self.temporal_index.sort_by_key(|(ts, _)| *ts);
63
64        Ok(())
65    }
66
67    /// Get episode by ID
68    pub fn get_episode(&self, id: MemoryId) -> Option<Episode> {
69        self.episodes.get(&id).cloned()
70    }
71
72    /// Query episodes by criteria
73    pub fn query(&self, query: &EpisodeQuery) -> Vec<Episode> {
74        self.episodes
75            .values()
76            .filter(|ep| {
77                // Filter by time range
78                if let Some(start) = query.start_time {
79                    if ep.timestamp < start {
80                        return false;
81                    }
82                }
83                if let Some(end) = query.end_time {
84                    if ep.timestamp > end {
85                        return false;
86                    }
87                }
88
89                // Filter by importance
90                if let Some(min_imp) = query.min_importance {
91                    if ep.importance < min_imp {
92                        return false;
93                    }
94                }
95
96                true
97            })
98            .take(query.limit)
99            .cloned()
100            .collect()
101    }
102
103    /// Get number of episodes
104    pub fn len(&self) -> usize {
105        self.episodes.len()
106    }
107
108    /// Check if empty
109    pub fn is_empty(&self) -> bool {
110        self.episodes.is_empty()
111    }
112
113    /// Get all episodes (for consolidation)
114    pub fn all_episodes(&self) -> Vec<&Episode> {
115        self.episodes.values().collect()
116    }
117}
118
119impl Default for EpisodicMemory {
120    fn default() -> Self {
121        Self::new()
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
129    use super::*;
130    use crate::current_timestamp;
131    use ronn_core::types::{DataType, TensorLayout};
132
133    #[test]
134    fn test_store_and_retrieve() -> Result<()> {
135        let mut em = EpisodicMemory::new();
136
137        let data = vec![1.0f32, 2.0, 3.0];
138        let tensor = Tensor::from_data(data, vec![1, 3], DataType::F32, TensorLayout::RowMajor)?;
139
140        let episode = Episode {
141            id: 1,
142            data: tensor,
143            timestamp: current_timestamp(),
144            importance: 0.8,
145        };
146
147        em.store_episode(episode.clone())?;
148
149        let retrieved = em.get_episode(1);
150        assert!(retrieved.is_some());
151        assert_eq!(retrieved.unwrap().id, 1);
152
153        Ok(())
154    }
155
156    #[test]
157    fn test_temporal_query() -> Result<()> {
158        let mut em = EpisodicMemory::new();
159        let base_time = current_timestamp();
160
161        // Store episodes at different times
162        for i in 0..5 {
163            let data = vec![i as f32; 2];
164            let tensor =
165                Tensor::from_data(data, vec![1, 2], DataType::F32, TensorLayout::RowMajor)?;
166
167            let episode = Episode {
168                id: i as MemoryId,
169                data: tensor,
170                timestamp: base_time + (i * 1000),
171                importance: 0.5 + (i as f64 * 0.1),
172            };
173
174            em.store_episode(episode)?;
175        }
176
177        // Query for episodes
178        let query = EpisodeQuery {
179            min_importance: Some(0.7),
180            ..Default::default()
181        };
182
183        let results = em.query(&query);
184        assert!(results.len() >= 2);
185
186        Ok(())
187    }
188}