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

Latest change 705e1ce035ea4b9fb4c413dd38f550650185b571 - release: v0.0.5 by Claude Code

use akurai_storage::BTree;
use akurai_vector::{decode, embed, embed_with_bearer, encode, EmbedError};
use notes_core::search::owner_matches;

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], owners: &[&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, owners) {
                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, owners: &[&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)) {
                    owner_matches(&v, owners)
                } else {
                    false
                }
            }
            _ => false,
        }
    }
}

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

pub fn embed_query_with_key(
    text: &str,
    url: &str,
    model: &str,
    api_key: Option<&str>,
) -> Result<Vec<f32>, EmbedError> {
    match api_key {
        Some(token) if !token.trim().is_empty() => embed_with_bearer(url, model, text, token),
        _ => embed_query(text, url, model),
    }
}

pub fn maybe_reembed(
    db: &mut BTree,
    embed_url: &Option<String>,
    embed_model: &Option<String>,
    embed_api_key: &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_with_key(body, url, model, embed_api_key.as_deref()) {
        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
}