do_memory_storage_redb/
embeddings.rs1use crate::{EMBEDDINGS_TABLE, RedbStorage};
4use do_memory_core::{Error, Result};
5use redb::ReadableDatabase;
6use std::sync::Arc;
7use tracing::debug;
8
9impl RedbStorage {
10 pub async fn store_embedding_raw(&self, id: &str, embedding: &[f32]) -> Result<()> {
12 debug!("Storing embedding: {}", id);
13 let db = Arc::clone(&self.db);
14 let id_str = id.to_string();
15 let embedding_bytes = postcard::to_allocvec(embedding)
16 .map_err(|e| Error::Storage(format!("Failed to serialize embedding: {}", e)))?;
17
18 tokio::task::spawn_blocking(move || {
19 let write_txn = db
20 .begin_write()
21 .map_err(|e| Error::Storage(format!("Failed to begin write transaction: {}", e)))?;
22
23 {
24 let mut table = write_txn.open_table(EMBEDDINGS_TABLE).map_err(|e| {
25 Error::Storage(format!("Failed to open embeddings table: {}", e))
26 })?;
27
28 table
29 .insert(id_str.as_str(), embedding_bytes.as_slice())
30 .map_err(|e| Error::Storage(format!("Failed to insert embedding: {}", e)))?;
31 }
32
33 write_txn
34 .commit()
35 .map_err(|e| Error::Storage(format!("Failed to commit transaction: {}", e)))?;
36
37 Ok::<(), Error>(())
38 })
39 .await
40 .map_err(|e| Error::Storage(format!("Task join error: {}", e)))??;
41
42 Ok(())
43 }
44
45 pub async fn get_embedding_raw(&self, id: &str) -> Result<Option<Vec<f32>>> {
47 debug!("Retrieving embedding: {}", id);
48 let db = Arc::clone(&self.db);
49 let id_str = id.to_string();
50
51 tokio::task::spawn_blocking(move || {
52 let read_txn = db
53 .begin_read()
54 .map_err(|e| Error::Storage(format!("Failed to begin read transaction: {}", e)))?;
55
56 let table = read_txn
57 .open_table(EMBEDDINGS_TABLE)
58 .map_err(|e| Error::Storage(format!("Failed to open embeddings table: {}", e)))?;
59
60 match table
61 .get(id_str.as_str())
62 .map_err(|e| Error::Storage(format!("Failed to get embedding: {}", e)))?
63 {
64 Some(bytes_guard) => {
65 let _bytes = bytes_guard.value();
66 let embedding: Vec<f32> =
67 postcard::from_bytes(bytes_guard.value()).map_err(|e| {
68 Error::Storage(format!("Failed to deserialize embedding: {}", e))
69 })?;
70 Ok(Some(embedding))
71 }
72 None => Ok(None),
73 }
74 })
75 .await
76 .map_err(|e| Error::Storage(format!("Task join error: {}", e)))?
77 }
78}