Skip to main content

elph_core/floppy/
query.rs

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    /// Extended store statistics (counts, categories, top memories, task metrics).
15    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    /// List memories, optionally filtered by category.
84    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    /// List recent tasks with retrievals and memories created during each task.
125    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    /// Merged timeline of tasks and memory events.
256    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    /// Read-only semantic search — no task record, no retrieval side effects.
328    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    /// Semantic search via full task lifecycle (creates a task record).
365    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}