@sigit / siti
Code Commits Issues Pulls Cloud Agent CI/CD
release/windows-microsoft-store
development fix/asc-key-path main release/windows-microsoft-store
rs 494 lines 18.4 KB
Raw
1 //! Chat feature module: fully on-device LLM inference for Siti AI.
2 //!
3 //! Each Tauri command lives in its own file (`command_*.rs`). Shared types,
4 //! the cached engine instance, model/sampling configuration, and helpers are
5 //! kept here in `mod.rs` so every command file can `use super::*`.
6 //!
7 //! Model loading, inference, and history management are delegated to
8 //! [`onde::inference::ChatEngine`], the shared on-device inference engine
9 //! from the `onde` crate. Siti runs the engine **fully offline**: it loads
10 //! the platform-default GGUF model directly (no operator-assigned model, no
11 //! app credentials) and pulse telemetry is disabled during app setup.
12
13 // ── Command submodules (one Tauri command per file) ──────────────────────────
14
15 pub mod command_clear_history;
16 pub mod command_get_history;
17 pub mod command_get_status;
18 pub mod command_list_models;
19 pub mod command_load_model;
20 pub mod command_remove_model;
21 pub mod command_send_message;
22 pub mod command_set_model;
23 pub mod command_unload_model;
24
25 // ── Re-exports so lib.rs can pull in commands with a flat path ───────────────
26
27 pub use command_clear_history::chat_clear_history;
28 pub use command_get_history::chat_get_history;
29 pub use command_get_status::chat_get_status;
30 pub use command_list_models::chat_list_models;
31 pub use command_load_model::chat_load_model;
32 pub use command_remove_model::chat_remove_model;
33 pub use command_send_message::chat_send_message;
34 pub use command_set_model::chat_set_model;
35 pub use command_unload_model::chat_unload_model;
36
37 // ── Imports shared with submodules via `super::` ─────────────────────────────
38
39 use {
40 crate::constants::ChatStatus,
41 serde::{Deserialize, Serialize},
42 };
43
44 #[cfg(any(
45 target_os = "macos",
46 target_os = "ios",
47 target_os = "android",
48 target_os = "windows"
49 ))]
50 use {
51 crate::events::{EVENT_CHAT_REPLY, EVENT_CHAT_STATUS_CHANGED},
52 log::error,
53 once_cell::sync::Lazy,
54 onde::inference::{ChatEngine, GgufModelConfig, InferenceError, SamplingConfig},
55 tauri::{AppHandle, Emitter},
56 };
57
58 // `IsqModelConfig` and the ISQ load path only exist on macOS (Metal). Gemma is
59 // offered exclusively there — see `ResolvedModel` below for the reasoning.
60 #[cfg(target_os = "macos")]
61 use onde::inference::IsqModelConfig;
62
63 // ── Response / payload types ─────────────────────────────────────────────────
64
65 /// Payload emitted with the `chat_status_changed` event.
66 #[derive(Debug, Clone, Serialize, Deserialize)]
67 pub struct ChatStatusPayload {
68 pub status: ChatStatus,
69 pub model_name: Option<String>,
70 pub error: Option<String>,
71 }
72
73 /// A single message in the conversation, returned to the frontend.
74 #[derive(Debug, Clone, Serialize, Deserialize)]
75 pub struct ChatMessagePayload {
76 pub role: String,
77 pub content: String,
78 }
79
80 /// Event payload emitted as `chat_reply` once inference completes.
81 ///
82 /// Using an event instead of a blocking command return prevents the WebView
83 /// from garbage-collecting the JS invoke-callback before the slow on-device
84 /// inference finishes.
85 #[derive(Debug, Clone, Serialize, Deserialize)]
86 pub struct ChatReplyPayload {
87 /// The assistant's reply text, or `None` on error.
88 pub reply: Option<String>,
89 /// Human-readable inference duration, or `None` on error.
90 pub duration: Option<String>,
91 /// Error message if inference failed, or `None` on success.
92 pub error: Option<String>,
93 }
94
95 /// Response returned by `chat_get_status`.
96 #[derive(Debug, Clone, Serialize, Deserialize)]
97 pub struct ChatStatusResponse {
98 pub status: ChatStatus,
99 pub model_name: Option<String>,
100 pub approx_memory: Option<String>,
101 pub history_length: usize,
102 }
103
104 /// A selectable on-device model, returned by `chat_list_models` and rendered
105 /// in the Settings model dropdown.
106 #[derive(Debug, Clone, Serialize, Deserialize)]
107 pub struct ModelInfo {
108 /// HuggingFace repo id; the value passed back to `chat_set_model`.
109 pub id: String,
110 /// Human-friendly display name, e.g. `"Qwen 2.5 1.5B (GGUF)"`.
111 pub name: String,
112 /// Publisher / organisation, e.g. `"Qwen / Alibaba"`.
113 pub org: String,
114 /// Short description of the model's purpose and footprint.
115 pub description: String,
116 /// Approximate memory footprint, e.g. `"~941 MB (GGUF Q4_K_M)"`.
117 pub approx_memory: String,
118 /// Approximate on-disk download size in bytes (for display).
119 pub size_bytes: Option<u64>,
120 /// Whether the model's weights are already downloaded to the local cache.
121 pub is_downloaded: bool,
122 /// Whether this model is the one currently selected for Siti.
123 pub is_selected: bool,
124 }
125
126 // ── Shared ChatEngine instance ───────────────────────────────────────────────
127 //
128 // `ChatEngine` from `onde` handles all model lifecycle, history, and inference.
129 // It is `Send + Sync` and manages its own internal mutex.
130
131 #[cfg(any(
132 target_os = "macos",
133 target_os = "ios",
134 target_os = "android",
135 target_os = "windows"
136 ))]
137 pub(crate) static ENGINE: Lazy<ChatEngine> = Lazy::new(ChatEngine::new);
138
139 // ── Model selection & configuration ──────────────────────────────────────────
140 //
141 // The user can switch models from Settings (`chat_set_model`). The chosen
142 // HuggingFace repo id is held in `SELECTED_MODEL` and drives both the model
143 // config and the sampling config. When nothing has been chosen yet, Siti
144 // falls back to a general-purpose chat model appropriate for the platform:
145 // - iOS / Android → Qwen 2.5 1.5B (~941 MB, fits mobile memory budgets)
146 // - macOS/Windows → Qwen 2.5 3B (~1.93 GB, more headroom on desktop)
147 //
148 // `config_for_model_id` handles per-platform details (e.g. `tok_model_id` on
149 // mobile) via the onde `GgufModelConfig` constructors.
150
151 /// The model id currently selected for Siti. Defaults to the platform-
152 /// appropriate general chat model; updated by `chat_set_model`.
153 #[cfg(any(
154 target_os = "macos",
155 target_os = "ios",
156 target_os = "android",
157 target_os = "windows"
158 ))]
159 pub(crate) static SELECTED_MODEL: Lazy<std::sync::Mutex<String>> =
160 Lazy::new(|| std::sync::Mutex::new(siti_default_config().model_id));
161
162 /// Siti's default general-purpose chat model for the current platform.
163 #[cfg(any(
164 target_os = "macos",
165 target_os = "ios",
166 target_os = "android",
167 target_os = "windows"
168 ))]
169 pub(crate) fn siti_default_config() -> GgufModelConfig {
170 if cfg!(any(target_os = "ios", target_os = "android")) {
171 GgufModelConfig::qwen25_1_5b()
172 } else {
173 GgufModelConfig::qwen25_3b()
174 }
175 }
176
177 /// Map a HuggingFace repo id to a loadable [`GgufModelConfig`], or `None` if
178 /// the id is not a model the onde engine knows how to load.
179 #[cfg(any(
180 target_os = "macos",
181 target_os = "ios",
182 target_os = "android",
183 target_os = "windows"
184 ))]
185 pub(crate) fn config_for_model_id(id: &str) -> Option<GgufModelConfig> {
186 use crate::inference::models as m;
187 let cfg = if id == m::BARTOWSKI_QWEN25_0_5B_INSTRUCT_GGUF {
188 GgufModelConfig::qwen25_0_5b()
189 } else if id == m::BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF {
190 GgufModelConfig::qwen25_1_5b()
191 } else if id == m::BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF {
192 GgufModelConfig::qwen25_3b()
193 } else if id == m::BARTOWSKI_QWEN25_CODER_1_5B_INSTRUCT_GGUF {
194 GgufModelConfig::qwen25_coder_1_5b()
195 } else if id == m::BARTOWSKI_QWEN25_CODER_3B_INSTRUCT_GGUF {
196 GgufModelConfig::qwen25_coder_3b()
197 } else if id == m::BARTOWSKI_QWEN25_CODER_7B_INSTRUCT_GGUF {
198 GgufModelConfig::qwen25_coder_7b()
199 } else if id == m::BARTOWSKI_QWEN3_0_6B_GGUF {
200 GgufModelConfig::qwen3_0_6b()
201 } else if id == m::BARTOWSKI_QWEN3_1_7B_GGUF {
202 GgufModelConfig::qwen3_1_7b()
203 } else if id == m::BARTOWSKI_QWEN3_4B_GGUF {
204 GgufModelConfig::qwen3_4b()
205 } else if id == m::BARTOWSKI_QWEN3_8B_GGUF {
206 GgufModelConfig::qwen3_8b()
207 } else if id == m::BARTOWSKI_QWEN3_14B_GGUF {
208 GgufModelConfig::qwen3_14b()
209 } else if id == m::BARTOWSKI_QWEN3_32B_GGUF {
210 GgufModelConfig::qwen3_32b()
211 } else if id == m::BARTOWSKI_QWEN3_4B_INSTRUCT_2507_GGUF {
212 GgufModelConfig::qwen3_4b_instruct_2507()
213 } else if id == m::BARTOWSKI_QWEN3_4B_THINKING_2507_GGUF {
214 GgufModelConfig::qwen3_4b_thinking_2507()
215 } else if id == m::BARTOWSKI_QWEN3_30B_A3B_INSTRUCT_2507_GGUF {
216 GgufModelConfig::qwen3_30b_a3b_instruct_2507()
217 } else if id == m::THEBLOKE_DEEPSEEK_CODER_6_7B_INSTRUCT_GGUF {
218 GgufModelConfig::deepseek_coder_6_7b()
219 } else {
220 return None;
221 };
222 Some(cfg)
223 }
224
225 // ── Gemma (ISQ safetensors path, macOS only) ─────────────────────────────────
226 //
227 // Gemma cannot be loaded through the GGUF path Siti uses for every other model:
228 // mistral.rs's GGUF loader (`GGUFArchitecture`) only recognises Llama / Qwen2 /
229 // Qwen3 / Mistral3 / etc. — it has **no Gemma architecture**, so a Gemma GGUF
230 // file fails to load with "Unknown GGUF architecture `gemma2`". mistral.rs does
231 // support Gemma, but only via its plain safetensors loader, which onde exposes
232 // through the ISQ (in-situ quantise) path (`load_isq_model`).
233 //
234 // That path is Metal-only, so Gemma is offered on **macOS only** — never on
235 // iOS/Android, whose builds have no `IsqModelConfig` load path and far tighter
236 // memory budgets than an ISQ load (full bf16 weights downloaded, then quantised
237 // in memory) allows.
238
239 /// HuggingFace repo id for the Gemma 2 2B Instruct model, loaded via ISQ.
240 ///
241 /// The `unsloth` re-upload is used instead of `google/gemma-2-2b-it` because the
242 /// official repo is gated (requires accepting Google's licence + an HF token),
243 /// which Siti's credential-free offline download cannot satisfy. This mirror
244 /// ships the same `Gemma2ForCausalLM` bf16 weights, ungated.
245 #[cfg(target_os = "macos")]
246 pub(crate) const GEMMA2_2B_IT_ISQ_ID: &str = "unsloth/gemma-2-2b-it";
247
248 /// ISQ config for Gemma 2 2B Instruct (4-bit, Metal). macOS only.
249 #[cfg(target_os = "macos")]
250 pub(crate) fn gemma2_2b_isq_config() -> IsqModelConfig {
251 IsqModelConfig {
252 model_id: GEMMA2_2B_IT_ISQ_ID.to_string(),
253 isq_bits: 4,
254 display_name: "Gemma 2 2B (ISQ 4-bit)".to_string(),
255 approx_memory: "~1.6 GB (ISQ Q4K, Metal)".to_string(),
256 }
257 }
258
259 /// A model selection resolved to the concrete engine config and load path it
260 /// needs. Most models are GGUF; Gemma is the sole ISQ model (macOS only).
261 #[cfg(any(
262 target_os = "macos",
263 target_os = "ios",
264 target_os = "android",
265 target_os = "windows"
266 ))]
267 pub(crate) enum ResolvedModel {
268 Gguf(GgufModelConfig),
269 #[cfg(target_os = "macos")]
270 Isq(IsqModelConfig),
271 }
272
273 #[cfg(any(
274 target_os = "macos",
275 target_os = "ios",
276 target_os = "android",
277 target_os = "windows"
278 ))]
279 impl ResolvedModel {
280 /// Human-friendly display name for the resolved model.
281 pub(crate) fn display_name(&self) -> String {
282 match self {
283 ResolvedModel::Gguf(c) => c.display_name.clone(),
284 #[cfg(target_os = "macos")]
285 ResolvedModel::Isq(c) => c.display_name.clone(),
286 }
287 }
288
289 /// Load the resolved model into the shared [`ENGINE`], dispatching to the
290 /// GGUF or ISQ load path as appropriate.
291 pub(crate) async fn load(
292 self,
293 system_prompt: Option<String>,
294 sampling: Option<SamplingConfig>,
295 ) -> Result<std::time::Duration, InferenceError> {
296 match self {
297 ResolvedModel::Gguf(c) => ENGINE.load_gguf_model(c, system_prompt, sampling).await,
298 #[cfg(target_os = "macos")]
299 ResolvedModel::Isq(c) => ENGINE.load_isq_model(c, system_prompt, sampling).await,
300 }
301 }
302 }
303
304 /// Resolve a HuggingFace repo id to a loadable [`ResolvedModel`], or `None` if
305 /// the id is not a model Siti can load on this platform.
306 #[cfg(any(
307 target_os = "macos",
308 target_os = "ios",
309 target_os = "android",
310 target_os = "windows"
311 ))]
312 pub(crate) fn resolve_model_id(id: &str) -> Option<ResolvedModel> {
313 #[cfg(target_os = "macos")]
314 if id == GEMMA2_2B_IT_ISQ_ID {
315 return Some(ResolvedModel::Isq(gemma2_2b_isq_config()));
316 }
317 config_for_model_id(id).map(ResolvedModel::Gguf)
318 }
319
320 /// Return the resolved model (GGUF or ISQ) for the current selection, falling
321 /// back to the platform-default GGUF model if the selection is unknown.
322 #[cfg(any(
323 target_os = "macos",
324 target_os = "ios",
325 target_os = "android",
326 target_os = "windows"
327 ))]
328 pub(crate) fn resolved_model_config() -> ResolvedModel {
329 let id = SELECTED_MODEL.lock().map(|g| g.clone()).unwrap_or_default();
330 resolve_model_id(&id).unwrap_or_else(|| ResolvedModel::Gguf(siti_default_config()))
331 }
332
333 // ── Download-status detection ────────────────────────────────────────────────
334
335 /// Fraction of `expected_size_bytes` that must be present on disk for a model
336 /// to count as fully downloaded (mirrors onde's own completeness threshold).
337 #[cfg(any(
338 target_os = "macos",
339 target_os = "ios",
340 target_os = "android",
341 target_os = "windows"
342 ))]
343 const DOWNLOAD_COMPLETE_THRESHOLD: f64 = 0.99;
344
345 /// Sum the byte size of every real file directly and recursively under `path`.
346 ///
347 /// Only used against the HF cache's `blobs/` directory, which holds the actual
348 /// downloaded files (not the `snapshots/` symlink/hard-link views), so nothing
349 /// is double-counted.
350 #[cfg(any(
351 target_os = "macos",
352 target_os = "ios",
353 target_os = "android",
354 target_os = "windows"
355 ))]
356 fn blobs_dir_size(path: &std::path::Path) -> u64 {
357 let mut total = 0;
358 if let Ok(entries) = std::fs::read_dir(path) {
359 for entry in entries.flatten() {
360 let p = entry.path();
361 match entry.metadata() {
362 Ok(md) if md.is_dir() => total += blobs_dir_size(&p),
363 Ok(md) => total += md.len(),
364 Err(_) => {}
365 }
366 }
367 }
368 total
369 }
370
371 /// Whether `id`'s weights are already downloaded to the local HF cache.
372 ///
373 /// A model counts as downloaded when its cache `blobs/` directory holds at
374 /// least [`DOWNLOAD_COMPLETE_THRESHOLD`] of the model's expected size, which
375 /// excludes partial/interrupted downloads.
376 #[cfg(any(
377 target_os = "macos",
378 target_os = "ios",
379 target_os = "android",
380 target_os = "windows"
381 ))]
382 pub(crate) fn is_model_downloaded(id: &str, expected_size_bytes: u64) -> bool {
383 if expected_size_bytes == 0 {
384 return false;
385 }
386 match onde::hf_cache::model_cache_path(id) {
387 Some(root) => {
388 let blobs = root.join("blobs");
389 blobs_dir_size(&blobs) as f64
390 >= expected_size_bytes as f64 * DOWNLOAD_COMPLETE_THRESHOLD
391 }
392 None => false,
393 }
394 }
395
396 /// Return the sampling config for the currently selected model.
397 ///
398 /// iOS and Android use conservative settings to bound KV cache memory and
399 /// keep worst-case latency reasonable on slower SoCs. Qwen 3 models get a
400 /// larger `max_tokens` budget because their `<think>…</think>` block can
401 /// consume hundreds of tokens before the visible reply begins, and too small
402 /// a budget yields empty replies.
403 #[cfg(any(
404 target_os = "macos",
405 target_os = "ios",
406 target_os = "android",
407 target_os = "windows"
408 ))]
409 pub(crate) fn sampling_config() -> SamplingConfig {
410 let id = SELECTED_MODEL.lock().map(|g| g.clone()).unwrap_or_default();
411 let is_qwen3 = id.contains("Qwen3") || id.contains("Qwen_Qwen3");
412 let mobile = cfg!(any(target_os = "ios", target_os = "android"));
413
414 let max_tokens = match (mobile, is_qwen3) {
415 (true, true) => 1024,
416 (true, false) => 256,
417 (false, true) => 4096,
418 (false, false) => 512,
419 };
420
421 if mobile {
422 SamplingConfig {
423 temperature: Some(0.7),
424 top_p: Some(0.95),
425 max_tokens: Some(max_tokens),
426 ..SamplingConfig::mobile()
427 }
428 } else {
429 SamplingConfig {
430 temperature: Some(0.7),
431 top_p: Some(0.95),
432 max_tokens: Some(max_tokens),
433 ..SamplingConfig::default()
434 }
435 }
436 }
437
438 /// Default system prompt for Siti, the private on-device personal assistant.
439 #[cfg(any(
440 target_os = "macos",
441 target_os = "ios",
442 target_os = "android",
443 target_os = "windows"
444 ))]
445 pub(crate) const CHAT_SYSTEM_PROMPT: &str = r#"You are Siti, a friendly and capable personal assistant that runs entirely on the user's device. Nothing the user says ever leaves the device.
446
447 You help with everyday things: answering questions, brainstorming, drafting and editing text, summarising, explaining ideas in plain terms, and planning.
448
449 Keep replies clear and concise, and prefer plain language. If a request is ambiguous, ask one short question instead of guessing. Use a short list only when it genuinely helps, and go easy on emoji.
450
451 You run locally, so you may not know about very recent events. When you are not sure, say so instead of inventing an answer."#;
452
453 // ── Helpers ──────────────────────────────────────────────────────────────────
454
455 /// Format a `Duration` as `Xm Ys` or just `Ys` when under a minute.
456 #[cfg(any(
457 target_os = "macos",
458 target_os = "ios",
459 target_os = "android",
460 target_os = "windows"
461 ))]
462 pub(crate) fn fmt_duration(d: std::time::Duration) -> String {
463 let total_secs = d.as_secs_f64();
464 let mins = (total_secs / 60.0).floor() as u64;
465 let secs = total_secs - (mins as f64 * 60.0);
466 if mins > 0 {
467 format!("{}m {:.1}s", mins, secs)
468 } else {
469 format!("{:.1}s", secs)
470 }
471 }
472
473 /// Emit a chat status event to the frontend.
474 #[cfg(any(
475 target_os = "macos",
476 target_os = "ios",
477 target_os = "android",
478 target_os = "windows"
479 ))]
480 pub(crate) fn emit_chat_status(
481 app: &AppHandle,
482 status: ChatStatus,
483 model_name: Option<&str>,
484 error_msg: Option<&str>,
485 ) {
486 let payload = ChatStatusPayload {
487 status,
488 model_name: model_name.map(|s| s.to_string()),
489 error: error_msg.map(|s| s.to_string()),
490 };
491 if let Err(e) = app.emit(EVENT_CHAT_STATUS_CHANGED, &payload) {
492 error!("Failed to emit chat status event: {:?}", e);
493 }
494 }