1use std::path::Path;
46
47use thiserror::Error;
48use turso::{Builder, Connection, Database};
49
50use crate::embedder::{
51 GEMINI_EMBED_MODEL, GEMINI_EMBED_PROVIDER, LOCAL_EMBED_MODEL, LOCAL_EMBED_PROVIDER,
52};
53use crate::types::{Drawer, Entity, Triple};
54
55#[derive(Debug, Error)]
56pub enum DbError {
57 #[error("Turso error: {0}")]
58 Turso(#[from] turso::Error),
59 #[error("I/O error: {0}")]
60 Io(#[from] std::io::Error),
61 #[error("Type conversion: {0}")]
62 Conversion(String),
63 #[error("Serialization error: {0}")]
64 Serde(#[from] serde_json::Error),
65 #[error("Embedding dimension mismatch: store uses {expected}, attempted {actual}")]
66 EmbeddingDimensionMismatch { expected: usize, actual: usize },
67 #[error(
68 "Embedding model mismatch: store uses {expected_provider}/{expected_model}, attempted {actual_provider}/{actual_model}"
69 )]
70 EmbeddingModelMismatch {
71 expected_provider: String,
72 expected_model: String,
73 actual_provider: String,
74 actual_model: String,
75 },
76}
77
78pub type DbResult<T> = Result<T, DbError>;
79
80#[derive(Debug, Clone, Default, PartialEq, Eq)]
81pub struct EmbeddingStoreProfile {
82 pub provider: Option<String>,
83 pub model: Option<String>,
84 pub dimension: Option<usize>,
85}
86
87const INIT_SQL: &str = "
89CREATE TABLE IF NOT EXISTS drawers (
90 id TEXT PRIMARY KEY,
91 wing TEXT NOT NULL,
92 room TEXT NOT NULL,
93 content TEXT NOT NULL,
94 parts TEXT NOT NULL DEFAULT '[]',
95 embedding BLOB,
96 source_file TEXT,
97 chunk_index INTEGER NOT NULL DEFAULT 0,
98 added_by TEXT NOT NULL DEFAULT 'aimem',
99 filed_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
100);
101
102CREATE INDEX IF NOT EXISTS idx_drawers_wing ON drawers(wing);
103CREATE INDEX IF NOT EXISTS idx_drawers_room ON drawers(room);
104CREATE INDEX IF NOT EXISTS idx_drawers_wing_room ON drawers(wing, room);
105CREATE INDEX IF NOT EXISTS idx_drawers_source ON drawers(source_file);
106
107CREATE TABLE IF NOT EXISTS aimem_meta (
108 key TEXT PRIMARY KEY,
109 value TEXT NOT NULL
110);
111
112CREATE TABLE IF NOT EXISTS entities (
113 id TEXT PRIMARY KEY,
114 name TEXT NOT NULL,
115 type TEXT NOT NULL DEFAULT 'unknown',
116 properties TEXT NOT NULL DEFAULT '{}',
117 created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
118);
119
120CREATE TABLE IF NOT EXISTS triples (
121 id TEXT PRIMARY KEY,
122 subject TEXT NOT NULL,
123 predicate TEXT NOT NULL,
124 object TEXT NOT NULL,
125 valid_from TEXT,
126 valid_to TEXT,
127 confidence REAL NOT NULL DEFAULT 1.0,
128 source_closet TEXT,
129 source_file TEXT,
130 extracted_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
131 FOREIGN KEY (subject) REFERENCES entities(id),
132 FOREIGN KEY (object) REFERENCES entities(id)
133);
134
135CREATE INDEX IF NOT EXISTS idx_triples_subject ON triples(subject);
136CREATE INDEX IF NOT EXISTS idx_triples_object ON triples(object);
137CREATE INDEX IF NOT EXISTS idx_triples_predicate ON triples(predicate);
138CREATE INDEX IF NOT EXISTS idx_triples_valid ON triples(valid_from, valid_to);
139";
140
141#[derive(Clone)]
145pub struct AimemDb {
146 db: Database,
147}
148
149impl std::fmt::Debug for AimemDb {
150 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151 f.debug_struct("AimemDb").finish_non_exhaustive()
152 }
153}
154
155impl AimemDb {
156 pub async fn open(path: impl AsRef<Path>) -> DbResult<Self> {
158 let path = path.as_ref();
159 if let Some(parent) = path.parent() {
160 std::fs::create_dir_all(parent)?;
161 }
162
163 let path_str = path.to_string_lossy();
164 let db = Builder::new_local(path_str.as_ref()).build().await?;
165 let aimem_db = Self { db };
166 aimem_db.migrate().await?;
167 Ok(aimem_db)
168 }
169
170 pub async fn memory() -> DbResult<Self> {
172 let db = Builder::new_local(":memory:").build().await?;
173 let aimem_db = Self { db };
174 aimem_db.migrate().await?;
175 Ok(aimem_db)
176 }
177
178 pub fn conn(&self) -> DbResult<Connection> {
180 Ok(self.db.connect()?)
181 }
182
183 async fn migrate(&self) -> DbResult<()> {
185 let conn = self.conn()?;
186 conn.execute_batch(INIT_SQL).await?;
187
188 if let Err(err) = conn
189 .execute(
190 "ALTER TABLE drawers ADD COLUMN parts TEXT NOT NULL DEFAULT '[]'",
191 (),
192 )
193 .await
194 {
195 let msg = err.to_string();
196 if !msg.contains("duplicate column name: parts") {
197 return Err(err.into());
198 }
199 }
200
201 Ok(())
202 }
203
204 async fn meta_value(&self, key: &str) -> DbResult<Option<String>> {
205 let conn = self.conn()?;
206 let mut rows = conn
207 .query("SELECT value FROM aimem_meta WHERE key = ?1 LIMIT 1", [key])
208 .await?;
209 let Some(row) = rows.next().await? else {
210 return Ok(None);
211 };
212 Ok(Some(val_to_string(&row, 0)?))
213 }
214
215 async fn set_meta_value(&self, key: &str, value: &str) -> DbResult<()> {
216 let conn = self.conn()?;
217 conn.execute(
218 "INSERT OR REPLACE INTO aimem_meta (key, value) VALUES (?1, ?2)",
219 [key, value],
220 )
221 .await?;
222 Ok(())
223 }
224
225 fn infer_profile_from_dimension(dim: usize) -> Option<(&'static str, &'static str)> {
226 match dim {
227 384 => Some((LOCAL_EMBED_PROVIDER, LOCAL_EMBED_MODEL)),
228 768 => Some((GEMINI_EMBED_PROVIDER, GEMINI_EMBED_MODEL)),
229 _ => None,
230 }
231 }
232
233 pub async fn embedding_dimension(&self) -> DbResult<Option<usize>> {
236 if let Some(raw) = self.meta_value("embedding_dim").await? {
237 let dim = raw.parse::<usize>().map_err(|err| {
238 DbError::Conversion(format!("invalid embedding_dim metadata {raw:?}: {err}"))
239 })?;
240 return Ok(Some(dim));
241 }
242
243 let conn = self.conn()?;
244
245 let mut rows = conn
246 .query(
247 "SELECT length(embedding) FROM drawers WHERE embedding IS NOT NULL LIMIT 1",
248 (),
249 )
250 .await?;
251 let Some(row) = rows.next().await? else {
252 return Ok(None);
253 };
254
255 let bytes = row.get_value(0)?.as_integer().copied().ok_or_else(|| {
256 DbError::Conversion("length(embedding) did not return an integer".to_string())
257 })?;
258 if bytes <= 0 || bytes % 4 != 0 {
259 return Err(DbError::Conversion(format!(
260 "unexpected embedding blob length {bytes}; cannot infer dimension"
261 )));
262 }
263
264 let dim = (bytes / 4) as usize;
265 self.set_meta_value("embedding_dim", &dim.to_string())
266 .await?;
267 Ok(Some(dim))
268 }
269
270 pub async fn embedding_profile(&self) -> DbResult<EmbeddingStoreProfile> {
271 let dimension = self.embedding_dimension().await?;
272 let mut provider = self.meta_value("embedding_provider").await?;
273 let mut model = self.meta_value("embedding_model").await?;
274
275 if let Some(dim) = dimension {
276 if let Some((inferred_provider, inferred_model)) =
277 Self::infer_profile_from_dimension(dim)
278 {
279 if provider.is_none() {
280 self.set_meta_value("embedding_provider", inferred_provider)
281 .await?;
282 provider = Some(inferred_provider.to_string());
283 }
284 if model.is_none() {
285 self.set_meta_value("embedding_model", inferred_model)
286 .await?;
287 model = Some(inferred_model.to_string());
288 }
289 }
290 }
291
292 Ok(EmbeddingStoreProfile {
293 provider,
294 model,
295 dimension,
296 })
297 }
298
299 pub async fn assert_embedding_dimension(&self, dim: usize) -> DbResult<()> {
301 if let Some(expected) = self.embedding_dimension().await? {
302 if expected != dim {
303 return Err(DbError::EmbeddingDimensionMismatch {
304 expected,
305 actual: dim,
306 });
307 }
308 }
309 Ok(())
310 }
311
312 pub async fn assert_embedding_profile(
313 &self,
314 dim: usize,
315 provider: &str,
316 model: &str,
317 ) -> DbResult<()> {
318 let profile = self.embedding_profile().await?;
319 if let Some(expected) = profile.dimension {
320 if expected != dim {
321 return Err(DbError::EmbeddingDimensionMismatch {
322 expected,
323 actual: dim,
324 });
325 }
326 }
327
328 if let (Some(expected_provider), Some(expected_model)) = (profile.provider, profile.model) {
329 if expected_provider != provider || expected_model != model {
330 return Err(DbError::EmbeddingModelMismatch {
331 expected_provider,
332 expected_model,
333 actual_provider: provider.to_string(),
334 actual_model: model.to_string(),
335 });
336 }
337 }
338
339 Ok(())
340 }
341
342 async fn ensure_embedding_dimension_for_write(&self, dim: usize) -> DbResult<()> {
344 let conn = self.conn()?;
345 if let Some(expected) = self.embedding_dimension().await? {
346 if expected != dim {
347 return Err(DbError::EmbeddingDimensionMismatch {
348 expected,
349 actual: dim,
350 });
351 }
352 return Ok(());
353 }
354
355 conn.execute(
356 "INSERT OR REPLACE INTO aimem_meta (key, value) VALUES ('embedding_dim', ?1)",
357 [dim.to_string()],
358 )
359 .await?;
360 Ok(())
361 }
362
363 async fn ensure_embedding_profile_for_write(
364 &self,
365 dim: usize,
366 provider: &str,
367 model: &str,
368 ) -> DbResult<()> {
369 let profile = self.embedding_profile().await?;
370 if let Some(expected) = profile.dimension {
371 if expected != dim {
372 return Err(DbError::EmbeddingDimensionMismatch {
373 expected,
374 actual: dim,
375 });
376 }
377 }
378
379 if let (Some(expected_provider), Some(expected_model)) =
380 (profile.provider.clone(), profile.model.clone())
381 {
382 if expected_provider != provider || expected_model != model {
383 return Err(DbError::EmbeddingModelMismatch {
384 expected_provider,
385 expected_model,
386 actual_provider: provider.to_string(),
387 actual_model: model.to_string(),
388 });
389 }
390 }
391
392 self.set_meta_value("embedding_dim", &dim.to_string())
393 .await?;
394 self.set_meta_value("embedding_provider", provider).await?;
395 self.set_meta_value("embedding_model", model).await?;
396 Ok(())
397 }
398
399 pub async fn insert_drawer(
405 &self,
406 drawer: &Drawer,
407 embedding: Option<&[f32]>,
408 ) -> DbResult<bool> {
409 if let Some(v) = embedding {
410 if let Some((provider, model)) = Self::infer_profile_from_dimension(v.len()) {
411 self.ensure_embedding_profile_for_write(v.len(), provider, model)
412 .await?;
413 } else {
414 self.ensure_embedding_dimension_for_write(v.len()).await?;
415 }
416 }
417
418 self.insert_drawer_inner(drawer, embedding).await
419 }
420
421 pub async fn insert_drawer_with_profile(
423 &self,
424 drawer: &Drawer,
425 embedding: Option<&[f32]>,
426 provider: &str,
427 model: &str,
428 ) -> DbResult<bool> {
429 if let Some(v) = embedding {
430 self.ensure_embedding_profile_for_write(v.len(), provider, model)
431 .await?;
432 }
433
434 self.insert_drawer_inner(drawer, embedding).await
435 }
436
437 async fn insert_drawer_inner(
438 &self,
439 drawer: &Drawer,
440 embedding: Option<&[f32]>,
441 ) -> DbResult<bool> {
442 let conn = self.conn()?;
443
444 let emb_json: Option<String> = embedding.map(|v| {
446 let nums: Vec<String> = v.iter().map(|f| f.to_string()).collect();
447 format!("[{}]", nums.join(","))
448 });
449 let parts_json = serde_json::to_string(&drawer.parts)?;
450
451 let sql = if emb_json.is_some() {
453 "INSERT OR IGNORE INTO drawers \
454 (id, wing, room, content, parts, embedding, source_file, chunk_index, added_by, filed_at) \
455 VALUES (?1, ?2, ?3, ?4, ?5, vector32(?6), ?7, ?8, ?9, ?10)"
456 } else {
457 "INSERT OR IGNORE INTO drawers \
458 (id, wing, room, content, parts, embedding, source_file, chunk_index, added_by, filed_at) \
459 VALUES (?1, ?2, ?3, ?4, ?5, NULL, ?6, ?7, ?8, ?9)"
460 };
461
462 let rows_affected = if let Some(ref emb) = emb_json {
463 conn.execute(
464 sql,
465 turso::params![
466 drawer.id.as_str(),
467 drawer.wing.as_str(),
468 drawer.room.as_str(),
469 drawer.content.as_str(),
470 parts_json.as_str(),
471 emb.as_str(),
472 drawer.source_file.as_deref(),
473 drawer.chunk_index,
474 drawer.added_by.as_str(),
475 drawer.filed_at.as_str(),
476 ],
477 )
478 .await?
479 } else {
480 conn.execute(
481 sql,
482 turso::params![
483 drawer.id.as_str(),
484 drawer.wing.as_str(),
485 drawer.room.as_str(),
486 drawer.content.as_str(),
487 parts_json.as_str(),
488 drawer.source_file.as_deref(),
489 drawer.chunk_index,
490 drawer.added_by.as_str(),
491 drawer.filed_at.as_str(),
492 ],
493 )
494 .await?
495 };
496
497 Ok(rows_affected > 0)
498 }
499
500 pub async fn source_already_mined(&self, source_file: &str) -> DbResult<bool> {
502 let conn = self.conn()?;
503 let mut rows = conn
504 .query(
505 "SELECT 1 FROM drawers WHERE source_file = ?1 LIMIT 1",
506 [source_file],
507 )
508 .await?;
509 Ok(rows.next().await?.is_some())
510 }
511
512 pub async fn delete_drawer(&self, id: &str) -> DbResult<bool> {
514 let conn = self.conn()?;
515 let n = conn
516 .execute("DELETE FROM drawers WHERE id = ?1", [id])
517 .await?;
518 Ok(n > 0)
519 }
520
521 pub async fn find_drawers_by_exact_content(
523 &self,
524 content: &str,
525 limit: usize,
526 ) -> DbResult<Vec<Drawer>> {
527 let conn = self.conn()?;
528 let mut rows = conn
529 .query(
530 "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
531 FROM drawers WHERE content = ?1 ORDER BY filed_at DESC LIMIT ?2",
532 turso::params![content, limit as i64],
533 )
534 .await?;
535
536 let mut drawers = Vec::new();
537 while let Some(row) = rows.next().await? {
538 drawers.push(row_to_drawer(&row)?);
539 }
540 Ok(drawers)
541 }
542
543 pub async fn drawer_count(&self) -> DbResult<i64> {
545 let conn = self.conn()?;
546 let mut rows = conn.query("SELECT COUNT(*) FROM drawers", ()).await?;
547 let row = rows
548 .next()
549 .await?
550 .ok_or_else(|| DbError::Conversion("COUNT returned no row".to_string()))?;
551 Ok(row.get_value(0)?.as_integer().copied().unwrap_or(0))
552 }
553
554 pub async fn taxonomy(&self) -> DbResult<(Vec<(String, i64)>, Vec<(String, i64)>)> {
556 let conn = self.conn()?;
557
558 let mut wing_rows = conn
559 .query(
560 "SELECT wing, COUNT(*) as cnt FROM drawers GROUP BY wing ORDER BY cnt DESC",
561 (),
562 )
563 .await?;
564 let mut wings: Vec<(String, i64)> = Vec::new();
565 while let Some(row) = wing_rows.next().await? {
566 let wing = match row.get_value(0)? {
567 turso::Value::Text(s) => s,
568 turso::Value::Null => String::new(),
569 v => format!("{v:?}"),
570 };
571 let cnt = row.get_value(1)?.as_integer().copied().unwrap_or(0);
572 wings.push((wing, cnt));
573 }
574
575 let mut room_rows = conn
576 .query(
577 "SELECT room, COUNT(*) as cnt FROM drawers GROUP BY room ORDER BY cnt DESC",
578 (),
579 )
580 .await?;
581 let mut rooms: Vec<(String, i64)> = Vec::new();
582 while let Some(row) = room_rows.next().await? {
583 let room = match row.get_value(0)? {
584 turso::Value::Text(s) => s,
585 turso::Value::Null => String::new(),
586 v => format!("{v:?}"),
587 };
588 let cnt = row.get_value(1)?.as_integer().copied().unwrap_or(0);
589 rooms.push((room, cnt));
590 }
591
592 Ok((wings, rooms))
593 }
594
595 pub async fn fetch_drawers(
597 &self,
598 wing: Option<&str>,
599 room: Option<&str>,
600 limit: usize,
601 ) -> DbResult<Vec<Drawer>> {
602 let conn = self.conn()?;
603 let limit = limit as i64;
604
605 let (sql, params_vec): (String, Vec<String>) = match (wing, room) {
606 (Some(w), Some(r)) => (
607 "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
608 FROM drawers WHERE wing = ?1 AND room = ?2 \
609 ORDER BY filed_at DESC LIMIT ?3"
610 .to_string(),
611 vec![w.to_string(), r.to_string(), limit.to_string()],
612 ),
613 (Some(w), None) => (
614 "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
615 FROM drawers WHERE wing = ?1 \
616 ORDER BY filed_at DESC LIMIT ?2"
617 .to_string(),
618 vec![w.to_string(), limit.to_string()],
619 ),
620 (None, Some(r)) => (
621 "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
622 FROM drawers WHERE room = ?1 \
623 ORDER BY filed_at DESC LIMIT ?2"
624 .to_string(),
625 vec![r.to_string(), limit.to_string()],
626 ),
627 (None, None) => (
628 "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
629 FROM drawers ORDER BY filed_at DESC LIMIT ?1"
630 .to_string(),
631 vec![limit.to_string()],
632 ),
633 };
634
635 let params: Vec<&str> = params_vec.iter().map(|s| s.as_str()).collect();
636 let mut rows = match params.len() {
637 1 => conn.query(&sql, [params[0]]).await?,
638 2 => conn.query(&sql, [params[0], params[1]]).await?,
639 3 => conn.query(&sql, [params[0], params[1], params[2]]).await?,
640 _ => unreachable!(),
641 };
642
643 let mut drawers = Vec::new();
644 while let Some(row) = rows.next().await? {
645 drawers.push(row_to_drawer(&row)?);
646 }
647 Ok(drawers)
648 }
649
650 pub async fn upsert_entity(&self, entity: &Entity) -> DbResult<()> {
656 let conn = self.conn()?;
657 conn.execute(
658 "INSERT OR REPLACE INTO entities (id, name, type, properties, created_at) \
659 VALUES (?1, ?2, ?3, ?4, ?5)",
660 turso::params![
661 entity.id.as_str(),
662 entity.name.as_str(),
663 entity.entity_type.as_str(),
664 entity.properties.as_str(),
665 entity.created_at.as_str(),
666 ],
667 )
668 .await?;
669 Ok(())
670 }
671
672 pub async fn insert_triple(&self, triple: &Triple) -> DbResult<bool> {
674 let conn = self.conn()?;
675 let n = conn
676 .execute(
677 "INSERT OR IGNORE INTO triples \
678 (id, subject, predicate, object, valid_from, valid_to, confidence, \
679 source_closet, source_file, extracted_at) \
680 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
681 turso::params![
682 triple.id.as_str(),
683 triple.subject.as_str(),
684 triple.predicate.as_str(),
685 triple.object.as_str(),
686 triple.valid_from.as_deref(),
687 triple.valid_to.as_deref(),
688 triple.confidence,
689 triple.source_closet.as_deref(),
690 triple.source_file.as_deref(),
691 triple.extracted_at.as_str(),
692 ],
693 )
694 .await?;
695 Ok(n > 0)
696 }
697
698 pub async fn invalidate_triple(
700 &self,
701 subject: &str,
702 predicate: &str,
703 object: &str,
704 ended: &str,
705 ) -> DbResult<u64> {
706 let conn = self.conn()?;
707 let n = conn
708 .execute(
709 "UPDATE triples SET valid_to = ?1 \
710 WHERE subject = ?2 AND predicate = ?3 AND object = ?4 AND valid_to IS NULL",
711 turso::params![ended, subject, predicate, object],
712 )
713 .await?;
714 Ok(n)
715 }
716}
717
718fn row_to_drawer(row: &turso::Row) -> DbResult<Drawer> {
723 Ok(Drawer {
724 id: val_to_string(row, 0)?,
725 wing: val_to_string(row, 1)?,
726 room: val_to_string(row, 2)?,
727 content: val_to_string(row, 3)?,
728 parts: parse_parts(&val_to_string(row, 4)?)?,
729 source_file: {
730 let s = val_to_string(row, 5)?;
731 if s.is_empty() { None } else { Some(s) }
732 },
733 chunk_index: row.get_value(6)?.as_integer().copied().unwrap_or(0),
734 added_by: val_to_string(row, 7)?,
735 filed_at: val_to_string(row, 8)?,
736 })
737}
738
739fn parse_parts(raw: &str) -> Result<Vec<crate::types::ContentPart>, serde_json::Error> {
740 if raw.is_empty() || raw == "[]" {
741 Ok(Vec::new())
742 } else {
743 serde_json::from_str(raw)
744 }
745}
746
747fn val_to_string(row: &turso::Row, idx: usize) -> DbResult<String> {
748 match row.get_value(idx)? {
749 turso::Value::Text(s) => Ok(s),
750 turso::Value::Null => Ok(String::new()),
751 turso::Value::Integer(i) => Ok(i.to_string()),
752 turso::Value::Real(f) => Ok(f.to_string()),
753 turso::Value::Blob(_) => Ok(String::new()),
754 }
755}