1use crate::{MemoryId, Result};
4use ronn_core::tensor::Tensor;
5use std::collections::HashMap;
6
7#[derive(Clone)]
9pub struct Episode {
10 pub id: MemoryId,
11 pub data: Tensor,
12 pub timestamp: u64,
13 pub importance: f64,
14}
15
16#[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
36pub struct EpisodicMemory {
38 episodes: HashMap<MemoryId, Episode>,
39 temporal_index: Vec<(u64, MemoryId)>, }
41
42impl EpisodicMemory {
43 pub fn new() -> Self {
45 Self {
46 episodes: HashMap::new(),
47 temporal_index: Vec::new(),
48 }
49 }
50
51 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 self.temporal_index.push((timestamp, id));
60
61 self.temporal_index.sort_by_key(|(ts, _)| *ts);
63
64 Ok(())
65 }
66
67 pub fn get_episode(&self, id: MemoryId) -> Option<Episode> {
69 self.episodes.get(&id).cloned()
70 }
71
72 pub fn query(&self, query: &EpisodeQuery) -> Vec<Episode> {
74 self.episodes
75 .values()
76 .filter(|ep| {
77 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 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 pub fn len(&self) -> usize {
105 self.episodes.len()
106 }
107
108 pub fn is_empty(&self) -> bool {
110 self.episodes.is_empty()
111 }
112
113 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 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 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}