aimem_core/
layers.rs

1//! 4-layer memory stack — mirrors Python's layers.py.
2//!
3//! ```text
4//! L0  Identity        (~100 tokens)   Always loaded — identity.txt
5//! L1  Essential Story (~500-800 tok)  Top drawers from AiMem
6//! L2  On-Demand       (~200-500 tok)  Wing/room filtered retrieval
7//! L3  Deep Search     (unlimited)     Turso vector_distance_cos semantic search
8//! ```
9
10use std::collections::{HashMap, HashSet};
11use std::path::Path;
12use std::sync::Arc;
13
14use thiserror::Error;
15
16use crate::{
17    config::Config,
18    db::{AimemDb, DbError},
19    embedder::Embedder,
20    search::{SearchError, Searcher},
21    types::{Drawer, DrawerFilingRequest},
22};
23
24#[derive(Debug, Error)]
25pub enum LayerError {
26    #[error("db: {0}")]
27    Db(#[from] DbError),
28    #[error("search: {0}")]
29    Search(#[from] SearchError),
30    #[error("io: {0}")]
31    Io(#[from] std::io::Error),
32}
33
34// ── L0 — Identity ────────────────────────────────────────────────────────────
35
36/// Layer 0: identity text (~100 tokens, always loaded).
37pub async fn l0_render(identity_path: &Path) -> String {
38    if identity_path.exists() {
39        match tokio::fs::read_to_string(identity_path).await {
40            Ok(text) => return text.trim().to_string(),
41            Err(_) => {}
42        }
43    }
44    "## L0 — IDENTITY\nNo identity configured. Create ~/.aimem/identity.txt".to_string()
45}
46
47// ── L1 — Essential Story ─────────────────────────────────────────────────────
48
49const L1_MAX_DRAWERS: usize = 15;
50const L1_MAX_CHARS: usize = 3200;
51
52/// Layer 1: top drawers from AiMem, grouped by room (~500-800 tokens).
53pub async fn l1_generate(db: &AimemDb, wing: Option<&str>) -> Result<String, LayerError> {
54    let drawers = db.fetch_drawers(wing, None, 200).await?;
55
56    if drawers.is_empty() {
57        return Ok("## L1 — No memories yet.".to_string());
58    }
59
60    // Group by room, take top L1_MAX_DRAWERS
61    let mut by_room: HashMap<String, Vec<Drawer>> = HashMap::new();
62    for d in drawers.into_iter().take(L1_MAX_DRAWERS) {
63        by_room.entry(d.room.clone()).or_default().push(d);
64    }
65
66    let mut lines = vec!["## L1 — ESSENTIAL STORY".to_string()];
67    let mut total_len = 0usize;
68
69    let mut rooms: Vec<_> = by_room.iter().collect();
70    rooms.sort_by_key(|(r, _)| r.as_str());
71
72    for (room, entries) in rooms {
73        let room_line = format!("\n[{}]", room);
74        lines.push(room_line.clone());
75        total_len += room_line.len();
76
77        for d in entries {
78            let snippet: String = d
79                .content
80                .trim()
81                .replace('\n', " ")
82                .chars()
83                .take(200)
84                .collect();
85            let snippet = if d.content.len() > 200 {
86                format!("{}...", snippet)
87            } else {
88                snippet
89            };
90
91            let source_tag = d
92                .source_file
93                .as_deref()
94                .and_then(|s| Path::new(s).file_name())
95                .map(|n| format!("  ({})", n.to_string_lossy()))
96                .unwrap_or_default();
97
98            let entry = format!("  - {}{}", snippet, source_tag);
99
100            if total_len + entry.len() > L1_MAX_CHARS {
101                lines.push("  ... (more in L3 search)".to_string());
102                return Ok(lines.join("\n"));
103            }
104            total_len += entry.len();
105            lines.push(entry);
106        }
107    }
108
109    Ok(lines.join("\n"))
110}
111
112// ── L2 — On-Demand ───────────────────────────────────────────────────────────
113
114/// Layer 2: on-demand wing/room filtered retrieval (~200-500 tokens).
115pub async fn l2_retrieve(
116    db: &AimemDb,
117    wing: Option<&str>,
118    room: Option<&str>,
119    n: usize,
120) -> Result<String, LayerError> {
121    let drawers = db.fetch_drawers(wing, room, n).await?;
122
123    if drawers.is_empty() {
124        let label = match (wing, room) {
125            (Some(w), Some(r)) => format!("wing={w} room={r}"),
126            (Some(w), None) => format!("wing={w}"),
127            (None, Some(r)) => format!("room={r}"),
128            (None, None) => "all".to_string(),
129        };
130        return Ok(format!("No drawers found for {label}."));
131    }
132
133    let mut lines = vec![format!("## L2 — ON-DEMAND ({} drawers)", drawers.len())];
134    for d in &drawers {
135        let snippet: String = d
136            .content
137            .trim()
138            .replace('\n', " ")
139            .chars()
140            .take(300)
141            .collect();
142        let snippet = if d.content.len() > 300 {
143            format!("{}...", snippet)
144        } else {
145            snippet
146        };
147        let source_tag = d
148            .source_file
149            .as_deref()
150            .and_then(|s| Path::new(s).file_name())
151            .map(|n| format!("  ({})", n.to_string_lossy()))
152            .unwrap_or_default();
153        lines.push(format!("  [{}] {}{}", d.room, snippet, source_tag));
154    }
155    Ok(lines.join("\n"))
156}
157
158// ── L3 — Deep Search ─────────────────────────────────────────────────────────
159
160/// Layer 3: full semantic search via Turso `vector_distance_cos`.
161pub async fn l3_search(
162    searcher: &Searcher,
163    query: &str,
164    wing: Option<&str>,
165    room: Option<&str>,
166    n: usize,
167) -> Result<String, LayerError> {
168    let results = searcher.vector_search(query, wing, room, n).await?;
169
170    if results.is_empty() {
171        return Ok("No results found.".to_string());
172    }
173
174    let mut lines = vec![format!("## L3 — SEARCH RESULTS for \"{}\"", query)];
175    for (i, r) in results.iter().enumerate() {
176        let snippet: String = r
177            .drawer
178            .content
179            .trim()
180            .replace('\n', " ")
181            .chars()
182            .take(300)
183            .collect();
184        let snippet = if r.drawer.content.len() > 300 {
185            format!("{}...", snippet)
186        } else {
187            snippet
188        };
189        let source = r
190            .drawer
191            .source_file
192            .as_deref()
193            .and_then(|s| Path::new(s).file_name())
194            .map(|n| n.to_string_lossy().to_string())
195            .unwrap_or_default();
196
197        lines.push(format!(
198            "  [{}] {}/{} (sim={:.3})",
199            i + 1,
200            r.drawer.wing,
201            r.drawer.room,
202            r.similarity,
203        ));
204        lines.push(format!("      {}", snippet));
205        if !source.is_empty() {
206            lines.push(format!("      src: {}", source));
207        }
208    }
209    Ok(lines.join("\n"))
210}
211
212// ── MemoryStack — unified interface ──────────────────────────────────────────
213
214/// Full 4-layer memory stack.
215#[derive(Debug, Clone)]
216pub struct MemoryStack {
217    db: AimemDb,
218    searcher: Searcher,
219    identity_path: std::path::PathBuf,
220}
221
222impl MemoryStack {
223    pub fn new(db: AimemDb, embedder: Arc<dyn Embedder>, cfg: &Config) -> Self {
224        let searcher = Searcher::new(db.clone(), embedder);
225        Self {
226            db,
227            searcher,
228            identity_path: cfg.identity_path.clone(),
229        }
230    }
231
232    /// Generate wake-up context: L0 (identity) + L1 (essential story).
233    /// ~600-900 tokens total. Inject into system prompt or first message.
234    pub async fn wake_up(&self, wing: Option<&str>) -> Result<String, LayerError> {
235        let l0 = l0_render(&self.identity_path).await;
236        let l1 = l1_generate(&self.db, wing).await?;
237        Ok(format!("{}\n\n{}", l0, l1))
238    }
239
240    /// On-demand L2 retrieval filtered by wing/room.
241    pub async fn recall(
242        &self,
243        wing: Option<&str>,
244        room: Option<&str>,
245    ) -> Result<String, LayerError> {
246        l2_retrieve(&self.db, wing, room, 10).await
247    }
248
249    /// Deep L3 semantic search.
250    pub async fn search(
251        &self,
252        query: &str,
253        wing: Option<&str>,
254        room: Option<&str>,
255    ) -> Result<String, LayerError> {
256        l3_search(&self.searcher, query, wing, room, 5).await
257    }
258
259    /// File a new memory (drawer) into the L1-L3 store.
260    /// This handles embedding generation and DB insertion automatically.
261    pub async fn file_drawer(
262        &self,
263        wing: &str,
264        room: &str,
265        content: String,
266        parts: Vec<crate::types::ContentPart>,
267        agent: &str,
268    ) -> Result<String, LayerError> {
269        let now = chrono::Utc::now().to_rfc3339();
270        let digest = md5::compute(format!("{wing}{room}{content}{now}").as_bytes());
271        let id = format!("mem_{wing}_{digest:x}");
272
273        self.file_drawer_with_id(&id, wing, room, content, parts, None, 0, agent)
274            .await?;
275        Ok(id)
276    }
277
278    /// File a new memory with a caller-supplied stable drawer ID and optional metadata.
279    /// Returns `true` when a new drawer row was inserted and `false` when the ID already existed.
280    pub async fn file_drawer_with_id(
281        &self,
282        id: &str,
283        wing: &str,
284        room: &str,
285        content: String,
286        parts: Vec<crate::types::ContentPart>,
287        source_file: Option<&str>,
288        chunk_index: i64,
289        agent: &str,
290    ) -> Result<bool, LayerError> {
291        let request = build_filing_request(
292            id,
293            wing,
294            room,
295            content,
296            parts,
297            source_file,
298            chunk_index,
299            agent,
300        );
301        let mut results = self.file_drawers_with_ids(&[request]).await?;
302        Ok(results.pop().unwrap_or(false))
303    }
304
305    /// File multiple stable-ID drawers in one embedding batch.
306    ///
307    /// This is useful for attachment-style ingestion where a caller has a set of
308    /// related chunks (for example one file's summary + chunk drawers) and wants
309    /// to avoid one remote embedding request per drawer.
310    ///
311    /// The returned vector is aligned with `requests`: `true` means a new drawer
312    /// row was inserted, `false` means the stable ID already existed.
313    ///
314    /// Existing IDs are checked before embedding so retries do not spend another
315    /// remote embedding call on already-filed drawers.
316    pub async fn file_drawers_with_ids(
317        &self,
318        requests: &[DrawerFilingRequest],
319    ) -> Result<Vec<bool>, LayerError> {
320        if requests.is_empty() {
321            return Ok(Vec::new());
322        }
323
324        let mut results = vec![false; requests.len()];
325        let mut pending_indexes = Vec::new();
326        let mut pending_parts = Vec::new();
327        let mut seen_ids = HashSet::new();
328
329        for (index, request) in requests.iter().enumerate() {
330            if !seen_ids.insert(request.id.clone()) {
331                continue;
332            }
333            if self.db.drawer_exists(&request.id).await? {
334                continue;
335            }
336            pending_indexes.push(index);
337            pending_parts.push(parts_for_embedding(request));
338        }
339
340        let embedder = self.searcher.embedder();
341        let embeddings = if let Some(ref embedder) = embedder {
342            if pending_parts.is_empty() {
343                Some(Vec::new())
344            } else {
345                let embeddings = embedder
346                    .embed(&pending_parts)
347                    .await
348                    .map_err(SearchError::from)?;
349                if embeddings.len() != pending_parts.len() {
350                    return Err(SearchError::EmbedBatchSizeMismatch {
351                        expected: pending_parts.len(),
352                        actual: embeddings.len(),
353                    }
354                    .into());
355                }
356                Some(embeddings)
357            }
358        } else {
359            None
360        };
361
362        let default_filed_at = chrono::Utc::now().to_rfc3339();
363
364        for (pending_offset, request_index) in pending_indexes.into_iter().enumerate() {
365            let request = &requests[request_index];
366            let drawer = drawer_from_request(request, &default_filed_at);
367            let inserted = if let (Some(embedder), Some(embeddings)) =
368                (embedder.as_ref(), embeddings.as_ref())
369            {
370                self.db
371                    .insert_drawer_with_profile(
372                        &drawer,
373                        Some(embeddings[pending_offset].as_slice()),
374                        embedder.provider_name(),
375                        embedder.model_name(),
376                    )
377                    .await?
378            } else {
379                self.db.insert_drawer(&drawer, None).await?
380            };
381            results[request_index] = inserted;
382        }
383
384        Ok(results)
385    }
386
387    /// File a plain-text memory without manually constructing a parts vector.
388    pub async fn file_text(
389        &self,
390        wing: &str,
391        room: &str,
392        content: impl Into<String>,
393        agent: &str,
394    ) -> Result<String, LayerError> {
395        self.file_drawer(wing, room, content.into(), Vec::new(), agent)
396            .await
397    }
398
399    /// Access the underlying searcher.
400    pub fn searcher(&self) -> &Searcher {
401        &self.searcher
402    }
403
404    /// Status of the whole stack.
405    pub async fn status(&self) -> Result<serde_json::Value, LayerError> {
406        let count = self.db.drawer_count().await?;
407        let (wings, rooms) = self.db.taxonomy().await?;
408        let l0_exists = self.identity_path.exists();
409        let l0_tokens = if l0_exists {
410            std::fs::read_to_string(&self.identity_path)
411                .map(|s| s.len() / 4)
412                .unwrap_or(0)
413        } else {
414            0
415        };
416
417        Ok(serde_json::json!({
418            "total_drawers": count,
419            "identity_path": self.identity_path.display().to_string(),
420            "identity_exists": l0_exists,
421            "identity_tokens_est": l0_tokens,
422            "wings": wings.into_iter().map(|(w, c)| serde_json::json!({"wing": w, "count": c})).collect::<Vec<_>>(),
423            "rooms": rooms.into_iter().map(|(r, c)| serde_json::json!({"room": r, "count": c})).collect::<Vec<_>>(),
424        }))
425    }
426}
427
428fn build_filing_request(
429    id: &str,
430    wing: &str,
431    room: &str,
432    content: String,
433    parts: Vec<crate::types::ContentPart>,
434    source_file: Option<&str>,
435    chunk_index: i64,
436    agent: &str,
437) -> DrawerFilingRequest {
438    let mut request = DrawerFilingRequest::new(id, wing, room, content, agent)
439        .with_parts(parts)
440        .with_chunk_index(chunk_index);
441    if let Some(source_file) = source_file {
442        request = request.with_source_file(source_file);
443    }
444    request
445}
446
447fn parts_for_embedding(request: &DrawerFilingRequest) -> Vec<crate::types::ContentPart> {
448    if request.parts.is_empty() {
449        vec![crate::types::ContentPart::text(request.content.clone())]
450    } else {
451        request.parts.clone()
452    }
453}
454
455fn drawer_from_request(request: &DrawerFilingRequest, default_filed_at: &str) -> Drawer {
456    let mut drawer = Drawer::new(
457        request.id.clone(),
458        request.wing.clone(),
459        request.room.clone(),
460        request.content.clone(),
461        request.added_by.clone(),
462    )
463    .with_parts(request.parts.clone())
464    .with_chunk_index(request.chunk_index)
465    .with_filed_at(
466        request
467            .filed_at
468            .clone()
469            .unwrap_or_else(|| default_filed_at.to_string()),
470    );
471
472    if let Some(source_file) = request.source_file.as_deref() {
473        drawer = drawer.with_source_file(source_file);
474    }
475
476    drawer
477}
aimem_core/layers.rs

aimem_core/
layers.rs
