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;
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,
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 parts_for_embedding = if parts.is_empty() {
270            vec![crate::types::ContentPart::text(content.clone())]
271        } else {
272            parts.clone()
273        };
274
275        let embedding = if let Some(ref emb) = self.searcher.embedder() {
276            let vecs = emb
277                .embed(&[parts_for_embedding])
278                .await
279                .map_err(SearchError::from)?;
280            vecs.into_iter().next()
281        } else {
282            None
283        };
284
285        let now = chrono::Utc::now().to_rfc3339();
286        let digest = md5::compute(format!("{wing}{room}{content}{now}").as_bytes());
287        let id = format!("mem_{wing}_{digest:x}");
288
289        let drawer = Drawer::new(id.clone(), wing, room, content, agent)
290            .with_parts(parts)
291            .with_filed_at(now);
292
293        if let (Some(embedding), Some(embedder)) = (embedding.as_deref(), self.searcher.embedder())
294        {
295            self.db
296                .insert_drawer_with_profile(
297                    &drawer,
298                    Some(embedding),
299                    embedder.provider_name(),
300                    embedder.model_name(),
301                )
302                .await?;
303        } else {
304            self.db.insert_drawer(&drawer, embedding.as_deref()).await?;
305        }
306        Ok(id)
307    }
308
309    /// File a plain-text memory without manually constructing a parts vector.
310    pub async fn file_text(
311        &self,
312        wing: &str,
313        room: &str,
314        content: impl Into<String>,
315        agent: &str,
316    ) -> Result<String, LayerError> {
317        self.file_drawer(wing, room, content.into(), Vec::new(), agent)
318            .await
319    }
320
321    /// Access the underlying searcher.
322    pub fn searcher(&self) -> &Searcher {
323        &self.searcher
324    }
325
326    /// Status of the whole stack.
327    pub async fn status(&self) -> Result<serde_json::Value, LayerError> {
328        let count = self.db.drawer_count().await?;
329        let (wings, rooms) = self.db.taxonomy().await?;
330        let l0_exists = self.identity_path.exists();
331        let l0_tokens = if l0_exists {
332            std::fs::read_to_string(&self.identity_path)
333                .map(|s| s.len() / 4)
334                .unwrap_or(0)
335        } else {
336            0
337        };
338
339        Ok(serde_json::json!({
340            "total_drawers": count,
341            "identity_path": self.identity_path.display().to_string(),
342            "identity_exists": l0_exists,
343            "identity_tokens_est": l0_tokens,
344            "wings": wings.into_iter().map(|(w, c)| serde_json::json!({"wing": w, "count": c})).collect::<Vec<_>>(),
345            "rooms": rooms.into_iter().map(|(r, c)| serde_json::json!({"room": r, "count": c})).collect::<Vec<_>>(),
346        }))
347    }
348}
aimem_core/layers.rs

aimem_core/
layers.rs
