elph_core/memz/
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 status (`elph memory status`).
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 (`elph memory list`).
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)?),
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 (`elph memory tasks`).
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 (`elph memory log`).
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().to_string();
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,
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 (`elph memory search` — creates a task).
365    pub async fn search(&self, query: &str) -> Result<super::types::StartTaskResult> {
366        let _ = self.embed_pending().await?;
367        self.start_task(query).await
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use super::super::*;
374    use crate::memz::types::{MemoryCategory, MemoryReportInput, ReportCorrectionInput, TaskEndInput, UserInputSource};
375    use std::sync::Arc;
376
377    fn mock_embed() -> EmbedFn {
378        Arc::new(|text: &str| {
379            let text = text.to_string();
380            Box::pin(async move {
381                let mut vec = vec![0.0f32; 4];
382                for (i, b) in text.bytes().enumerate() {
383                    vec[i % 4] += b as f32;
384                }
385                let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
386                if norm > 0.0 {
387                    for x in &mut vec {
388                        *x /= norm;
389                    }
390                }
391                Ok(vec)
392            })
393        })
394    }
395
396    struct TestCtx {
397        _dir: tempfile::TempDir,
398        store: MemoryStore,
399    }
400
401    impl TestCtx {
402        fn new() -> Self {
403            let dir = tempfile::tempdir().expect("tempdir");
404            let db_path = dir.path().join("memory.db").to_string_lossy().into_owned();
405            let store = create_memory_store(MemzConfig::new(db_path, "test").top_k(3).dimensions(4), mock_embed());
406            Self { _dir: dir, store }
407        }
408    }
409
410    #[tokio::test]
411    async fn get_status_includes_categories() {
412        let ctx = TestCtx::new();
413        ctx.store
414            .report_user_input(crate::memz::ReportUserInput {
415                lesson: "use pnpm".into(),
416                source: UserInputSource::UserInput,
417            })
418            .await
419            .expect("report");
420
421        let status = ctx.store.get_status().await.expect("status");
422        assert_eq!(status.total_memories, 1);
423        assert_eq!(status.categories.len(), 1);
424        assert_eq!(status.categories[0].category, MemoryCategory::User);
425    }
426
427    #[tokio::test]
428    async fn list_memories_filters_category() {
429        let ctx = TestCtx::new();
430        ctx.store
431            .insert_raw_memory("insight note", MemoryCategory::Insight, 1.0)
432            .await
433            .expect("insight");
434        ctx.store
435            .report_user_input(crate::memz::ReportUserInput {
436                lesson: "user note".into(),
437                source: UserInputSource::UserInput,
438            })
439            .await
440            .expect("user");
441
442        let all = ctx.store.list_memories(None).await.expect("all");
443        assert_eq!(all.len(), 2);
444
445        let users = ctx
446            .store
447            .list_memories(Some(MemoryCategory::User))
448            .await
449            .expect("users");
450        assert_eq!(users.len(), 1);
451        assert_eq!(users[0].category, MemoryCategory::User);
452    }
453
454    #[tokio::test]
455    async fn search_memories_is_read_only() {
456        let ctx = TestCtx::new();
457        let mem_id = ctx
458            .store
459            .report_user_input(crate::memz::ReportUserInput {
460                lesson: "auth middleware path".into(),
461                source: UserInputSource::UserCorrection,
462            })
463            .await
464            .expect("report");
465
466        let hits = ctx.store.search_memories("auth middleware").await.expect("search");
467        assert!(hits.iter().any(|m| m.id == mem_id));
468
469        let tasks = ctx.store.list_tasks(10).await.expect("tasks");
470        assert!(tasks.is_empty(), "search_memories must not create tasks");
471    }
472
473    #[tokio::test]
474    async fn report_unified_insight_and_end_with_decay() {
475        let ctx = TestCtx::new();
476        let id = ctx
477            .store
478            .report(MemoryReportInput::insight("VDBE architecture"))
479            .await
480            .expect("insight");
481        assert!(!id.is_empty());
482
483        let start = ctx.store.start_task("explore vm").await.expect("start");
484        let result = ctx
485            .store
486            .end_task_with_decay(
487                &start.task_id,
488                TaskEndInput {
489                    tokens_used: 100,
490                    tool_calls: 1,
491                    errors: 0,
492                    user_corrections: 0,
493                    completed: true,
494                    self_report: None,
495                },
496            )
497            .await
498            .expect("end+decay");
499        assert!(result.decay.decayed >= 1);
500
501        let timeline = ctx.store.get_timeline(10).await.expect("timeline");
502        assert!(!timeline.is_empty());
503    }
504
505    #[tokio::test]
506    async fn contradict_wrapper_returns_struct() {
507        let ctx = TestCtx::new();
508        let id = ctx
509            .store
510            .report_correction(ReportCorrectionInput {
511                lesson: "old".into(),
512                what_failed: "a".into(),
513                what_worked: "b".into(),
514                tokens_wasted: None,
515                tools_wasted: None,
516            })
517            .await
518            .expect("report");
519
520        let result = ctx.store.contradict(&id, Some("corrected")).await.expect("contradict");
521        assert!(result.deleted);
522        assert!(result.correction_id.is_some());
523    }
524
525    #[test]
526    fn memz_paths_project_local() {
527        let paths = MemzPaths::project_local();
528        assert!(paths.db_path().ends_with(".memz/memory.db"));
529        let cfg = paths.config("sess");
530        assert_eq!(cfg.session_id, "sess");
531    }
532}
elph_core/memz/query.rs

elph_core/memz/
query.rs
