AkurAI-Notes / crates/notes-api/src/handlers/search.rs

Latest change 161c7c40bf25c3770f7422b84193c3dd6fbf0aac - AKURAI-NOTES-7 emit privacy-safe organization activation and adoption facts by Ólafur Búi Ólafsson

use crate::auth::session;
use crate::embed::{embed_query_with_key, EmbedStore};
use crate::handlers::RouteHandler;
use crate::router_setup::{json_response, unauthorized, upper_bound, NotesState};
use akurai_http::{Request, Response};
use akurai_json::Value;
use akurai_storage::BTree;
use notes_core::search as note_search;
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 = decode_query_component(value),
                    "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, embed config, and telemetry in one locked scope.
        let (user, embed_url, embed_model, embed_api_key, telemetry) = {
            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 unauthorized();
                }
            };
            let user = match session::get_user(&mut db, &sid) {
                Some(u) => u,
                None => {
                    return unauthorized();
                }
            };
            (
                user,
                state_lock.embed_url.clone(),
                state_lock.embed_model.clone(),
                state_lock.embed_api_key.clone(),
                state_lock.telemetry.clone(),
            )
        };

        let now = NotesState::now();

        // Helper: emit active_use telemetry for non-empty results.
        let emit_search_results = |res: &[Value]| {
            if let Some(ref t) = telemetry {
                if !res.is_empty() {
                    let _ = t.emit_active_use(now, res.len() as i64);
                }
            }
        };

        // Try semantic search
        if let Some(url) = &embed_url {
            let model = embed_model.as_deref().unwrap_or("qwen3-embedding-0.6b");
            match embed_query_with_key(&q, url, model, embed_api_key.as_deref()) {
                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 owners = [user.sub.as_str(), user.email.as_str()];
                    let ranked = store.rank_notes(&qv, &owners, k);
                    let results: Vec<Value> = ranked
                        .into_iter()
                        .map(|(id, score)| note_to_result(&mut *store.db, id, score))
                        .collect();
                    if !results.is_empty() {
                        emit_search_results(&results);
                        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()),
                    )]));
                }
            };
            let owners = [user.sub.as_str(), user.email.as_str()];
            substring_search(&mut db, &q, &owners, k)
        };

        emit_search_results(&results);
        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, owners: &[&str], k: usize) -> Vec<Value> {
    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, i64, 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 !note_search::owner_matches(&v, owners) {
            continue;
        }
        let title = note_search::note_title(&v).to_string();
        let body = note_search::note_body(&v).to_string();
        if let Some(score) = note_search::lexical_score(q, &title, &body) {
            let updated_at = note_search::note_updated_at(&v);
            scored.push((id, score, updated_at, title, body));
        }
    }

    scored.sort_by(|a, b| {
        b.1.partial_cmp(&a.1)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| b.2.cmp(&a.2))
            .then_with(|| b.0.cmp(&a.0))
    });
    scored.truncate(k);

    scored
        .into_iter()
        .map(|(id, score, _updated_at, title, body)| {
            Value::Object(vec![
                ("id".into(), Value::Int(id)),
                ("title".into(), Value::Str(title)),
                (
                    "snippet".into(),
                    Value::Str(note_search::snippet(&body, 180)),
                ),
                ("score".into(), Value::Float(score as f64)),
            ])
        })
        .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 = note_search::snippet(&body, 180);
    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)),
    ])
}

fn decode_query_component(input: &str) -> String {
    let bytes = input.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;
    while i < bytes.len() {
        match bytes[i] {
            b'+' => {
                out.push(b' ');
                i += 1;
            }
            b'%' if i + 2 < bytes.len() => {
                if let (Some(hi), Some(lo)) = (hex_value(bytes[i + 1]), hex_value(bytes[i + 2])) {
                    out.push((hi << 4) | lo);
                    i += 3;
                } else {
                    out.push(bytes[i]);
                    i += 1;
                }
            }
            b => {
                out.push(b);
                i += 1;
            }
        }
    }
    String::from_utf8_lossy(&out).to_string()
}

fn hex_value(byte: u8) -> Option<u8> {
    match byte {
        b'0'..=b'9' => Some(byte - b'0'),
        b'a'..=b'f' => Some(byte - b'a' + 10),
        b'A'..=b'F' => Some(byte - b'A' + 10),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::decode_query_component;

    #[test]
    fn decodes_browser_query_components() {
        assert_eq!(
            decode_query_component("AkurAI+Platform+Skills%20CLI%20Tools"),
            "AkurAI Platform Skills CLI Tools"
        );
    }
}