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;
12
13use thiserror::Error;
14
15use crate::{
16    config::Config,
17    db::{AimemDb, DbError},
18    embedder::Embedder,
19    search::{SearchError, Searcher},
20    types::Drawer,
21};
22
23#[derive(Debug, Error)]
24pub enum LayerError {
25    #[error("db: {0}")]
26    Db(#[from] DbError),
27    #[error("search: {0}")]
28    Search(#[from] SearchError),
29    #[error("io: {0}")]
30    Io(#[from] std::io::Error),
31}
32
33// ── L0 — Identity ────────────────────────────────────────────────────────────
34
35/// Layer 0: identity text (~100 tokens, always loaded).
36pub async fn l0_render(identity_path: &Path) -> String {
37    if identity_path.exists() {
38        match tokio::fs::read_to_string(identity_path).await {
39            Ok(text) => return text.trim().to_string(),
40            Err(_) => {}
41        }
42    }
43    "## L0 — IDENTITY\nNo identity configured. Create ~/.aimem/identity.txt".to_string()
44}
45
46// ── L1 — Essential Story ─────────────────────────────────────────────────────
47
48const L1_MAX_DRAWERS: usize = 15;
49const L1_MAX_CHARS: usize = 3200;
50
51/// Layer 1: top drawers from AiMem, grouped by room (~500-800 tokens).
52pub async fn l1_generate(db: &AimemDb, wing: Option<&str>) -> Result<String, LayerError> {
53    let drawers = db.fetch_drawers(wing, None, 200).await?;
54
55    if drawers.is_empty() {
56        return Ok("## L1 — No memories yet.".to_string());
57    }
58
59    // Group by room, take top L1_MAX_DRAWERS
60    let mut by_room: HashMap<String, Vec<Drawer>> = HashMap::new();
61    for d in drawers.into_iter().take(L1_MAX_DRAWERS) {
62        by_room.entry(d.room.clone()).or_default().push(d);
63    }
64
65    let mut lines = vec!["## L1 — ESSENTIAL STORY".to_string()];
66    let mut total_len = 0usize;
67
68    let mut rooms: Vec<_> = by_room.iter().collect();
69    rooms.sort_by_key(|(r, _)| r.as_str());
70
71    for (room, entries) in rooms {
72        let room_line = format!("\n[{}]", room);
73        lines.push(room_line.clone());
74        total_len += room_line.len();
75
76        for d in entries {
77            let snippet: String = d
78                .content
79                .trim()
80                .replace('\n', " ")
81                .chars()
82                .take(200)
83                .collect();
84            let snippet = if d.content.len() > 200 {
85                format!("{}...", snippet)
86            } else {
87                snippet
88            };
89
90            let source_tag = d
91                .source_file
92                .as_deref()
93                .and_then(|s| Path::new(s).file_name())
94                .map(|n| format!("  ({})", n.to_string_lossy()))
95                .unwrap_or_default();
96
97            let entry = format!("  - {}{}", snippet, source_tag);
98
99            if total_len + entry.len() > L1_MAX_CHARS {
100                lines.push("  ... (more in L3 search)".to_string());
101                return Ok(lines.join("\n"));
102            }
103            total_len += entry.len();
104            lines.push(entry);
105        }
106    }
107
108    Ok(lines.join("\n"))
109}
110
111// ── L2 — On-Demand ───────────────────────────────────────────────────────────
112
113/// Layer 2: on-demand wing/room filtered retrieval (~200-500 tokens).
114pub async fn l2_retrieve(
115    db: &AimemDb,
116    wing: Option<&str>,
117    room: Option<&str>,
118    n: usize,
119) -> Result<String, LayerError> {
120    let drawers = db.fetch_drawers(wing, room, n).await?;
121
122    if drawers.is_empty() {
123        let label = match (wing, room) {
124            (Some(w), Some(r)) => format!("wing={w} room={r}"),
125            (Some(w), None) => format!("wing={w}"),
126            (None, Some(r)) => format!("room={r}"),
127            (None, None) => "all".to_string(),
128        };
129        return Ok(format!("No drawers found for {label}."));
130    }
131
132    let mut lines = vec![format!("## L2 — ON-DEMAND ({} drawers)", drawers.len())];
133    for d in &drawers {
134        let snippet: String = d
135            .content
136            .trim()
137            .replace('\n', " ")
138            .chars()
139            .take(300)
140            .collect();
141        let snippet = if d.content.len() > 300 {
142            format!("{}...", snippet)
143        } else {
144            snippet
145        };
146        let source_tag = d
147            .source_file
148            .as_deref()
149            .and_then(|s| Path::new(s).file_name())
150            .map(|n| format!("  ({})", n.to_string_lossy()))
151            .unwrap_or_default();
152        lines.push(format!("  [{}] {}{}", d.room, snippet, source_tag));
153    }
154    Ok(lines.join("\n"))
155}
156
157// ── L3 — Deep Search ─────────────────────────────────────────────────────────
158
159/// Layer 3: full semantic search via Turso `vector_distance_cos`.
160pub async fn l3_search(
161    searcher: &Searcher,
162    query: &str,
163    wing: Option<&str>,
164    room: Option<&str>,
165    n: usize,
166) -> Result<String, LayerError> {
167    let results = searcher.vector_search(query, wing, room, n).await?;
168
169    if results.is_empty() {
170        return Ok("No results found.".to_string());
171    }
172
173    let mut lines = vec![format!("## L3 — SEARCH RESULTS for \"{}\"", query)];
174    for (i, r) in results.iter().enumerate() {
175        let snippet: String = r
176            .drawer
177            .content
178            .trim()
179            .replace('\n', " ")
180            .chars()
181            .take(300)
182            .collect();
183        let snippet = if r.drawer.content.len() > 300 {
184            format!("{}...", snippet)
185        } else {
186            snippet
187        };
188        let source = r
189            .drawer
190            .source_file
191            .as_deref()
192            .and_then(|s| Path::new(s).file_name())
193            .map(|n| n.to_string_lossy().to_string())
194            .unwrap_or_default();
195
196        lines.push(format!(
197            "  [{}] {}/{} (sim={:.3})",
198            i + 1,
199            r.drawer.wing,
200            r.drawer.room,
201            r.similarity,
202        ));
203        lines.push(format!("      {}", snippet));
204        if !source.is_empty() {
205            lines.push(format!("      src: {}", source));
206        }
207    }
208    Ok(lines.join("\n"))
209}
210
211// ── MemoryStack — unified interface ──────────────────────────────────────────
212
213/// Full 4-layer memory stack.
214///
215/// ```rust,no_run
216/// use aimem_core::{AimemDb, Config, Embedder, MemoryStack};
217/// # async fn run() -> anyhow::Result<()> {
218/// let cfg = Config::load()?;
219/// let db = AimemDb::open(&cfg.db_path).await?;
220/// let embedder = Embedder::new()?;
221/// let stack = MemoryStack::new(db, embedder, &cfg);
222///
223/// println!("{}", stack.wake_up(None).await?);
224/// println!("{}", stack.recall(Some("my_project"), None).await?);
225/// println!("{}", stack.search("pricing discussion", None, None).await?);
226/// # Ok(())
227/// # }
228/// ```
229#[derive(Debug, Clone)]
230pub struct MemoryStack {
231    db: AimemDb,
232    searcher: Searcher,
233    identity_path: std::path::PathBuf,
234}
235
236impl MemoryStack {
237    pub fn new(db: AimemDb, embedder: Embedder, cfg: &Config) -> Self {
238        let searcher = Searcher::new(db.clone(), embedder);
239        Self {
240            db,
241            searcher,
242            identity_path: cfg.identity_path.clone(),
243        }
244    }
245
246    /// Generate wake-up context: L0 (identity) + L1 (essential story).
247    /// ~600-900 tokens total. Inject into system prompt or first message.
248    pub async fn wake_up(&self, wing: Option<&str>) -> Result<String, LayerError> {
249        let l0 = l0_render(&self.identity_path).await;
250        let l1 = l1_generate(&self.db, wing).await?;
251        Ok(format!("{}\n\n{}", l0, l1))
252    }
253
254    /// On-demand L2 retrieval filtered by wing/room.
255    pub async fn recall(
256        &self,
257        wing: Option<&str>,
258        room: Option<&str>,
259    ) -> Result<String, LayerError> {
260        l2_retrieve(&self.db, wing, room, 10).await
261    }
262
263    /// Deep L3 semantic search.
264    pub async fn search(
265        &self,
266        query: &str,
267        wing: Option<&str>,
268        room: Option<&str>,
269    ) -> Result<String, LayerError> {
270        l3_search(&self.searcher, query, wing, room, 5).await
271    }
272
273    /// Status of the whole stack.
274    pub async fn status(&self) -> Result<serde_json::Value, LayerError> {
275        let count = self.db.drawer_count().await?;
276        let (wings, rooms) = self.db.taxonomy().await?;
277        let l0_exists = self.identity_path.exists();
278        let l0_tokens = if l0_exists {
279            std::fs::read_to_string(&self.identity_path)
280                .map(|s| s.len() / 4)
281                .unwrap_or(0)
282        } else {
283            0
284        };
285
286        Ok(serde_json::json!({
287            "total_drawers": count,
288            "identity_path": self.identity_path.display().to_string(),
289            "identity_exists": l0_exists,
290            "identity_tokens_est": l0_tokens,
291            "wings": wings.into_iter().map(|(w, c)| serde_json::json!({"wing": w, "count": c})).collect::<Vec<_>>(),
292            "rooms": rooms.into_iter().map(|(r, c)| serde_json::json!({"room": r, "count": c})).collect::<Vec<_>>(),
293        }))
294    }
295}
aimem_core/layers.rs

aimem_core/
layers.rs
