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 hybrid keyword + vector 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: hybrid keyword + semantic search via Turso FTS + vectors.
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.hybrid_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            "  [{}] {}/{} (score={:.3}{}{})",
199            i + 1,
200            r.drawer.wing,
201            r.drawer.room,
202            r.score,
203            r.semantic_similarity
204                .map(|sim| format!(", sim={sim:.3}"))
205                .unwrap_or_default(),
206            r.keyword_score
207                .map(|score| format!(", kw={score:.3}"))
208                .unwrap_or_default(),
209        ));
210        lines.push(format!("      {}", snippet));
211        if !source.is_empty() {
212            lines.push(format!("      src: {}", source));
213        }
214    }
215    Ok(lines.join("\n"))
216}
217
218// ── MemoryStack — unified interface ──────────────────────────────────────────
219
220/// Full 4-layer memory stack.
221#[derive(Debug, Clone)]
222pub struct MemoryStack {
223    db: AimemDb,
224    searcher: Searcher,
225    identity_path: std::path::PathBuf,
226}
227
228impl MemoryStack {
229    pub fn new(db: AimemDb, embedder: Arc<dyn Embedder>, cfg: &Config) -> Self {
230        let searcher = Searcher::new(db.clone(), embedder);
231        Self {
232            db,
233            searcher,
234            identity_path: cfg.identity_path.clone(),
235        }
236    }
237
238    /// Generate wake-up context: L0 (identity) + L1 (essential story).
239    /// ~600-900 tokens total. Inject into system prompt or first message.
240    pub async fn wake_up(&self, wing: Option<&str>) -> Result<String, LayerError> {
241        let l0 = l0_render(&self.identity_path).await;
242        let l1 = l1_generate(&self.db, wing).await?;
243        Ok(format!("{}\n\n{}", l0, l1))
244    }
245
246    /// On-demand L2 retrieval filtered by wing/room.
247    pub async fn recall(
248        &self,
249        wing: Option<&str>,
250        room: Option<&str>,
251    ) -> Result<String, LayerError> {
252        l2_retrieve(&self.db, wing, room, 10).await
253    }
254
255    /// Deep L3 hybrid search.
256    pub async fn search(
257        &self,
258        query: &str,
259        wing: Option<&str>,
260        room: Option<&str>,
261    ) -> Result<String, LayerError> {
262        l3_search(&self.searcher, query, wing, room, 5).await
263    }
264
265    /// File a new memory (drawer) into the L1-L3 store.
266    /// This handles embedding generation and DB insertion automatically.
267    pub async fn file_drawer(
268        &self,
269        wing: &str,
270        room: &str,
271        content: String,
272        parts: Vec<crate::types::ContentPart>,
273        agent: &str,
274    ) -> Result<String, LayerError> {
275        let now = chrono::Utc::now().to_rfc3339();
276        let digest = md5::compute(format!("{wing}{room}{content}{now}").as_bytes());
277        let id = format!("mem_{wing}_{digest:x}");
278
279        self.file_drawer_with_id(&id, wing, room, content, parts, None, 0, agent)
280            .await?;
281        Ok(id)
282    }
283
284    /// File a new memory with a caller-supplied stable drawer ID and optional metadata.
285    /// Returns `true` when a new drawer row was inserted and `false` when the ID already existed.
286    pub async fn file_drawer_with_id(
287        &self,
288        id: &str,
289        wing: &str,
290        room: &str,
291        content: String,
292        parts: Vec<crate::types::ContentPart>,
293        source_file: Option<&str>,
294        chunk_index: i64,
295        agent: &str,
296    ) -> Result<bool, LayerError> {
297        let request = build_filing_request(
298            id,
299            wing,
300            room,
301            content,
302            parts,
303            source_file,
304            chunk_index,
305            agent,
306        );
307        let mut results = self.file_drawers_with_ids(&[request]).await?;
308        Ok(results.pop().unwrap_or(false))
309    }
310
311    /// File multiple stable-ID drawers in one embedding batch.
312    ///
313    /// This is useful for attachment-style ingestion where a caller has a set of
314    /// related chunks (for example one file's summary + chunk drawers) and wants
315    /// to avoid one remote embedding request per drawer.
316    ///
317    /// The returned vector is aligned with `requests`: `true` means a new drawer
318    /// row was inserted, `false` means the stable ID already existed.
319    ///
320    /// Existing IDs are checked before embedding so retries do not spend another
321    /// remote embedding call on already-filed drawers.
322    pub async fn file_drawers_with_ids(
323        &self,
324        requests: &[DrawerFilingRequest],
325    ) -> Result<Vec<bool>, LayerError> {
326        if requests.is_empty() {
327            return Ok(Vec::new());
328        }
329
330        let mut results = vec![false; requests.len()];
331        let mut pending_indexes = Vec::new();
332        let mut pending_parts = Vec::new();
333        let mut seen_ids = HashSet::new();
334
335        for (index, request) in requests.iter().enumerate() {
336            if !seen_ids.insert(request.id.clone()) {
337                continue;
338            }
339            if self.db.drawer_exists(&request.id).await? {
340                continue;
341            }
342            pending_indexes.push(index);
343            pending_parts.push(parts_for_embedding(request));
344        }
345
346        let embedder = self.searcher.embedder();
347        let embeddings = if let Some(ref embedder) = embedder {
348            if pending_parts.is_empty() {
349                Some(Vec::new())
350            } else {
351                let embeddings = embedder
352                    .embed(&pending_parts)
353                    .await
354                    .map_err(SearchError::from)?;
355                if embeddings.len() != pending_parts.len() {
356                    return Err(SearchError::EmbedBatchSizeMismatch {
357                        expected: pending_parts.len(),
358                        actual: embeddings.len(),
359                    }
360                    .into());
361                }
362                Some(embeddings)
363            }
364        } else {
365            None
366        };
367
368        let default_filed_at = chrono::Utc::now().to_rfc3339();
369
370        for (pending_offset, request_index) in pending_indexes.into_iter().enumerate() {
371            let request = &requests[request_index];
372            let drawer = drawer_from_request(request, &default_filed_at);
373            let inserted = if let (Some(embedder), Some(embeddings)) =
374                (embedder.as_ref(), embeddings.as_ref())
375            {
376                self.db
377                    .insert_drawer_with_profile(
378                        &drawer,
379                        Some(embeddings[pending_offset].as_slice()),
380                        embedder.provider_name(),
381                        embedder.model_name(),
382                    )
383                    .await?
384            } else {
385                self.db.insert_drawer(&drawer, None).await?
386            };
387            results[request_index] = inserted;
388        }
389
390        Ok(results)
391    }
392
393    /// File a plain-text memory without manually constructing a parts vector.
394    pub async fn file_text(
395        &self,
396        wing: &str,
397        room: &str,
398        content: impl Into<String>,
399        agent: &str,
400    ) -> Result<String, LayerError> {
401        self.file_drawer(wing, room, content.into(), Vec::new(), agent)
402            .await
403    }
404
405    /// Access the underlying searcher.
406    pub fn searcher(&self) -> &Searcher {
407        &self.searcher
408    }
409
410    /// Status of the whole stack.
411    pub async fn status(&self) -> Result<serde_json::Value, LayerError> {
412        let count = self.db.drawer_count().await?;
413        let (wings, rooms) = self.db.taxonomy().await?;
414        let l0_exists = self.identity_path.exists();
415        let l0_tokens = if l0_exists {
416            std::fs::read_to_string(&self.identity_path)
417                .map(|s| s.len() / 4)
418                .unwrap_or(0)
419        } else {
420            0
421        };
422
423        Ok(serde_json::json!({
424            "total_drawers": count,
425            "identity_path": self.identity_path.display().to_string(),
426            "identity_exists": l0_exists,
427            "identity_tokens_est": l0_tokens,
428            "wings": wings.into_iter().map(|(w, c)| serde_json::json!({"wing": w, "count": c})).collect::<Vec<_>>(),
429            "rooms": rooms.into_iter().map(|(r, c)| serde_json::json!({"room": r, "count": c})).collect::<Vec<_>>(),
430        }))
431    }
432}
433
434fn build_filing_request(
435    id: &str,
436    wing: &str,
437    room: &str,
438    content: String,
439    parts: Vec<crate::types::ContentPart>,
440    source_file: Option<&str>,
441    chunk_index: i64,
442    agent: &str,
443) -> DrawerFilingRequest {
444    let mut request = DrawerFilingRequest::new(id, wing, room, content, agent)
445        .with_parts(parts)
446        .with_chunk_index(chunk_index);
447    if let Some(source_file) = source_file {
448        request = request.with_source_file(source_file);
449    }
450    request
451}
452
453fn parts_for_embedding(request: &DrawerFilingRequest) -> Vec<crate::types::ContentPart> {
454    if request.parts.is_empty() {
455        vec![crate::types::ContentPart::text(request.content.clone())]
456    } else {
457        request.parts.clone()
458    }
459}
460
461fn drawer_from_request(request: &DrawerFilingRequest, default_filed_at: &str) -> Drawer {
462    let mut drawer = Drawer::new(
463        request.id.clone(),
464        request.wing.clone(),
465        request.room.clone(),
466        request.content.clone(),
467        request.added_by.clone(),
468    )
469    .with_parts(request.parts.clone())
470    .with_chunk_index(request.chunk_index)
471    .with_filed_at(
472        request
473            .filed_at
474            .clone()
475            .unwrap_or_else(|| default_filed_at.to_string()),
476    );
477
478    if let Some(source_file) = request.source_file.as_deref() {
479        drawer = drawer.with_source_file(source_file);
480    }
481
482    drawer
483}
aimem_core/layers.rs

aimem_core/
layers.rs
