Skip to main content

do_memory_storage_redb/
embeddings.rs

1//! Embedding storage operations for redb cache
2
3use crate::{EMBEDDINGS_TABLE, RedbStorage};
4use do_memory_core::{Error, Result};
5use redb::ReadableDatabase;
6use std::sync::Arc;
7use tracing::debug;
8
9impl RedbStorage {
10    /// Store an embedding vector (internal method)
11    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    /// Retrieve an embedding vector (internal method)
46    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}