falcon_server/
ai.rs

1//! Eagle's `/api/v2/aiSearch/*` routes. The work is done by whatever the app
2//! hands over ([`crate::Handle::set_ai_search`]), so this crate doesn't build
3//! the model; with nothing handed over (the `falcon` command line, or the
4//! app before it starts) they answer that AI Search isn't installed.
5//!
6//! PARITY.md knows these routes by name only, so the shapes are falcon's
7//! reading of the names: the four `is*` routes answer a boolean, and searches
8//! answer the v2 list shape (`data`, `total`, `offset`, `limit`), each item
9//! with a `score`. Searching runs without the library lock.
10
11use std::sync::{Arc, PoisonError};
12
13use base64::Engine;
14use serde_json::{Map, Value, json};
15
16use crate::{Reply, Shared, State, item_json, lock, number, paginate, s};
17
18/// How AI Search stands.
19#[derive(Clone, Debug, Default, PartialEq, Eq)]
20pub struct AiStatus {
21    /// The model is on this machine and turned on.
22    pub installed: bool,
23    /// Searches can be answered now.
24    pub ready: bool,
25    /// The model is loading.
26    pub starting: bool,
27    /// Items are being indexed.
28    pub syncing: bool,
29    /// Items in the library that have been looked at, and how many there are.
30    pub done: usize,
31    pub total: usize,
32    /// The model failed.
33    pub error: Option<String>,
34}
35
36/// What to search by.
37pub enum AiQuery {
38    Text(String),
39    /// The bytes of a picture file.
40    Image(Vec<u8>),
41    Item(String),
42}
43
44/// AI Search as the app provides it.
45pub trait AiSearch: Send + Sync {
46    fn status(&self) -> AiStatus;
47    /// Item ids and scores, best first.
48    fn search(&self, query: AiQuery, limit: usize) -> Result<Vec<(String, f32)>, String>;
49}
50
51const NOT_INSTALLED: &str = "AI Search is not installed. Install its model in Preferences > AI Search.";
52const NOT_READY: &str = "AI Search is not ready";
53
54/// Answers one `aiSearch` route.
55pub(crate) fn route(lib: &Shared, state: &State, path: &str, args: &Map<String, Value>) -> Reply {
56    let engine: Option<Arc<dyn AiSearch>> = state.ai.lock().unwrap_or_else(PoisonError::into_inner).clone();
57    let status = engine.as_ref().map(|e| e.status()).unwrap_or_default();
58    let name = path.rsplit('/').next().unwrap_or("");
59    match name {
60        "isInstalled" => Reply::Data(json!(status.installed)),
61        "isReady" => Reply::Data(json!(status.ready)),
62        "isStarting" => Reply::Data(json!(status.starting)),
63        "isSyncing" => Reply::Data(json!(status.syncing)),
64        "getSyncStatus" => Reply::Data(json!({
65            "isSyncing": status.syncing,
66            "done": status.done,
67            "total": status.total,
68            "pending": status.total.saturating_sub(status.done),
69        })),
70        "checkServiceHealth" => {
71            let word = match (&status.error, status.installed, status.starting, status.syncing) {
72                (Some(_), ..) => "error",
73                (None, false, ..) => "notInstalled",
74                (None, true, true, _) => "starting",
75                (None, true, false, true) => "syncing",
76                (None, true, false, false) => "ready",
77            };
78            Reply::Data(json!({"healthy": status.ready && status.error.is_none(), "status": word, "error": status.error}))
79        }
80        "searchByText" | "searchByBase64" | "searchByItemId" => {
81            let query = match name {
82                "searchByText" => s(args, "query").or_else(|| s(args, "text")).map(|t| AiQuery::Text(t.to_string())),
83                "searchByItemId" => s(args, "itemId").or_else(|| s(args, "id")).map(|i| AiQuery::Item(i.to_string())),
84                _ => match s(args, "base64").or_else(|| s(args, "image")).map(decode_base64) {
85                    Some(Ok(bytes)) => Some(AiQuery::Image(bytes)),
86                    Some(Err(message)) => return Reply::Error(message),
87                    None => None,
88                },
89            };
90            let Some(query) = query else { return Reply::Error(format!("{name}: missing what to search by")) };
91            let Some(engine) = engine.filter(|_| status.installed) else { return Reply::Error(NOT_INSTALLED.to_string()) };
92            if !status.ready {
93                return Reply::Error(NOT_READY.to_string());
94            }
95            let limit = (number(args, "limit").filter(|l| *l > 0).unwrap_or(50) as usize).min(1000);
96            match engine.search(query, limit) {
97                Ok(hits) => {
98                    let lib = lock(lib);
99                    let items = hits.into_iter().filter_map(|(id, score)| {
100                        let mut value = item_json(lib.item(&id)?);
101                        if let Value::Object(m) = &mut value {
102                            m.insert("score".into(), json!(score));
103                        }
104                        Some(value)
105                    });
106                    Reply::Data(paginate(items.collect(), args))
107                }
108                Err(message) => Reply::Error(message),
109            }
110        }
111        _ => Reply::NotFound,
112    }
113}
114
115/// The bytes of a picture sent as Base64, with or without a `data:` prefix.
116fn decode_base64(text: &str) -> Result<Vec<u8>, String> {
117    let body = if text.starts_with("data:") { text.split_once(',').map_or("", |(_, b)| b) } else { text };
118    base64::engine::general_purpose::STANDARD.decode(body.trim()).map_err(|e| format!("base64: {e}"))
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn base64_comes_with_or_without_a_data_prefix() {
127        assert_eq!(decode_base64("aGk=").unwrap(), b"hi");
128        assert_eq!(decode_base64("data:image/png;base64,aGk=").unwrap(), b"hi");
129        assert!(decode_base64("not base64!").is_err());
130    }
131}
falcon_server/ai.rs

falcon_server/
ai.rs
