do_memory_storage_turso/cache/
wrapper_backend.rs1use super::wrapper::CachedTursoStorage;
6use async_trait::async_trait;
7use do_memory_core::memory::attribution::{
8 RecommendationFeedback, RecommendationSession, RecommendationStats,
9};
10use do_memory_core::{
11 Episode, Error, Heuristic, Pattern, Result, StorageBackend, episode::PatternId,
12};
13use uuid::Uuid;
14
15#[async_trait]
16impl StorageBackend for CachedTursoStorage {
17 async fn store_episode(&self, episode: &Episode) -> Result<()> {
18 self.store_episode_cached(episode)
19 .await
20 .map_err(|e| Error::Storage(format!("Cache store error: {}", e)))
21 }
22
23 async fn get_episode(&self, id: Uuid) -> Result<Option<Episode>> {
24 self.get_episode_cached(id)
25 .await
26 .map_err(|e| Error::Storage(format!("Cache get error: {}", e)))
27 }
28
29 async fn delete_episode(&self, id: Uuid) -> Result<()> {
30 self.delete_episode_cached(id)
31 .await
32 .map_err(|e| Error::Storage(format!("Cache delete error: {}", e)))
33 }
34
35 async fn store_pattern(&self, pattern: &Pattern) -> Result<()> {
36 self.store_pattern_cached(pattern)
37 .await
38 .map_err(|e| Error::Storage(format!("Cache store error: {}", e)))
39 }
40
41 async fn get_pattern(&self, id: PatternId) -> Result<Option<Pattern>> {
42 self.get_pattern_cached(id)
43 .await
44 .map_err(|e| Error::Storage(format!("Cache get error: {}", e)))
45 }
46
47 async fn store_heuristic(&self, heuristic: &Heuristic) -> Result<()> {
48 self.store_heuristic_cached(heuristic)
49 .await
50 .map_err(|e| Error::Storage(format!("Cache store error: {}", e)))
51 }
52
53 async fn get_heuristic(&self, id: Uuid) -> Result<Option<Heuristic>> {
54 self.get_heuristic_cached(id)
55 .await
56 .map_err(|e| Error::Storage(format!("Cache get error: {}", e)))
57 }
58
59 async fn query_episodes_since(
60 &self,
61 since: chrono::DateTime<chrono::Utc>,
62 limit: Option<usize>,
63 ) -> Result<Vec<Episode>> {
64 self.storage
65 .query_episodes_since(since, limit)
66 .await
67 .map_err(|e| Error::Storage(format!("Query error: {}", e)))
68 }
69
70 async fn query_episodes_by_metadata(
71 &self,
72 key: &str,
73 value: &str,
74 limit: Option<usize>,
75 ) -> Result<Vec<Episode>> {
76 self.storage
77 .query_episodes_by_metadata(key, value, limit)
78 .await
79 .map_err(|e| Error::Storage(format!("Query error: {}", e)))
80 }
81
82 async fn store_embedding(&self, id: &str, embedding: Vec<f32>) -> Result<()> {
83 self.storage
84 .store_embedding(id, embedding)
85 .await
86 .map_err(|e| Error::Storage(format!("Store embedding error: {}", e)))
87 }
88
89 async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>> {
90 self.storage
91 .get_embedding(id)
92 .await
93 .map_err(|e| Error::Storage(format!("Get embedding error: {}", e)))
94 }
95
96 async fn delete_embedding(&self, id: &str) -> Result<bool> {
97 self.storage
98 .delete_embedding(id)
99 .await
100 .map_err(|e| Error::Storage(format!("Delete embedding error: {}", e)))
101 }
102
103 async fn store_embeddings_batch(&self, embeddings: Vec<(String, Vec<f32>)>) -> Result<()> {
104 self.storage
105 .store_embeddings_batch(embeddings)
106 .await
107 .map_err(|e| Error::Storage(format!("Batch store embeddings error: {}", e)))
108 }
109
110 async fn get_embeddings_batch(&self, ids: &[String]) -> Result<Vec<Option<Vec<f32>>>> {
111 self.storage
112 .get_embeddings_batch(ids)
113 .await
114 .map_err(|e| Error::Storage(format!("Batch get embeddings error: {}", e)))
115 }
116
117 async fn store_recommendation_session(&self, session: &RecommendationSession) -> Result<()> {
118 self.storage
119 .store_recommendation_session(session)
120 .await
121 .map_err(|e| Error::Storage(format!("Store recommendation session error: {}", e)))
122 }
123
124 async fn get_recommendation_session(
125 &self,
126 session_id: Uuid,
127 ) -> Result<Option<RecommendationSession>> {
128 self.storage
129 .get_recommendation_session(session_id)
130 .await
131 .map_err(|e| Error::Storage(format!("Get recommendation session error: {}", e)))
132 }
133
134 async fn get_recommendation_session_for_episode(
135 &self,
136 episode_id: Uuid,
137 ) -> Result<Option<RecommendationSession>> {
138 self.storage
139 .get_recommendation_session_for_episode(episode_id)
140 .await
141 .map_err(|e| {
142 Error::Storage(format!("Get recommendation session (episode) error: {}", e))
143 })
144 }
145
146 async fn store_recommendation_feedback(&self, feedback: &RecommendationFeedback) -> Result<()> {
147 self.storage
148 .store_recommendation_feedback(feedback)
149 .await
150 .map_err(|e| Error::Storage(format!("Store recommendation feedback error: {}", e)))
151 }
152
153 async fn get_recommendation_feedback(
154 &self,
155 session_id: Uuid,
156 ) -> Result<Option<RecommendationFeedback>> {
157 self.storage
158 .get_recommendation_feedback(session_id)
159 .await
160 .map_err(|e| Error::Storage(format!("Get recommendation feedback error: {}", e)))
161 }
162
163 async fn get_recommendation_stats(&self) -> Result<RecommendationStats> {
164 self.storage
165 .get_recommendation_stats()
166 .await
167 .map_err(|e| Error::Storage(format!("Get recommendation stats error: {}", e)))
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174 use crate::TursoStorage;
175 use crate::cache::config::CacheConfig;
176 use crate::cache::wrapper::CachedTursoStorage;
177 use do_memory_core::{Evidence, TaskContext, TaskOutcome, TaskType};
178 use libsql::Builder;
179 use tempfile::TempDir;
180
181 async fn setup() -> (CachedTursoStorage, TempDir) {
182 let dir = TempDir::new().unwrap();
183 let db_path = dir.path().join("test_backend.db");
184 let db = Builder::new_local(&db_path).build().await.unwrap();
185 let storage = TursoStorage::from_database(db).unwrap();
186 storage.initialize_schema().await.unwrap();
187 let cached = CachedTursoStorage::new(storage, CacheConfig::default());
188 (cached, dir)
189 }
190
191 fn test_episode() -> Episode {
192 Episode {
193 episode_id: Uuid::new_v4(),
194 task_type: TaskType::CodeGeneration,
195 task_description: "backend test".to_string(),
196 context: TaskContext {
197 domain: "test".to_string(),
198 ..Default::default()
199 },
200 steps: vec![],
201 outcome: None,
202 reward: None,
203 reflection: None,
204 patterns: vec![],
205 heuristics: vec![],
206 applied_patterns: vec![],
207 salient_features: None,
208 start_time: chrono::Utc::now(),
209 end_time: None,
210 metadata: std::collections::HashMap::new(),
211 tags: vec![],
212 checkpoints: vec![],
213 }
214 }
215
216 #[tokio::test]
217 async fn test_storage_backend_episode_lifecycle() {
218 let (cached, _dir) = setup().await;
219 let ep = test_episode();
220
221 StorageBackend::store_episode(&cached, &ep).await.unwrap();
222 let got = StorageBackend::get_episode(&cached, ep.episode_id)
223 .await
224 .unwrap();
225 assert_eq!(got.unwrap().episode_id, ep.episode_id);
226
227 StorageBackend::delete_episode(&cached, ep.episode_id)
228 .await
229 .unwrap();
230 let gone = StorageBackend::get_episode(&cached, ep.episode_id)
231 .await
232 .unwrap();
233 assert!(gone.is_none());
234 }
235
236 #[tokio::test]
237 async fn test_storage_backend_pattern_and_heuristic() {
238 let (cached, _dir) = setup().await;
239
240 let pattern = Pattern::ToolSequence {
241 id: Uuid::new_v4(),
242 tools: vec!["t1".into()],
243 context: TaskContext {
244 domain: "test".into(),
245 ..Default::default()
246 },
247 success_rate: 0.5,
248 avg_latency: chrono::Duration::milliseconds(10),
249 occurrence_count: 1,
250 effectiveness: Default::default(),
251 };
252 StorageBackend::store_pattern(&cached, &pattern)
253 .await
254 .unwrap();
255 let got = StorageBackend::get_pattern(&cached, pattern.id())
256 .await
257 .unwrap();
258 assert!(got.is_some());
259
260 let h = Heuristic {
261 heuristic_id: Uuid::new_v4(),
262 condition: "c".to_string(),
263 action: "a".to_string(),
264 confidence: 0.9,
265 evidence: Evidence {
266 episode_ids: vec![],
267 success_rate: 0.9,
268 sample_size: 1,
269 },
270 created_at: chrono::Utc::now(),
271 updated_at: chrono::Utc::now(),
272 };
273 StorageBackend::store_heuristic(&cached, &h).await.unwrap();
274 let got_h = StorageBackend::get_heuristic(&cached, h.heuristic_id)
275 .await
276 .unwrap();
277 assert!(got_h.is_some());
278 }
279
280 #[tokio::test]
281 async fn test_storage_backend_queries() {
282 let (cached, _dir) = setup().await;
283 let ep = test_episode();
284 StorageBackend::store_episode(&cached, &ep).await.unwrap();
285
286 let since = chrono::Utc::now() - chrono::Duration::hours(1);
287 let results = StorageBackend::query_episodes_since(&cached, since, Some(10))
288 .await
289 .unwrap();
290 assert!(!results.is_empty());
291
292 let _ =
293 StorageBackend::query_episodes_by_metadata(&cached, "domain", "test", Some(10)).await;
294 }
295
296 #[tokio::test]
297 async fn test_storage_backend_embeddings() {
298 let (cached, _dir) = setup().await;
299
300 let embedding: Vec<f32> = (0..384).map(|i| i as f32 / 384.0).collect();
301 StorageBackend::store_embedding(&cached, "e1", embedding)
302 .await
303 .unwrap();
304 let got = StorageBackend::get_embedding(&cached, "e1").await.unwrap();
305 assert!(got.is_some());
306
307 let deleted = StorageBackend::delete_embedding(&cached, "e1")
308 .await
309 .unwrap();
310 assert!(deleted);
311
312 let emb1: Vec<f32> = (0..384).map(|i| i as f32 / 384.0).collect();
313 let emb2: Vec<f32> = (0..384).map(|i| (384 - i) as f32 / 384.0).collect();
314 StorageBackend::store_embeddings_batch(
315 &cached,
316 vec![("b1".into(), emb1), ("b2".into(), emb2)],
317 )
318 .await
319 .unwrap();
320
321 let batch = StorageBackend::get_embeddings_batch(&cached, &["b1".into(), "b2".into()])
322 .await
323 .unwrap();
324 assert_eq!(batch.len(), 2);
325 }
326
327 #[tokio::test]
328 async fn test_storage_backend_recommendations() {
329 let (cached, _dir) = setup().await;
330
331 let session = RecommendationSession {
332 session_id: Uuid::new_v4(),
333 episode_id: Uuid::new_v4(),
334 timestamp: chrono::Utc::now(),
335 recommended_pattern_ids: vec!["p1".into()],
336 recommended_playbook_ids: vec![],
337 };
338
339 StorageBackend::store_recommendation_session(&cached, &session)
340 .await
341 .unwrap();
342 let got = StorageBackend::get_recommendation_session(&cached, session.session_id)
343 .await
344 .unwrap();
345 assert!(got.is_some());
346
347 let by_ep =
348 StorageBackend::get_recommendation_session_for_episode(&cached, session.episode_id)
349 .await
350 .unwrap();
351 assert!(by_ep.is_some());
352
353 let feedback = RecommendationFeedback {
354 session_id: session.session_id,
355 applied_pattern_ids: vec!["p1".into()],
356 consulted_episode_ids: vec![],
357 outcome: TaskOutcome::Success {
358 verdict: "done".into(),
359 artifacts: vec![],
360 },
361 agent_rating: Some(0.9),
362 };
363 StorageBackend::store_recommendation_feedback(&cached, &feedback)
364 .await
365 .unwrap();
366 let got_fb = StorageBackend::get_recommendation_feedback(&cached, session.session_id)
367 .await
368 .unwrap();
369 assert!(got_fb.is_some());
370
371 let stats = StorageBackend::get_recommendation_stats(&cached)
372 .await
373 .unwrap();
374 assert!(stats.total_sessions >= 1);
375 }
376}