aimem_core/convo/
mod.rs

1//! Conversation miner — mine chat exports into AiMem.
2//!
3//! Supported formats (via normalize module):
4//!   - Plain text with `>` markers
5//!   - Claude.ai JSON export
6//!   - ChatGPT conversations.json
7//!   - Claude Code JSONL
8//!   - Slack JSON export
9
10pub mod normalize;
11
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14
15use thiserror::Error;
16use walkdir::WalkDir;
17
18use crate::{
19    db::{AimemDb, DbError},
20    embedder::{EmbedError, Embedder},
21    miner::{MIN_CHUNK, chunk_text},
22    types::{ContentPart, Drawer},
23};
24
25#[derive(Debug, Error)]
26pub enum ConvoMinerError {
27    #[error("db: {0}")]
28    Db(#[from] DbError),
29    #[error("embed: {0}")]
30    Embed(#[from] EmbedError),
31    #[error("io: {0}")]
32    Io(#[from] std::io::Error),
33}
34
35const CONVO_EXTENSIONS: &[&str] = &["txt", "md", "json", "jsonl"];
36const SKIP_DIRS: &[&str] = &[
37    ".git",
38    "node_modules",
39    "__pycache__",
40    ".venv",
41    "venv",
42    "env",
43    "dist",
44    "build",
45    ".aimem",
46    "target",
47];
48
49#[derive(Debug, Default)]
50pub struct ConvoMineStats {
51    pub files_scanned: usize,
52    pub files_skipped: usize,
53    pub drawers_added: usize,
54}
55
56/// Conversation miner.
57#[derive(Clone)]
58pub struct ConvoMiner {
59    db: AimemDb,
60    embedder: Option<Arc<dyn Embedder>>,
61}
62
63impl std::fmt::Debug for ConvoMiner {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("ConvoMiner")
66            .field("db", &self.db)
67            .field(
68                "embedder",
69                &self.embedder.as_ref().map(|_| "Some(Arc<dyn Embedder>)"),
70            )
71            .finish()
72    }
73}
74
75impl ConvoMiner {
76    pub fn new(db: AimemDb, embedder: Option<Arc<dyn Embedder>>) -> Self {
77        Self { db, embedder }
78    }
79
80    /// Mine a directory of conversation exports into AiMem.
81    pub async fn mine(
82        &self,
83        convo_dir: impl AsRef<Path>,
84        wing: &str,
85        room: &str,
86        agent: &str,
87        limit: usize,
88        dry_run: bool,
89    ) -> Result<ConvoMineStats, ConvoMinerError> {
90        let convo_dir = convo_dir.as_ref();
91        let files = scan_convo_files(convo_dir);
92        let files: Vec<_> = if limit > 0 {
93            files.into_iter().take(limit).collect()
94        } else {
95            files
96        };
97
98        let mut stats = ConvoMineStats::default();
99        stats.files_scanned = files.len();
100
101        for filepath in &files {
102            let source = filepath.to_string_lossy().to_string();
103
104            if !dry_run && self.db.source_already_mined(&source).await? {
105                stats.files_skipped += 1;
106                continue;
107            }
108
109            let content = match std::fs::read_to_string(filepath) {
110                Ok(c) => c,
111                Err(_) => continue,
112            };
113
114            // Normalize to plain transcript
115            let normalized = normalize::normalize_content(&content, filepath);
116            if normalized.trim().len() < MIN_CHUNK {
117                continue;
118            }
119
120            let chunks = chunk_exchanges(&normalized);
121            if chunks.is_empty() {
122                continue;
123            }
124
125            if dry_run {
126                tracing::info!(
127                    "[DRY RUN] {} → {}/{} ({} chunks)",
128                    filepath.file_name().unwrap_or_default().to_string_lossy(),
129                    wing,
130                    room,
131                    chunks.len()
132                );
133                stats.drawers_added += chunks.len();
134                continue;
135            }
136
137            let embeddings: Vec<Option<Vec<f32>>> = if let Some(ref emb) = self.embedder {
138                let inputs: Vec<Vec<ContentPart>> =
139                    chunks.iter().map(|c| vec![ContentPart::text(c)]).collect();
140                match emb.embed(&inputs).await {
141                    Ok(vecs) => vecs.into_iter().map(Some).collect(),
142                    Err(e) => {
143                        tracing::warn!("embedding failed for {source}: {e}");
144                        vec![None; chunks.len()]
145                    }
146                }
147            } else {
148                vec![None; chunks.len()]
149            };
150
151            let now = chrono::Utc::now().to_rfc3339();
152            for (i, (chunk, embedding)) in chunks.iter().zip(embeddings.iter()).enumerate() {
153                let id = {
154                    let input = format!("{wing}{room}{source}{i}");
155                    let digest = md5::compute(input.as_bytes());
156                    format!("drawer_{wing}_{room}_{digest:x}")
157                };
158                let drawer = Drawer::new(
159                    id,
160                    wing.to_string(),
161                    room.to_string(),
162                    chunk.clone(),
163                    agent.to_string(),
164                )
165                .with_source_file(source.clone())
166                .with_chunk_index(i as i64)
167                .with_filed_at(now.clone());
168                let inserted = if let (Some(embedding), Some(embedder)) =
169                    (embedding.as_deref(), self.embedder.as_ref())
170                {
171                    self.db
172                        .insert_drawer_with_profile(
173                            &drawer,
174                            Some(embedding),
175                            embedder.provider_name(),
176                            embedder.model_name(),
177                        )
178                        .await?
179                } else {
180                    self.db.insert_drawer(&drawer, embedding.as_deref()).await?
181                };
182                if inserted {
183                    stats.drawers_added += 1;
184                }
185            }
186        }
187
188        Ok(stats)
189    }
190}
191
192/// Chunk a conversation by exchange pairs (user turn + AI response = 1 unit).
193pub fn chunk_exchanges(content: &str) -> Vec<String> {
194    let lines: Vec<&str> = content.lines().collect();
195    let quote_count = lines
196        .iter()
197        .filter(|l| l.trim_start().starts_with('>'))
198        .count();
199
200    if quote_count >= 3 {
201        chunk_by_exchange(&lines)
202    } else {
203        chunk_text(content)
204    }
205}
206
207fn chunk_by_exchange(lines: &[&str]) -> Vec<String> {
208    let mut chunks = Vec::new();
209    let mut i = 0;
210
211    while i < lines.len() {
212        let line = lines[i].trim();
213        if line.starts_with('>') {
214            let user_turn = line.to_string();
215            i += 1;
216
217            let mut ai_lines = Vec::new();
218            while i < lines.len() {
219                let next = lines[i].trim();
220                if next.starts_with('>') || next == "---" {
221                    break;
222                }
223                ai_lines.push(lines[i]);
224                i += 1;
225            }
226
227            let ai_response = ai_lines.join("\n");
228            let chunk = format!("{}\n{}", user_turn, ai_response);
229            if chunk.trim().len() >= MIN_CHUNK {
230                chunks.push(chunk);
231            }
232        } else {
233            i += 1;
234        }
235    }
236
237    // Fallback: if we got nothing, just use paragraph chunking
238    if chunks.is_empty() {
239        chunk_text(&lines.join("\n"))
240    } else {
241        chunks
242    }
243}
244
245fn scan_convo_files(dir: &Path) -> Vec<PathBuf> {
246    WalkDir::new(dir)
247        .follow_links(false)
248        .into_iter()
249        .filter_entry(|e| {
250            let name = e.file_name().to_string_lossy();
251            !SKIP_DIRS.contains(&name.as_ref())
252        })
253        .filter_map(|e| e.ok())
254        .filter(|e| e.file_type().is_file())
255        .filter(|e| {
256            let ext = e
257                .path()
258                .extension()
259                .map(|x| x.to_string_lossy().to_lowercase())
260                .unwrap_or_default();
261            CONVO_EXTENSIONS.contains(&ext.as_str())
262        })
263        .map(|e| e.into_path())
264        .collect()
265}
aimem_core/convo/mod.rs

aimem_core/convo/
mod.rs
