Menu
Public repository / 5591f87904dfed6a2e00c83c31b44e5e0b8063e1
AkurAI-Notes / crates/notes-api/src/handlers/search.rs
Latest change 3594ca5510bcac2b26e45d9235a28905856ebf61 - E05: On-device embeddings + semantic search — embed.rs + search.rs + re-embed on update by Claude Code
use crate::auth::session;
use crate::embed::{embed_query, EmbedStore};
use crate::handlers::RouteHandler;
use crate::router_setup::{json_response, upper_bound, NotesState};
use akurai_http::{Request, Response};
use akurai_json::Value;
use akurai_storage::BTree;
use std::sync::{Arc, Mutex};
pub fn search_route(state: Arc<Mutex<NotesState>>) -> RouteHandler {
Box::new(move |req: &Request| -> Response {
let query_str = req.query.as_deref().unwrap_or("");
let mut q = String::new();
let mut k: usize = 20;
for pair in query_str.split('&') {
if let Some((key, value)) = pair.split_once('=') {
match key {
"q" => q = value.to_string(),
"k" => k = value.parse().unwrap_or(20),
_ => {}
}
}
}
if q.is_empty() {
return json_response(Value::Object(vec![
("query".into(), Value::Str("".into())),
("mode".into(), Value::Str("none".into())),
("results".into(), Value::Array(vec![])),
]));
}
// Extract session and embed config in one locked scope
let (user, embed_url, embed_model) = {
let state_lock = match state.lock() {
Ok(s) => s,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
let mut db = match state_lock.db.lock() {
Ok(db) => db,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
let sid = match session::extract_session_id(req) {
Some(s) => s,
None => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("unauthorized".into())),
]));
}
};
let user = match session::get_user(&mut db, &sid) {
Some(u) => u,
None => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("unauthorized".into())),
]));
}
};
(user, state_lock.embed_url.clone(), state_lock.embed_model.clone())
};
// Try semantic search
if let Some(ref url) = embed_url {
let model = embed_model.as_deref().unwrap_or("qwen3-embedding-0.6b");
match embed_query(&q, url, model) {
Ok(qv) => {
let state_lock = match state.lock() {
Ok(s) => s,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
let mut db = match state_lock.db.lock() {
Ok(db) => db,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
let mut store = EmbedStore { db: &mut *db };
let ranked = store.rank_notes(&qv, &user.sub, k);
let results: Vec<Value> = ranked
.into_iter()
.map(|(id, score)| note_to_result(&mut *store.db, id, score))
.collect();
return json_response(Value::Object(vec![
("query".into(), Value::Str(q)),
("mode".into(), Value::Str("semantic".into())),
("results".into(), Value::Array(results)),
]));
}
Err(e) => {
eprintln!("embed query failed: {e} — falling back to substring");
}
}
}
// Substring fallback
let results = {
let state_lock = match state.lock() {
Ok(s) => s,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
let mut db = match state_lock.db.lock() {
Ok(db) => db,
Err(_) => {
return json_response(Value::Object(vec![
("error".into(), Value::Str("internal error".into())),
]));
}
};
substring_search(&mut *db, &q, &user.sub, k)
};
json_response(Value::Object(vec![
("query".into(), Value::Str(q)),
("mode".into(), Value::Str("substring".into())),
("results".into(), Value::Array(results)),
]))
})
}
fn substring_search(
db: &mut BTree,
q: &str,
owner_sub: &str,
k: usize,
) -> Vec<Value> {
let q_lower = q.to_lowercase();
let prefix = b"notes:";
let end = upper_bound(prefix);
let entries = match db.range(prefix, &end) {
Ok(e) => e,
Err(_) => return Vec::new(),
};
let mut scored: Vec<(i64, f32, String, String)> = Vec::new();
for entry in entries {
let (key, val) = entry;
let rest = &key[prefix.len()..];
let id: i64 = match rest.try_into() {
Ok(bytes) => i64::from_be_bytes(bytes),
Err(_) => continue,
};
let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&val)) else {
continue;
};
if v.get("ownerSub").and_then(|x| x.as_str()) != Some(owner_sub) {
continue;
}
let title = v
.get("title")
.and_then(|x| x.as_str())
.unwrap_or("")
.to_string();
let body = v
.get("body")
.and_then(|x| x.as_str())
.unwrap_or("")
.to_string();
let combined = format!("{} {}", title, body).to_lowercase();
if combined.contains(&q_lower) {
let count = combined.matches(&q_lower).count() as f32;
scored.push((id, count, title, body));
}
}
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(k);
scored
.into_iter()
.map(|(id, _score, title, body)| {
let snippet = if body.len() > 180 {
format!("{}…", &body[..180])
} else {
body.clone()
};
Value::Object(vec![
("id".into(), Value::Int(id)),
("title".into(), Value::Str(title)),
("snippet".into(), Value::Str(snippet)),
("score".into(), Value::Float(0.0)),
])
})
.collect()
}
fn note_to_result(db: &mut BTree, id: i64, score: f32) -> Value {
let mut key = b"notes:".to_vec();
key.extend_from_slice(&(id as u64).to_be_bytes());
let mut title = String::new();
let mut body = String::new();
let mut notebook_id = 0i64;
let mut updated_at = 0i64;
if let Ok(Some(bytes)) = db.get(&key) {
if let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&bytes)) {
title = v
.get("title")
.and_then(|x| x.as_str())
.unwrap_or("")
.to_string();
body = v
.get("body")
.and_then(|x| x.as_str())
.unwrap_or("")
.to_string();
notebook_id = v.get("notebookId").and_then(|x| x.as_i64()).unwrap_or(0);
updated_at = v.get("updatedAt").and_then(|x| x.as_i64()).unwrap_or(0);
}
}
let snippet = if body.len() > 180 {
format!("{}…", &body[..180])
} else {
body
};
Value::Object(vec![
("id".into(), Value::Int(id)),
("title".into(), Value::Str(title)),
("snippet".into(), Value::Str(snippet)),
("score".into(), Value::Float(score as f64)),
("notebookId".into(), Value::Int(notebook_id)),
("updatedAt".into(), Value::Int(updated_at)),
])
}