1use std::collections::HashMap;
2
3use anyhow::{Context, Result};
4use turso::params;
5
6use super::store::MemoryStore;
7use super::types::{
8 CategoryCount, Memory, MemoryCategory, MemoryRecord, StoreStatus, TaskCreatedMemory, TaskRecord, TaskRetrieval,
9 TaskStatus, TimelineEvent, TimelineEventKind, TopMemory,
10};
11use super::util::{category_from_str, drain_rows, embedding_status, vec_buf};
12
13impl MemoryStore {
14 pub async fn get_status(&self) -> Result<StoreStatus> {
16 self.init().await?;
17 self.with_db(|conn| async move {
18 let (total_memories, completed_tasks, total_tasks, avg_score) = {
19 let mut rows = conn
20 .query(
21 "SELECT
22 (SELECT COUNT(*) FROM memories),
23 (SELECT COUNT(*) FROM tasks WHERE finished_at IS NOT NULL),
24 (SELECT COUNT(*) FROM tasks),
25 (SELECT AVG(task_score) FROM tasks WHERE task_score IS NOT NULL)",
26 (),
27 )
28 .await?;
29 let row = rows.next().await?.context("no status row")?;
30 let stats = (
31 row.get::<i64>(0)?,
32 row.get::<i64>(1)?,
33 row.get::<i64>(2)?,
34 row.get::<Option<f64>>(3)?.unwrap_or(0.0),
35 );
36 drain_rows(&mut rows).await?;
37 stats
38 };
39
40 let mut cat_rows = conn
41 .query(
42 "SELECT category, COUNT(*) as c FROM memories GROUP BY category ORDER BY c DESC",
43 (),
44 )
45 .await?;
46 let mut categories = Vec::new();
47 while let Some(row) = cat_rows.next().await? {
48 categories.push(CategoryCount {
49 category: category_from_str(&row.get::<String>(0)?),
50 count: row.get::<i64>(1)? as u32,
51 });
52 }
53 drain_rows(&mut cat_rows).await?;
54
55 let mut top_rows = conn
56 .query(
57 "SELECT content, weight, retrieval_count FROM memories ORDER BY weight DESC LIMIT 5",
58 (),
59 )
60 .await?;
61 let mut top_memories = Vec::new();
62 while let Some(row) = top_rows.next().await? {
63 top_memories.push(TopMemory {
64 content: row.get(0)?,
65 weight: row.get(1)?,
66 retrieval_count: row.get(2)?,
67 });
68 }
69 drain_rows(&mut top_rows).await?;
70
71 Ok(StoreStatus {
72 total_memories: total_memories as u32,
73 completed_tasks: completed_tasks as u32,
74 total_tasks: total_tasks as u32,
75 avg_task_score: avg_score,
76 categories,
77 top_memories,
78 })
79 })
80 .await
81 }
82
83 pub async fn list_memories(&self, category: Option<MemoryCategory>) -> Result<Vec<MemoryRecord>> {
85 self.init().await?;
86 let filter = category.map(super::util::category_str);
87 self.with_db(move |conn| async move {
88 let (sql, params): (String, Vec<String>) = if let Some(cat) = filter {
89 (
90 "SELECT id, content, category, weight, retrieval_count, created_at, length(embedding) as emb_len FROM memories WHERE category = ? ORDER BY created_at DESC".into(),
91 vec![cat.to_string()],
92 )
93 } else {
94 (
95 "SELECT id, content, category, weight, retrieval_count, created_at, length(embedding) as emb_len FROM memories ORDER BY created_at DESC".into(),
96 vec![],
97 )
98 };
99
100 let mut rows = if params.is_empty() {
101 conn.query(&sql, ()).await?
102 } else {
103 conn.query(&sql, params![params[0].as_str()]).await?
104 };
105
106 let mut out = Vec::new();
107 while let Some(row) = rows.next().await? {
108 out.push(MemoryRecord {
109 id: row.get(0)?,
110 content: row.get(1)?,
111 category: category_from_str(&row.get::<String>(2)?),
112 weight: row.get(3)?,
113 retrieval_count: row.get(4)?,
114 created_at: row.get(5)?,
115 embedding_status: embedding_status(row.get::<Option<i64>>(6)?, self.dimensions()),
116 });
117 }
118 drain_rows(&mut rows).await?;
119 Ok(out)
120 })
121 .await
122 }
123
124 pub async fn list_tasks(&self, limit: u32) -> Result<Vec<TaskRecord>> {
126 self.init().await?;
127 self.with_db(move |conn| async move {
128 let mut rows = conn
129 .query(
130 r#"
131 SELECT id, description, tokens_used, tool_calls, errors, user_corrections,
132 completed, task_score, started_at, finished_at
133 FROM tasks
134 ORDER BY started_at DESC
135 LIMIT ?
136 "#,
137 params![limit],
138 )
139 .await?;
140
141 struct TaskRow {
142 id: String,
143 description: Option<String>,
144 tokens_used: Option<i64>,
145 tool_calls: Option<i64>,
146 errors: Option<i64>,
147 user_corrections: Option<i64>,
148 completed: Option<i64>,
149 task_score: Option<f64>,
150 started_at: Option<i64>,
151 finished_at: Option<i64>,
152 }
153
154 let mut task_rows = Vec::new();
155 let mut task_ids = Vec::new();
156 while let Some(row) = rows.next().await? {
157 let id: String = row.get(0)?;
158 task_ids.push(id.clone());
159 task_rows.push(TaskRow {
160 id,
161 description: row.get(1)?,
162 tokens_used: row.get(2)?,
163 tool_calls: row.get(3)?,
164 errors: row.get(4)?,
165 user_corrections: row.get(5)?,
166 completed: row.get(6)?,
167 task_score: row.get(7)?,
168 started_at: row.get(8)?,
169 finished_at: row.get(9)?,
170 });
171 }
172 drain_rows(&mut rows).await?;
173
174 let mut retrievals_by_task: HashMap<String, Vec<TaskRetrieval>> = HashMap::new();
175 let mut created_by_task: HashMap<String, Vec<TaskCreatedMemory>> = HashMap::new();
176
177 if !task_ids.is_empty() {
178 let placeholders = std::iter::repeat_n("?", task_ids.len())
179 .collect::<Vec<_>>()
180 .join(", ");
181
182 let retrieval_sql = format!(
183 r#"
184 SELECT r.task_id, r.memory_id, r.similarity, r.self_report, r.credit,
185 substr(m.content, 1, 80) as preview, m.category
186 FROM memory_retrievals r
187 JOIN memories m ON r.memory_id = m.id
188 WHERE r.task_id IN ({placeholders})
189 "#
190 );
191 let mut ret_rows = conn
192 .query(&retrieval_sql, turso::params_from_iter(task_ids.iter().map(String::as_str)))
193 .await?;
194 while let Some(r) = ret_rows.next().await? {
195 let task_id: String = r.get(0)?;
196 let self_report: Option<f64> = r.get(3)?;
197 retrievals_by_task
198 .entry(task_id)
199 .or_default()
200 .push(TaskRetrieval {
201 memory_id: r.get(1)?,
202 similarity: r.get::<Option<f64>>(2)?,
203 self_report: self_report.map(|s| s.round() as u8),
204 credit: r.get(4)?,
205 preview: r.get(5)?,
206 category: category_from_str(&r.get::<String>(6)?),
207 });
208 }
209 drain_rows(&mut ret_rows).await?;
210
211 let created_sql = format!(
212 "SELECT source_task, category, substr(content, 1, 60) as preview FROM memories WHERE source_task IN ({placeholders})"
213 );
214 let mut created_rows = conn
215 .query(&created_sql, turso::params_from_iter(task_ids.iter().map(String::as_str)))
216 .await?;
217 while let Some(c) = created_rows.next().await? {
218 let task_id: String = c.get(0)?;
219 created_by_task.entry(task_id).or_default().push(TaskCreatedMemory {
220 category: category_from_str(&c.get::<String>(1)?),
221 preview: c.get(2)?,
222 });
223 }
224 drain_rows(&mut created_rows).await?;
225 }
226
227 let mut tasks = Vec::with_capacity(task_rows.len());
228 for row in task_rows {
229 let status = match row.finished_at {
230 None => TaskStatus::InProgress,
231 Some(_) if row.completed == Some(1) => TaskStatus::Completed,
232 Some(_) => TaskStatus::Failed,
233 };
234
235 tasks.push(TaskRecord {
236 id: row.id.clone(),
237 description: row.description,
238 tokens_used: row.tokens_used.map(|n| n as u32),
239 tool_calls: row.tool_calls.map(|n| n as u32),
240 errors: row.errors.map(|n| n as u32),
241 user_corrections: row.user_corrections.map(|n| n as u32),
242 status,
243 task_score: row.task_score,
244 started_at: row.started_at,
245 finished_at: row.finished_at,
246 retrievals: retrievals_by_task.remove(&row.id).unwrap_or_default(),
247 created_memories: created_by_task.remove(&row.id).unwrap_or_default(),
248 });
249 }
250 Ok(tasks)
251 })
252 .await
253 }
254
255 pub async fn get_timeline(&self, limit: u32) -> Result<Vec<TimelineEvent>> {
257 self.init().await?;
258 self.with_db(move |conn| async move {
259 let mut events = Vec::new();
260
261 let mut task_rows = conn
262 .query(
263 r#"
264 SELECT description, task_score, tokens_used, errors, completed, started_at
265 FROM tasks ORDER BY started_at DESC LIMIT ?
266 "#,
267 params![limit],
268 )
269 .await?;
270 while let Some(row) = task_rows.next().await? {
271 let started_at: i64 = row.get(5)?;
272 let completed: Option<i64> = row.get(4)?;
273 let status = if completed == Some(1) { "OK" } else { "FAIL" };
274 let score = row
275 .get::<Option<f64>>(1)?
276 .map(|s| format!("{s:.2}"))
277 .unwrap_or_else(|| "?".into());
278 let tokens = row
279 .get::<Option<i64>>(2)?
280 .map(|n| n.to_string())
281 .unwrap_or_else(|| "?".into());
282 let errors = row.get::<Option<i64>>(3)?.unwrap_or(0);
283 let desc: String = row.get::<Option<String>>(0)?.unwrap_or_default();
284 let desc = if desc.len() > 80 {
285 format!("{}...", &desc[..80])
286 } else {
287 desc
288 };
289 events.push(TimelineEvent {
290 timestamp: started_at,
291 kind: TimelineEventKind::Task,
292 summary: format!("TASK [{status}] score={score} {tokens}tok {errors}err — {desc}"),
293 });
294 }
295 drain_rows(&mut task_rows).await?;
296
297 let mut mem_rows = conn
298 .query(
299 "SELECT content, category, weight, created_at FROM memories ORDER BY created_at DESC LIMIT ?",
300 params![limit],
301 )
302 .await?;
303 while let Some(row) = mem_rows.next().await? {
304 let created_at: i64 = row.get(3)?;
305 let category: String = row.get(1)?;
306 let weight: f64 = row.get(2)?;
307 let content: String = row.get(0)?;
308 let preview = if content.len() > 80 {
309 format!("{}...", &content[..80])
310 } else {
311 content
312 };
313 events.push(TimelineEvent {
314 timestamp: created_at,
315 kind: TimelineEventKind::Memory,
316 summary: format!("MEM [{category}] w={weight:.2} — {preview}"),
317 });
318 }
319 drain_rows(&mut mem_rows).await?;
320
321 events.sort_by_key(|e| e.timestamp);
322 Ok(events)
323 })
324 .await
325 }
326
327 pub async fn search_memories(&self, query: &str) -> Result<Vec<Memory>> {
329 self.init().await?;
330 let embedding = (self.embed_fn())(query).await?;
331 let emb_buf = vec_buf(&embedding);
332 let sql = self.retrieval_sql();
333 let decay_rate = self.decay_rate();
334 let top_k = self.top_k();
335 let now = super::store::now_secs();
336
337 self.with_db(move |conn| async move {
338 let mut rows = conn
339 .query(
340 sql.as_ref(),
341 params![emb_buf.as_slice(), emb_buf.as_slice(), decay_rate, now, top_k],
342 )
343 .await?;
344
345 let mut mems = Vec::new();
346 while let Some(row) = rows.next().await? {
347 let distance: f64 = row.get(6)?;
348 mems.push(Memory {
349 id: row.get(0)?,
350 content: row.get(1)?,
351 category: category_from_str(&row.get::<String>(2)?),
352 weight: row.get(3)?,
353 score: 1.0 - distance,
354 created_at: row.get(4)?,
355 retrieval_count: row.get(5)?,
356 });
357 }
358 drain_rows(&mut rows).await?;
359 Ok(mems)
360 })
361 .await
362 }
363
364 pub async fn search(&self, query: &str) -> Result<super::types::StartTaskResult> {
366 self.start_task(query).await
367 }
368}
369
370#[cfg(test)]
371mod tests {
372 use super::super::*;
373 use crate::floppy::types::{
374 MemoryCategory, MemoryReportInput, ReportCorrectionInput, TaskEndInput, UserInputSource,
375 };
376 use std::sync::Arc;
377
378 fn mock_embed() -> EmbedFn {
379 Arc::new(|text: &str| {
380 let text = text.to_string();
381 Box::pin(async move {
382 let mut vec = vec![0.0f32; 4];
383 for (i, b) in text.bytes().enumerate() {
384 vec[i % 4] += b as f32;
385 }
386 let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
387 if norm > 0.0 {
388 for x in &mut vec {
389 *x /= norm;
390 }
391 }
392 Ok(vec)
393 })
394 })
395 }
396
397 struct TestCtx {
398 _dir: tempfile::TempDir,
399 store: MemoryStore,
400 }
401
402 impl TestCtx {
403 fn new() -> Self {
404 let dir = tempfile::tempdir().expect("tempdir");
405 let db_path = dir.path().join("memory.db").to_string_lossy().into_owned();
406 let store = create_memory_store(FloppyConfig::new(db_path, "test").top_k(3).dimensions(4), mock_embed());
407 Self { _dir: dir, store }
408 }
409 }
410
411 #[tokio::test]
412 async fn get_status_includes_categories() {
413 let ctx = TestCtx::new();
414 ctx.store
415 .report_user_input(crate::floppy::ReportUserInput {
416 lesson: "use pnpm".into(),
417 source: UserInputSource::UserInput,
418 })
419 .await
420 .expect("report");
421
422 let status = ctx.store.get_status().await.expect("status");
423 assert_eq!(status.total_memories, 1);
424 assert_eq!(status.categories.len(), 1);
425 assert_eq!(status.categories[0].category, MemoryCategory::User);
426 }
427
428 #[tokio::test]
429 async fn list_memories_filters_category() {
430 let ctx = TestCtx::new();
431 ctx.store
432 .insert_raw_memory("insight note", MemoryCategory::Insight, 1.0)
433 .await
434 .expect("insight");
435 ctx.store
436 .report_user_input(crate::floppy::ReportUserInput {
437 lesson: "user note".into(),
438 source: UserInputSource::UserInput,
439 })
440 .await
441 .expect("user");
442
443 let all = ctx.store.list_memories(None).await.expect("all");
444 assert_eq!(all.len(), 2);
445
446 let users = ctx
447 .store
448 .list_memories(Some(MemoryCategory::User))
449 .await
450 .expect("users");
451 assert_eq!(users.len(), 1);
452 assert_eq!(users[0].category, MemoryCategory::User);
453 }
454
455 #[tokio::test]
456 async fn search_memories_is_read_only() {
457 let ctx = TestCtx::new();
458 let mem_id = ctx
459 .store
460 .report_user_input(crate::floppy::ReportUserInput {
461 lesson: "auth middleware path".into(),
462 source: UserInputSource::UserCorrection,
463 })
464 .await
465 .expect("report");
466
467 let hits = ctx.store.search_memories("auth middleware").await.expect("search");
468 assert!(hits.iter().any(|m| m.id == mem_id));
469
470 let tasks = ctx.store.list_tasks(10).await.expect("tasks");
471 assert!(tasks.is_empty(), "search_memories must not create tasks");
472 }
473
474 #[tokio::test]
475 async fn report_unified_insight_and_end_with_decay() {
476 let ctx = TestCtx::new();
477 let id = ctx
478 .store
479 .report(MemoryReportInput::insight("VDBE architecture"))
480 .await
481 .expect("insight");
482 assert!(!id.is_empty());
483
484 let start = ctx.store.start_task("explore vm").await.expect("start");
485 let result = ctx
486 .store
487 .end_task_with_decay(
488 &start.task_id,
489 TaskEndInput {
490 tokens_used: 100,
491 tool_calls: 1,
492 errors: 0,
493 user_corrections: 0,
494 completed: true,
495 self_report: None,
496 },
497 )
498 .await
499 .expect("end+decay");
500 assert!(result.decay.decayed >= 1);
501
502 let timeline = ctx.store.get_timeline(10).await.expect("timeline");
503 assert!(!timeline.is_empty());
504 }
505
506 #[tokio::test]
507 async fn contradict_wrapper_returns_struct() {
508 let ctx = TestCtx::new();
509 let id = ctx
510 .store
511 .report_correction(ReportCorrectionInput {
512 lesson: "old".into(),
513 what_failed: "a".into(),
514 what_worked: "b".into(),
515 tokens_wasted: None,
516 tools_wasted: None,
517 })
518 .await
519 .expect("report");
520
521 let result = ctx.store.contradict(&id, Some("corrected")).await.expect("contradict");
522 assert!(result.deleted);
523 assert!(result.correction_id.is_some());
524 }
525
526 #[test]
527 fn floppy_paths_project_local() {
528 let paths = FloppyPaths::project_local();
529 assert!(paths.db_path().ends_with(".floppy/memory.db"));
530 let cfg = paths.config("sess");
531 assert_eq!(cfg.session_id, "sess");
532 }
533}