AkurAI-Notes / crates/notes-api/src/embed.rs

Latest change 3594ca5510bcac2b26e45d9235a28905856ebf61 - E05: On-device embeddings + semantic search — embed.rs + search.rs + re-embed on update by Claude Code

use akurai_storage::BTree;
use akurai_vector::{decode, embed, encode, EmbedError};

pub struct EmbedStore<'a> {
    pub db: &'a mut BTree,
}

impl<'a> EmbedStore<'a> {
    pub fn store(&mut self, note_id: i64, vec: &[f32]) -> Result<(), String> {
        let key = format!("note_emb:{}", note_id);
        let bytes = encode(vec);
        self.db
            .insert(key.as_bytes(), &bytes)
            .map_err(|e| e.to_string())?;
        self.db.commit().map_err(|e| e.to_string())?;
        Ok(())
    }

    pub fn load(&mut self, note_id: i64) -> Option<Vec<f32>> {
        let key = format!("note_emb:{}", note_id);
        match self.db.get(key.as_bytes()) {
            Ok(Some(bytes)) => decode(&bytes),
            _ => None,
        }
    }

    pub fn remove(&mut self, note_id: i64) -> Result<(), String> {
        let key = format!("note_emb:{}", note_id);
        self.db
            .delete(key.as_bytes())
            .map_err(|e| e.to_string())?;
        self.db.commit().map_err(|e| e.to_string())?;
        Ok(())
    }

    pub fn rank_notes(
        &mut self,
        qv: &[f32],
        owner_sub: &str,
        k: usize,
    ) -> Vec<(i64, f32)> {
        let prefix = b"note_emb:";
        let end = upper_bound(prefix);

        let entries = match self.db.range(prefix, &end) {
            Ok(e) => e,
            Err(_) => return Vec::new(),
        };

        let mut candidates: Vec<(u64, Vec<f32>)> = Vec::new();
        for (key, val) in entries {
            let id_str = match std::str::from_utf8(&key[b"note_emb:".len()..]) {
                Ok(s) => s,
                Err(_) => continue,
            };
            let id: i64 = match id_str.parse() {
                Ok(i) => i,
                Err(_) => continue,
            };
            if !self.note_is_owned_by(id, owner_sub) {
                continue;
            }
            if let Some(doc) = decode(&val) {
                candidates.push((id as u64, doc));
            }
        }

        let scored = akurai_vector::rank(qv, &candidates, k);
        scored
            .into_iter()
            .map(|(id, score)| (id as i64, score))
            .collect()
    }

    fn note_is_owned_by(&mut self, id: i64, want: &str) -> bool {
        let mut key = b"notes:".to_vec();
        key.extend_from_slice(&(id as u64).to_be_bytes());
        match self.db.get(&key) {
            Ok(Some(bytes)) => {
                if let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&bytes)) {
                    v.get("ownerSub").and_then(|x| x.as_str()) == Some(want)
                } else {
                    false
                }
            }
            _ => false,
        }
    }
}

pub fn embed_query(text: &str, url: &str, model: &str) -> Result<Vec<f32>, EmbedError> {
    embed(url, model, text)
}

pub fn maybe_reembed(
    db: &mut BTree,
    embed_url: &Option<String>,
    embed_model: &Option<String>,
    note_id: i64,
    body: &str,
) {
    let url = match embed_url {
        Some(u) => u,
        None => return,
    };
    let model = embed_model.as_deref().unwrap_or("qwen3-embedding-0.6b");
    match embed_query(body, url, model) {
        Ok(vec) => {
            let mut store = EmbedStore { db };
            let _ = store.store(note_id, &vec);
        }
        Err(e) => eprintln!("re-embed failed for note {note_id}: {e}"),
    }
}

fn upper_bound(prefix: &[u8]) -> Vec<u8> {
    let mut bound = prefix.to_vec();
    bound.push(0xff);
    bound
}