@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
feature/diffusion-vision
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 2,445 lines 91.2 KB
Raw
1 // Copyright 2026 Splitfire AB (Onde Inference). All rights reserved.
2 // SPDX-License-Identifier: MIT OR Apache-2.0
3
4 //! On-device LLM chat inference engine powered by [mistral.rs](https://github.com/EricLBuehler/mistral.rs).
5 //!
6 //! `ChatEngine` provides a high-level, framework-agnostic API for:
7 //!
8 //! - Loading GGUF-quantized models (the primary format for mobile/desktop)
9 //! - Multi-turn chat with conversation history management
10 //! - Both blocking (non-streaming) and streaming inference
11 //! - Model lifecycle management (load / unload / status)
12 //!
13 //! # Architecture
14 //!
15 //! ```text
16 //! ┌──────────────────────────────────────────┐
17 //! │ App / UniFFI binding / test │
18 //! └──────────────┬───────────────────────────┘
19 //! │ (framework-agnostic API)
20 //! ▼
21 //! ┌──────────────────────────────────────────┐
22 //! │ ChatEngine │
23 //! │ ┌────────────────────────────────────┐ │
24 //! │ │ Mutex<Option<LoadedModel>> │ │
25 //! │ │ · Arc<mistralrs::Model> │ │
26 //! │ │ · config (GgufModelConfig) │ │
27 //! │ │ · history (Vec<ChatMessage>) │ │
28 //! │ │ · sampling (SamplingConfig) │ │
29 //! │ └────────────────────────────────────┘ │
30 //! └──────────────┬───────────────────────────┘
31 //! │ (delegates to)
32 //! ▼
33 //! ┌──────────────────────────────────────────┐
34 //! │ mistralrs::Model │
35 //! │ (wraps Arc<MistralRs> engine thread) │
36 //! └──────────────────────────────────────────┘
37 //! ```
38 //!
39 //! # Example
40 //!
41 //! ```rust,ignore
42 //! use onde::inference::engine::ChatEngine;
43 //! use onde::inference::types::*;
44 //!
45 //! let engine = ChatEngine::new();
46 //!
47 //! let config = GgufModelConfig {
48 //! model_id: "bartowski/Qwen2.5-1.5B-Instruct-GGUF".into(),
49 //! files: vec!["Qwen2.5-1.5B-Instruct-Q4_K_M.gguf".into()],
50 //! tok_model_id: None,
51 //! display_name: "Qwen 2.5 1.5B".into(),
52 //! approx_memory: "~941 MB".into(),
53 //! chat_template: None,
54 //! };
55 //!
56 //! engine.load_gguf_model(config, None).await?;
57 //! engine.set_system_prompt("You are a helpful assistant.");
58 //!
59 //! let result = engine.send_message("Hello!").await?;
60 //! println!("{}", result.text);
61 //! ```
62
63 #[cfg(any(
64 target_os = "macos",
65 target_os = "ios",
66 target_os = "tvos",
67 target_os = "visionos",
68 target_os = "watchos",
69 target_os = "windows",
70 target_os = "linux",
71 target_os = "android"
72 ))]
73 use std::sync::Arc;
74
75 #[cfg(any(
76 target_os = "macos",
77 target_os = "ios",
78 target_os = "tvos",
79 target_os = "visionos",
80 target_os = "watchos",
81 target_os = "windows",
82 target_os = "linux",
83 target_os = "android"
84 ))]
85 use tokio::sync::Mutex;
86
87 use super::types::*;
88
89 // ── Platform-gated mistralrs imports ─────────────────────────────────────────
90
91 #[cfg(any(
92 target_os = "macos",
93 target_os = "ios",
94 target_os = "tvos",
95 target_os = "visionos",
96 target_os = "watchos",
97 target_os = "windows",
98 target_os = "linux",
99 target_os = "android"
100 ))]
101 use mistralrs::{
102 CalledFunction, Function, GgufModelBuilder, Model, RequestBuilder, TextMessageRole,
103 Tool as MistralTool, ToolCallResponse, ToolCallType, ToolChoice, ToolType,
104 };
105
106 #[cfg(any(
107 target_os = "macos",
108 target_os = "ios",
109 target_os = "tvos",
110 target_os = "visionos",
111 target_os = "watchos",
112 target_os = "windows",
113 target_os = "linux",
114 target_os = "android"
115 ))]
116 use std::collections::HashMap;
117
118 // ISQ types are only used by load_isq_model, which is macOS-only.
119 #[cfg(target_os = "macos")]
120 use mistralrs::{IsqBits, TextModelBuilder};
121
122 // ── Internals ────────────────────────────────────────────────────────────────
123
124 /// How a model was loaded — carries the data needed for status reporting.
125 #[cfg(any(
126 target_os = "macos",
127 target_os = "ios",
128 target_os = "tvos",
129 target_os = "visionos",
130 target_os = "watchos",
131 target_os = "windows",
132 target_os = "linux",
133 target_os = "android"
134 ))]
135 enum LoadedModelConfig {
136 Gguf(GgufModelConfig),
137 #[cfg(target_os = "macos")]
138 Isq(IsqModelConfig),
139 }
140
141 #[cfg(any(
142 target_os = "macos",
143 target_os = "ios",
144 target_os = "tvos",
145 target_os = "visionos",
146 target_os = "watchos",
147 target_os = "windows",
148 target_os = "linux",
149 target_os = "android"
150 ))]
151 impl LoadedModelConfig {
152 fn display_name(&self) -> &str {
153 match self {
154 LoadedModelConfig::Gguf(c) => &c.display_name,
155 #[cfg(target_os = "macos")]
156 LoadedModelConfig::Isq(c) => &c.display_name,
157 }
158 }
159
160 fn approx_memory(&self) -> &str {
161 match self {
162 LoadedModelConfig::Gguf(c) => &c.approx_memory,
163 #[cfg(target_os = "macos")]
164 LoadedModelConfig::Isq(c) => &c.approx_memory,
165 }
166 }
167 }
168
169 /// State of a loaded model held inside the engine's mutex.
170 #[cfg(any(
171 target_os = "macos",
172 target_os = "ios",
173 target_os = "tvos",
174 target_os = "visionos",
175 target_os = "watchos",
176 target_os = "windows",
177 target_os = "linux",
178 target_os = "android"
179 ))]
180 struct LoadedModel {
181 /// The mistral.rs model handle. `Model` wraps `Arc<MistralRs>` so
182 /// cloning is a cheap pointer copy — safe to snapshot before inference
183 /// while releasing the mutex.
184 model: Arc<Model>,
185 /// Configuration used to load this model (kept for status reporting).
186 config: LoadedModelConfig,
187 /// Conversation history (system prompt is stored separately).
188 history: Vec<HistoryEntry>,
189 /// System prompt prepended to every request.
190 system_prompt: Option<String>,
191 /// Sampling parameters for generation.
192 sampling: SamplingConfig,
193 }
194
195 /// An entry in the conversation history (internal only).
196 /// Supports tool-related messages alongside regular text.
197 #[cfg(any(
198 target_os = "macos",
199 target_os = "ios",
200 target_os = "tvos",
201 target_os = "visionos",
202 target_os = "watchos",
203 target_os = "windows",
204 target_os = "linux",
205 target_os = "android"
206 ))]
207 #[derive(Debug, Clone)]
208 enum HistoryEntry {
209 /// A regular text message (user, assistant, system).
210 Text(ChatMessage),
211 /// An assistant response that includes tool calls.
212 AssistantToolCall {
213 content: String,
214 tool_calls: Vec<ToolCallRequest>,
215 },
216 /// A tool execution result.
217 ToolResult {
218 tool_call_id: String,
219 content: String,
220 },
221 }
222
223 // ═════════════════════════════════════════════════════════════════════════════
224 // ChatEngine — the public API
225 // ═════════════════════════════════════════════════════════════════════════════
226
227 /// A reusable on-device LLM chat inference engine.
228 ///
229 /// Thread-safe (`Send + Sync`) — safe to store in a `once_cell::sync::Lazy`,
230 /// `tokio::sync::OnceCell`, or shared application state.
231 ///
232 /// All mutating operations acquire an internal `tokio::sync::Mutex`. The
233 /// mutex is released *before* the actual (potentially slow) model inference
234 /// runs, so other tasks can still query status or history while generation
235 /// is in progress.
236 #[cfg(any(
237 target_os = "macos",
238 target_os = "ios",
239 target_os = "tvos",
240 target_os = "visionos",
241 target_os = "watchos",
242 target_os = "windows",
243 target_os = "linux",
244 target_os = "android"
245 ))]
246 pub struct ChatEngine {
247 inner: Mutex<Option<LoadedModel>>,
248 pulse: std::sync::OnceLock<Option<crate::pulse::PulseClient>>,
249 /// The Onde app UUID this engine reports telemetry for. Threaded into every
250 /// pulse event so the document store associates rows with the right app.
251 /// `None` for open-ended SDK consumers (UniFFI `OndeChatEngine`).
252 onde_app_id: Option<String>,
253 }
254
255 #[cfg(any(
256 target_os = "macos",
257 target_os = "ios",
258 target_os = "tvos",
259 target_os = "visionos",
260 target_os = "watchos",
261 target_os = "windows",
262 target_os = "linux",
263 target_os = "android"
264 ))]
265 impl ChatEngine {
266 // ── Construction ─────────────────────────────────────────────────────
267
268 /// Create a new engine with no model loaded and no Onde app association.
269 /// Pulse telemetry is still enabled, but events carry no `onde_app_id`.
270 pub fn new() -> Self {
271 Self::with_app_id(None)
272 }
273
274 /// Create a new engine that reports telemetry for a specific Onde app.
275 /// The app UUID is threaded into every pulse event so the document store
276 /// associates rows with the right app. Consumer apps that know their ID
277 /// call `ChatEngine::with_app_id(Some(smbcloud::onde_app_id()))`; the
278 /// UniFFI `OndeChatEngine` keeps calling `new()` (→ `None`).
279 pub fn with_app_id(onde_app_id: Option<String>) -> Self {
280 crate::install_panic_hook_once();
281 Self {
282 inner: Mutex::new(None),
283 pulse: std::sync::OnceLock::new(),
284 onde_app_id,
285 }
286 }
287
288 /// Lazily initialise the pulse telemetry client.
289 ///
290 /// The underlying GresIQ SDK creates a `reqwest::Client`, which needs
291 /// a Tokio runtime. `ChatEngine::new()` is called from a sync UniFFI
292 /// constructor where no runtime exists yet, so we defer creation to the
293 /// first async call site (e.g. `load_gguf_model` or `send_message`).
294 fn pulse(&self) -> Option<&crate::pulse::PulseClient> {
295 self.pulse
296 .get_or_init(|| {
297 let environment = Self::pulse_environment();
298 let edge_id =
299 std::env::var("ONDE_EDGE_ID").unwrap_or_else(|_| "onde-unknown".to_string());
300 // Prefer the constructor-supplied app id; fall back to the
301 // ONDE_APP_ID env var so UniFFI consumers can opt in without a
302 // new constructor signature.
303 let onde_app_id = self.onde_app_id.clone().or_else(|| {
304 std::env::var("ONDE_APP_ID")
305 .ok()
306 .filter(|id| !id.is_empty())
307 });
308 let client = crate::pulse::PulseClient::new(environment, edge_id, onde_app_id);
309
310 match &client {
311 Some(_) => log::info!(
312 "ChatEngine: pulse telemetry enabled (environment={environment})"
313 ),
314 None if crate::pulse::PulseClient::disabled_by_env() => {
315 log::info!("ChatEngine: pulse telemetry disabled by ONDE_DISABLE_PULSE")
316 }
317 None => log::info!(
318 "ChatEngine: pulse telemetry disabled \
319 (GRESIQ_API_KEY / GRESIQ_API_SECRET not embedded at SDK build time)"
320 ),
321 }
322
323 client
324 })
325 .as_ref()
326 }
327
328 /// Resolve the pulse environment from the `GRESIQ_ENVIRONMENT` env var.
329 /// Defaults to `Production` when the var is absent or unrecognised.
330 fn pulse_environment() -> smbcloud_gresiq_sdk::Environment {
331 match std::env::var("GRESIQ_ENVIRONMENT").as_deref() {
332 Ok("dev") => smbcloud_gresiq_sdk::Environment::Dev,
333 _ => smbcloud_gresiq_sdk::Environment::Production,
334 }
335 }
336
337 // ── Model lifecycle ──────────────────────────────────────────────────
338
339 /// Load a GGUF model into the engine.
340 ///
341 /// If a model is already loaded it will be unloaded first (the previous
342 /// `Model` is dropped, which terminates its engine thread).
343 ///
344 /// # Arguments
345 ///
346 /// * `config` — Which model to load (repo ID, filename, etc.).
347 /// * `system_prompt` — Optional system prompt to prepend to every request.
348 /// * `sampling` — Sampling parameters; pass `None` for platform-aware
349 /// defaults (mobile gets [`SamplingConfig::mobile`],
350 /// desktop gets [`SamplingConfig::default`]).
351 ///
352 /// # Errors
353 ///
354 /// Returns [`InferenceError::ModelBuild`] if the model fails to download
355 /// or load.
356 pub async fn load_gguf_model(
357 &self,
358 config: GgufModelConfig,
359 system_prompt: Option<String>,
360 sampling: Option<SamplingConfig>,
361 ) -> Result<std::time::Duration, InferenceError> {
362 use log::info;
363 use std::time::Instant;
364
365 info!(
366 "ChatEngine: loading GGUF model {} (files: {:?})",
367 config.model_id, config.files
368 );
369
370 // ── Sandboxed platforms: seed GLOBAL_HF_CACHE ────────────────────
371 //
372 // On sandboxed platforms the default `~/.cache/huggingface/hub` path
373 // is either inaccessible or non-existent:
374 // - Android: `dirs::home_dir()` returns `None` → `Cache::default()` panics.
375 // - iOS/tvOS: `~/.cache` is outside the app container → os error 1.
376 //
377 // The `get_paths_gguf!` macro in mistralrs falls back to
378 // `Cache::default()` when `GLOBAL_HF_CACHE` (a OnceLock) is empty.
379 //
380 // `HF_HOME` must be set by the host app (via `configure_cache_dir`
381 // or `download_model(app_data_dir:)`) before any model load.
382 // `get_or_init` is a no-op if already seeded — safe to call repeatedly.
383 #[cfg(any(
384 target_os = "android",
385 target_os = "ios",
386 target_os = "tvos",
387 target_os = "visionos",
388 target_os = "watchos"
389 ))]
390 {
391 let hf_home = std::env::var("HF_HOME")
392 .map(std::path::PathBuf::from)
393 .map_err(|_| InferenceError::ModelBuild {
394 reason: "HF_HOME is not set — cannot initialise HF cache. \
395 On iOS/tvOS/Android, call configure_cache_dir() or \
396 download_model(app_data_dir:) before load_gguf_model()."
397 .to_string(),
398 })?;
399 let hf_hub_cache = hf_home.join("hub");
400 if let Err(e) = std::fs::create_dir_all(&hf_hub_cache) {
401 log::warn!(
402 "ChatEngine: could not create HF hub cache dir {}: {}",
403 hf_hub_cache.display(),
404 e
405 );
406 }
407 mistralrs_core::GLOBAL_HF_CACHE
408 .get_or_init(|| hf_hub::Cache::new(hf_hub_cache.clone()));
409 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
410 log::debug!(
411 "ChatEngine: GLOBAL_HF_CACHE seeded at {}",
412 hf_hub_cache.display()
413 );
414 }
415
416 // Clean up stale HF cache artefacts before loading.
417 crate::hf_cache::clean_stale_lock_files(&config.model_id);
418 crate::hf_cache::repair_hf_cache_symlinks(&config.model_id);
419
420 // If the blob was downloaded but the process was killed before hf-hub
421 // created the snapshot entry, hf-hub will re-download. Create the
422 // missing snapshot hard links so the cache hit succeeds.
423 let file_refs: Vec<&str> = config.files.iter().map(|s| s.as_str()).collect();
424 crate::hf_cache::ensure_snapshot_entries(&config.model_id, &file_refs);
425
426 let start = Instant::now();
427
428 let mut builder = GgufModelBuilder::new(&config.model_id, config.files.clone())
429 .with_token_source(super::token::hf_token_source())
430 .with_logging();
431
432 // tvOS: Apple TV 4K has only 4 GB RAM. Reduce memory footprint by
433 // limiting concurrent sequences, disabling the prefix cache,
434 // disabling the KV cache, and capping context length at 512 tokens
435 // so the auto device mapper allocates far less memory.
436 #[cfg(target_os = "tvos")]
437 {
438 use mistralrs::DeviceMapSetting;
439 use mistralrs_core::AutoDeviceMapParams;
440
441 builder = builder
442 .with_max_num_seqs(1)
443 .with_prefix_cache_n(None)
444 .with_no_kv_cache()
445 .with_device_mapping(DeviceMapSetting::Auto(AutoDeviceMapParams::Text {
446 max_seq_len: 512,
447 max_batch_size: 1,
448 }));
449 }
450
451 // On Android the GGUF embedded tokenizer is not supported by the
452 // candle backend — an explicit tok_model_id is required.
453 if let Some(ref tok_id) = config.tok_model_id {
454 builder = builder.with_tok_model_id(tok_id);
455 }
456
457 // Some older GGUF files (e.g. TheBloke) do not embed a chat template.
458 if let Some(ref template) = config.chat_template {
459 builder = builder.with_chat_template(template);
460 }
461
462 let model = builder
463 .build()
464 .await
465 .map_err(|e| InferenceError::ModelBuild {
466 reason: format!("Failed to build {} model: {}", config.display_name, e),
467 })?;
468
469 let elapsed = start.elapsed();
470
471 let sampling = sampling.unwrap_or_else(|| {
472 if config.uses_extended_thinking() {
473 SamplingConfig::reasoning()
474 } else if cfg!(any(
475 target_os = "ios",
476 target_os = "tvos",
477 target_os = "visionos",
478 target_os = "watchos",
479 target_os = "android"
480 )) {
481 SamplingConfig::mobile()
482 } else {
483 SamplingConfig::default()
484 }
485 });
486
487 info!(
488 "ChatEngine: model {} loaded in {} (sampling: temp={:?}, max_tokens={:?})",
489 config.display_name,
490 format_duration(elapsed),
491 sampling.temperature,
492 sampling.max_tokens,
493 );
494
495 // Capture before config is moved into LoadedModel.
496 let pulse_model_id = config.model_id.clone();
497 let pulse_model_name = config.display_name.clone();
498
499 // Swap the new model in and take the old one out — all under the lock
500 // so the engine is never seen as None by a concurrent prompt. Then
501 // drop the old weights *outside* the lock so we don't block inference
502 // while the allocator frees potentially several GB of tensors.
503 let old_model = {
504 let mut guard = self.inner.lock().await;
505 guard.replace(LoadedModel {
506 model: Arc::new(model),
507 config: LoadedModelConfig::Gguf(config),
508 history: Vec::new(),
509 system_prompt,
510 sampling,
511 })
512 };
513 drop(old_model); // free old weights outside the lock
514
515 if let Some(pulse) = self.pulse() {
516 pulse
517 .record_model_loaded(pulse_model_id, pulse_model_name, elapsed.as_millis() as u64)
518 .await;
519 }
520
521 Ok(elapsed)
522 }
523
524 /// Fetch the model assigned to this Onde app from the SDK model-config
525 /// endpoint and load it.
526 ///
527 /// Authenticates using the app's own `onde_app_id` + `onde_app_secret`
528 /// (the SDK credentials shown in the ondeinference.com dashboard).
529 /// No end-user JWT is required — model assignment is an operator-level
530 /// configuration that is independent of which user is currently signed in.
531 ///
532 /// **Fallback behaviour:**
533 /// - HTTP 404 (no model assigned) → loads [`GgufModelConfig::platform_default()`].
534 /// - Missing `hf_repo_id` or `gguf_file` in response → loads platform default.
535 /// - Any HTTP or network error → returns [`InferenceError::ModelBuild`] so
536 /// the caller can apply its own fallback.
537 pub async fn load_assigned_model(
538 &self,
539 environment: smbcloud_gresiq_sdk::Environment,
540 onde_app_id: &str,
541 onde_app_secret: &str,
542 system_prompt: Option<String>,
543 sampling: Option<SamplingConfig>,
544 ) -> Result<std::time::Duration, InferenceError> {
545 use log::{info, warn};
546
547 #[derive(serde::Deserialize)]
548 struct ModelConfigResponse {
549 hf_repo_id: Option<String>,
550 gguf_file: Option<String>,
551 name: Option<String>,
552 approx_size_bytes: Option<i64>,
553 }
554
555 let url = format!(
556 "{}://{}/v1/client/onde_sdk/model_config?app_id={}&app_secret={}",
557 environment.api_protocol(),
558 environment.api_host(),
559 onde_app_id,
560 onde_app_secret,
561 );
562
563 let response = reqwest::Client::new()
564 .get(&url)
565 .header("Content-Type", "application/json")
566 .send()
567 .await
568 .map_err(|e| InferenceError::ModelBuild {
569 reason: format!("SDK model_config request failed: {e}"),
570 })?;
571
572 if response.status().as_u16() == 404 {
573 warn!(
574 "ChatEngine: no model assigned to Onde app {onde_app_id}; loading platform default."
575 );
576 return self
577 .load_gguf_model(GgufModelConfig::platform_default(), system_prompt, sampling)
578 .await;
579 }
580
581 if !response.status().is_success() {
582 return Err(InferenceError::ModelBuild {
583 reason: format!(
584 "SDK model_config returned HTTP {}",
585 response.status().as_u16()
586 ),
587 });
588 }
589
590 let resp: ModelConfigResponse =
591 response
592 .json()
593 .await
594 .map_err(|e| InferenceError::ModelBuild {
595 reason: format!("Failed to parse model_config response: {e}"),
596 })?;
597
598 let hf_repo_id = resp.hf_repo_id.as_deref().unwrap_or_default();
599 let gguf_file = resp.gguf_file.as_deref().unwrap_or_default();
600
601 if hf_repo_id.is_empty() || gguf_file.is_empty() {
602 warn!(
603 "ChatEngine: model_config response missing hf_repo_id or gguf_file; loading platform default."
604 );
605 return self
606 .load_gguf_model(GgufModelConfig::platform_default(), system_prompt, sampling)
607 .await;
608 }
609
610 #[cfg(target_os = "android")]
611 let tok_model_id = super::models::tok_model_id_for_repo(hf_repo_id).map(|s| s.to_string());
612 #[cfg(not(target_os = "android"))]
613 let tok_model_id: Option<String> = None;
614
615 let approx_memory = resp
616 .approx_size_bytes
617 .map(|b| {
618 let gb = b as f64 / 1_073_741_824.0;
619 if gb >= 1.0 {
620 format!("~{:.2} GB", gb)
621 } else {
622 format!("~{} MB", b / 1_048_576)
623 }
624 })
625 .unwrap_or_else(|| "—".to_string());
626
627 info!(
628 "ChatEngine: resolved SDK model assignment → {} / {} ({})",
629 hf_repo_id, gguf_file, approx_memory
630 );
631
632 let config = GgufModelConfig {
633 model_id: hf_repo_id.to_string(),
634 files: vec![gguf_file.to_string()],
635 tok_model_id,
636 display_name: resp.name.unwrap_or_else(|| hf_repo_id.to_string()),
637 approx_memory,
638 chat_template: None,
639 };
640
641 self.load_gguf_model(config, system_prompt, sampling).await
642 }
643
644 /// Load an ISQ (in-situ quantised) model into the engine.
645 ///
646 /// Unlike [`load_gguf_model`], this downloads the full-precision safetensors
647 /// from HuggingFace and quantises the weights in-situ on Metal (macOS only).
648 /// The quantisation happens once at load time; subsequent inference uses the
649 /// compressed weights.
650 ///
651 /// # macOS only
652 ///
653 /// ISQ with Metal requires the `metal` feature on mistral.rs and is therefore
654 /// restricted to macOS at the Rust level. Other platforms should continue to
655 /// use [`load_gguf_model`] with pre-quantised GGUF files.
656 ///
657 /// # Arguments
658 ///
659 /// * `config` — Which model to load and with how many ISQ bits.
660 /// * `system_prompt` — Optional system prompt prepended to every request.
661 /// * `sampling` — Sampling parameters; `None` uses [`SamplingConfig::default`].
662 ///
663 /// # Errors
664 ///
665 /// Returns [`InferenceError::ModelBuild`] if the model fails to download
666 /// or quantise.
667 #[cfg(target_os = "macos")]
668 pub async fn load_isq_model(
669 &self,
670 config: IsqModelConfig,
671 system_prompt: Option<String>,
672 sampling: Option<SamplingConfig>,
673 ) -> Result<std::time::Duration, InferenceError> {
674 use log::info;
675 use std::time::Instant;
676
677 info!(
678 "ChatEngine: loading ISQ model {} (bits={})",
679 config.model_id, config.isq_bits
680 );
681
682 // Clean up stale HF cache artefacts before loading.
683 crate::hf_cache::clean_stale_lock_files(&config.model_id);
684 crate::hf_cache::repair_hf_cache_symlinks(&config.model_id);
685
686 let start = Instant::now();
687
688 // Choose ISQ bit width.
689 let isq_bits = match config.isq_bits {
690 8 => IsqBits::Eight,
691 _ => IsqBits::Four, // default to 4-bit
692 };
693
694 let model = TextModelBuilder::new(&config.model_id)
695 .with_token_source(super::token::hf_token_source())
696 .with_auto_isq(isq_bits)
697 .with_logging()
698 .build()
699 .await
700 .map_err(|e| InferenceError::ModelBuild {
701 reason: format!("Failed to build ISQ model {}: {}", config.display_name, e),
702 })?;
703
704 let elapsed = start.elapsed();
705
706 let sampling = sampling.unwrap_or_default();
707
708 info!(
709 "ChatEngine: ISQ model {} loaded in {} (sampling: temp={:?}, max_tokens={:?})",
710 config.display_name,
711 format_duration(elapsed),
712 sampling.temperature,
713 sampling.max_tokens,
714 );
715
716 let mut guard = self.inner.lock().await;
717 *guard = Some(LoadedModel {
718 model: Arc::new(model),
719 config: LoadedModelConfig::Isq(config),
720 history: Vec::new(),
721 system_prompt,
722 sampling,
723 });
724
725 Ok(elapsed)
726 }
727
728 /// Unload the current model, freeing all memory.
729 ///
730 /// Dropping the `Model` sends a `Terminate` message to the mistral.rs
731 /// engine thread, which tears down the KV cache, activations, and model
732 /// weights.
733 ///
734 /// Returns the display name of the model that was unloaded, or `None`
735 /// if no model was loaded.
736 pub async fn unload_model(&self) -> Option<String> {
737 let mut guard = self.inner.lock().await;
738 if let Some(loaded) = guard.take() {
739 let name = loaded.config.display_name().to_string();
740 log::info!("ChatEngine: unloading model {}", name);
741 // `loaded` is dropped here → Model → Arc<MistralRs> (last ref)
742 // → MistralRs::drop() → engine thread termination.
743 Some(name)
744 } else {
745 log::debug!("ChatEngine: unload_model called but no model was loaded.");
746 None
747 }
748 }
749
750 /// Check whether a model is currently loaded.
751 pub async fn is_loaded(&self) -> bool {
752 self.inner.lock().await.is_some()
753 }
754
755 /// Get a snapshot of the engine's current state.
756 pub async fn info(&self) -> EngineInfo {
757 let guard = self.inner.lock().await;
758 match guard.as_ref() {
759 Some(loaded) => EngineInfo {
760 status: EngineStatus::Ready,
761 model_name: Some(loaded.config.display_name().to_string()),
762 approx_memory: Some(loaded.config.approx_memory().to_string()),
763 history_length: loaded.history.len() as u64,
764 },
765 None => EngineInfo {
766 status: EngineStatus::Unloaded,
767 model_name: None,
768 approx_memory: None,
769 history_length: 0u64,
770 },
771 }
772 }
773
774 // ── System prompt ────────────────────────────────────────────────────
775
776 /// Set or replace the system prompt.
777 ///
778 /// The system prompt is prepended to every inference request (it is NOT
779 /// stored in the conversation history).
780 pub async fn set_system_prompt(&self, prompt: impl Into<String>) {
781 if let Some(loaded) = self.inner.lock().await.as_mut() {
782 loaded.system_prompt = Some(prompt.into());
783 }
784 }
785
786 /// Clear the system prompt.
787 pub async fn clear_system_prompt(&self) {
788 if let Some(loaded) = self.inner.lock().await.as_mut() {
789 loaded.system_prompt = None;
790 }
791 }
792
793 // ── Sampling configuration ───────────────────────────────────────────
794
795 /// Replace the sampling configuration.
796 pub async fn set_sampling(&self, sampling: SamplingConfig) {
797 if let Some(loaded) = self.inner.lock().await.as_mut() {
798 loaded.sampling = sampling;
799 }
800 }
801
802 // ── Conversation history ─────────────────────────────────────────────
803
804 /// Get a clone of the full conversation history.
805 ///
806 /// Tool-related history entries are converted back to plain
807 /// [`ChatMessage`] values (assistant tool-call entries become assistant
808 /// messages with just their text content; tool results are omitted).
809 pub async fn history(&self) -> Vec<ChatMessage> {
810 let guard = self.inner.lock().await;
811 match guard.as_ref() {
812 Some(loaded) => loaded
813 .history
814 .iter()
815 .filter_map(|entry| match entry {
816 HistoryEntry::Text(msg) => Some(msg.clone()),
817 HistoryEntry::AssistantToolCall { content, .. } => {
818 Some(ChatMessage::assistant(content))
819 }
820 HistoryEntry::ToolResult { .. } => None,
821 })
822 .collect(),
823 None => Vec::new(),
824 }
825 }
826
827 /// Clear the conversation history but keep the model loaded.
828 ///
829 /// Returns the number of turns that were removed.
830 pub async fn clear_history(&self) -> usize {
831 let mut guard = self.inner.lock().await;
832 match guard.as_mut() {
833 Some(loaded) => {
834 let count = loaded.history.len();
835 loaded.history.clear();
836 log::info!("ChatEngine: cleared {} history turns.", count);
837 count
838 }
839 None => 0,
840 }
841 }
842
843 /// Append a message to history without running inference.
844 ///
845 /// Useful for restoring a saved conversation or injecting context.
846 pub async fn push_history(&self, message: ChatMessage) {
847 if let Some(loaded) = self.inner.lock().await.as_mut() {
848 loaded.history.push(HistoryEntry::Text(message));
849 }
850 }
851
852 // ── Non-streaming inference ──────────────────────────────────────────
853
854 /// Send a user message and receive a complete assistant reply.
855 ///
856 /// The user message and assistant reply are automatically appended to
857 /// the conversation history on success.
858 ///
859 /// # Errors
860 ///
861 /// - [`InferenceError::NoModelLoaded`] if no model is loaded.
862 /// - [`InferenceError::Inference`] if the model fails to generate.
863 pub async fn send_message(
864 &self,
865 user_message: impl Into<String>,
866 ) -> Result<InferenceResult, InferenceError> {
867 let user_message = user_message.into();
868
869 // ── 1. Snapshot model handle + build request, then release lock ──
870 let (model, request, pulse_model_id) = {
871 let guard = self.inner.lock().await;
872 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
873 let request = self::build_request(loaded, &user_message);
874 let pulse_model_id = match &loaded.config {
875 LoadedModelConfig::Gguf(c) => c.model_id.clone(),
876 #[cfg(target_os = "macos")]
877 LoadedModelConfig::Isq(c) => c.model_id.clone(),
878 };
879 (loaded.model.clone(), request, pulse_model_id)
880 }; // ← mutex released before inference
881
882 log::info!(
883 "ChatEngine: inference START — message: \"{}\"",
884 truncate_for_log(&user_message, 100)
885 );
886
887 // ── 2. Run inference (potentially slow — mutex is NOT held) ──────
888 let start = std::time::Instant::now();
889 let response =
890 model
891 .send_chat_request(request)
892 .await
893 .map_err(|e| InferenceError::Inference {
894 reason: e.to_string(),
895 })?;
896 let elapsed = start.elapsed();
897
898 let reply = response.choices[0]
899 .message
900 .content
901 .as_ref()
902 .map(|c| c.trim().to_string())
903 .unwrap_or_else(|| "(empty response)".to_string());
904
905 let finish_reason = response.choices[0].finish_reason.clone();
906
907 log::info!(
908 "ChatEngine: inference END — {} — reply: \"{}\"",
909 format_duration(elapsed),
910 truncate_for_log(&reply, 100)
911 );
912
913 // ── 3. Persist turns to history (brief lock re-acquisition) ──────
914 {
915 let mut guard = self.inner.lock().await;
916 if let Some(loaded) = guard.as_mut() {
917 loaded
918 .history
919 .push(HistoryEntry::Text(ChatMessage::user(user_message)));
920 loaded
921 .history
922 .push(HistoryEntry::Text(ChatMessage::assistant(reply.clone())));
923 }
924 }
925
926 if let Some(pulse) = self.pulse() {
927 pulse.record_inference(
928 pulse_model_id,
929 crate::pulse::next_request_id(),
930 elapsed.as_millis() as u64,
931 "success".to_string(),
932 );
933 }
934
935 Ok(InferenceResult {
936 text: reply,
937 duration_secs: elapsed.as_secs_f64(),
938 duration_display: format_duration(elapsed),
939 finish_reason,
940 tool_calls: vec![],
941 })
942 }
943
944 /// Run inference on an explicit list of messages WITHOUT modifying the
945 /// engine's internal history.
946 ///
947 /// This is useful for one-shot prompts (e.g. prompt enhancement) where
948 /// you don't want side-effects on the main conversation.
949 ///
950 /// The system prompt set on the engine is still prepended.
951 pub async fn generate(
952 &self,
953 messages: Vec<ChatMessage>,
954 sampling: Option<SamplingConfig>,
955 ) -> Result<InferenceResult, InferenceError> {
956 let (model, request) = {
957 let guard = self.inner.lock().await;
958 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
959
960 let sampling = sampling.as_ref().unwrap_or(&loaded.sampling);
961 let mut req = RequestBuilder::new();
962
963 // Apply sampling parameters.
964 req = apply_sampling(req, sampling);
965
966 // System prompt.
967 if let Some(ref sp) = loaded.system_prompt {
968 req = req.add_message(TextMessageRole::System, sp);
969 }
970
971 // Provided messages.
972 for msg in &messages {
973 req = req.add_message(chat_role_to_mistral(&msg.role), &msg.content);
974 }
975
976 (loaded.model.clone(), req)
977 };
978
979 let start = std::time::Instant::now();
980 let response =
981 model
982 .send_chat_request(request)
983 .await
984 .map_err(|e| InferenceError::Inference {
985 reason: e.to_string(),
986 })?;
987 let elapsed = start.elapsed();
988
989 let reply = response.choices[0]
990 .message
991 .content
992 .as_ref()
993 .map(|c| c.trim().to_string())
994 .unwrap_or_else(|| "(empty response)".to_string());
995
996 let finish_reason = response.choices[0].finish_reason.clone();
997
998 Ok(InferenceResult {
999 text: reply,
1000 duration_secs: elapsed.as_secs_f64(),
1001 duration_display: format_duration(elapsed),
1002 finish_reason,
1003 tool_calls: vec![],
1004 })
1005 }
1006
1007 // ── Streaming inference ──────────────────────────────────────────────
1008
1009 /// Send a user message and receive an async stream of token chunks.
1010 ///
1011 /// The user message and full assembled reply are automatically appended
1012 /// to the conversation history once the stream finishes.
1013 ///
1014 /// # Returns
1015 ///
1016 /// A `tokio::sync::mpsc::Receiver<StreamChunk>`. The caller should
1017 /// receive from it until `chunk.done == true` or the channel closes.
1018 ///
1019 /// # Errors
1020 ///
1021 /// Returns immediately with [`InferenceError::NoModelLoaded`] if no
1022 /// model is loaded. Inference errors are delivered as the final chunk
1023 /// with `done = true` and an empty delta.
1024 pub async fn stream_message(
1025 &self,
1026 user_message: impl Into<String>,
1027 ) -> Result<tokio::sync::mpsc::Receiver<StreamChunk>, InferenceError> {
1028 let user_message = user_message.into();
1029
1030 let (model, request) = {
1031 let guard = self.inner.lock().await;
1032 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
1033 let request = self::build_request(loaded, &user_message);
1034 (loaded.model.clone(), request)
1035 };
1036
1037 let (tx, rx) = tokio::sync::mpsc::channel::<StreamChunk>(64);
1038
1039 // We need a reference to `self` inside the spawned task so we can
1040 // update history when the stream completes. Since ChatEngine is
1041 // behind a shared reference (callers hold &self or Arc<ChatEngine>),
1042 // we borrow the inner Mutex via a raw pointer. This is safe because
1043 // ChatEngine's lifetime exceeds the spawned task (it's typically in
1044 // a Lazy static or shared application state).
1045 //
1046 // Alternative: require ChatEngine to always be wrapped in Arc and
1047 // accept `self: Arc<Self>`. We avoid that to keep the API simple.
1048 let inner_ptr = &self.inner as *const Mutex<Option<LoadedModel>>;
1049 // SAFETY: ChatEngine is stored in a Lazy static or equivalent with
1050 // 'static lifetime. The spawned task cannot outlive the engine.
1051 let inner_ref: &'static Mutex<Option<LoadedModel>> = unsafe { &*inner_ptr };
1052
1053 let user_msg_clone = user_message.clone();
1054
1055 tokio::task::spawn(async move {
1056 // `model` is an `Arc<Model>` moved into this task, so it's owned.
1057 // `stream_chat_request` borrows `&Model` — the borrow is scoped
1058 // to this async block and does NOT need to be `'static`.
1059 let stream_result = model.stream_chat_request(request).await;
1060
1061 match stream_result {
1062 Ok(mut stream) => {
1063 let mut assembled = String::new();
1064 let mut last_finish_reason = None;
1065
1066 // `Stream::next()` is an inherent async method on
1067 // mistralrs::Stream that returns `Option<Response>`.
1068 // No `futures::StreamExt` import needed.
1069 while let Some(response) = stream.next().await {
1070 match response {
1071 mistralrs::Response::Chunk(chunk) => {
1072 if let Some(choice) = chunk.choices.first() {
1073 if let Some(ref text) = choice.delta.content {
1074 assembled.push_str(text);
1075 let _ = tx
1076 .send(StreamChunk {
1077 delta: text.clone(),
1078 done: false,
1079 finish_reason: None,
1080 })
1081 .await;
1082 }
1083 if let Some(ref reason) = choice.finish_reason {
1084 last_finish_reason = Some(reason.clone());
1085 }
1086 }
1087 }
1088 mistralrs::Response::Done(_) => {
1089 // Non-streaming response arrived on a streaming
1090 // channel — should not happen, but handle gracefully.
1091 break;
1092 }
1093 mistralrs::Response::InternalError(e) => {
1094 log::error!("ChatEngine stream internal error: {}", e);
1095 break;
1096 }
1097 mistralrs::Response::ValidationError(e) => {
1098 log::error!("ChatEngine stream validation error: {}", e);
1099 break;
1100 }
1101 mistralrs::Response::ModelError(msg, _) => {
1102 log::error!("ChatEngine stream model error: {}", msg);
1103 break;
1104 }
1105 _ => {
1106 // Completion / Image / Speech variants — not expected
1107 // for chat streaming.
1108 break;
1109 }
1110 }
1111 }
1112
1113 // Persist turns to history.
1114 {
1115 let mut guard = inner_ref.lock().await;
1116 if let Some(loaded) = guard.as_mut() {
1117 loaded
1118 .history
1119 .push(HistoryEntry::Text(ChatMessage::user(user_msg_clone)));
1120 loaded
1121 .history
1122 .push(HistoryEntry::Text(ChatMessage::assistant(assembled.trim())));
1123 }
1124 }
1125
1126 // Send the final "done" chunk.
1127 let _ = tx
1128 .send(StreamChunk {
1129 delta: String::new(),
1130 done: true,
1131 finish_reason: last_finish_reason,
1132 })
1133 .await;
1134 }
1135 Err(e) => {
1136 log::error!("ChatEngine: stream_chat_request failed: {}", e);
1137 let _ = tx
1138 .send(StreamChunk {
1139 delta: String::new(),
1140 done: true,
1141 finish_reason: Some(format!("error: {}", e)),
1142 })
1143 .await;
1144 }
1145 }
1146 });
1147
1148 Ok(rx)
1149 }
1150
1151 // ── Tool-aware inference (Rust-only) ─────────────────────────────────
1152
1153 /// Send a user message with tool definitions available. Non-streaming.
1154 ///
1155 /// If the model decides to call tools, [`ToolAwareResult::tool_calls`]
1156 /// will be non-empty and `finish_reason` will typically be `"tool_calls"`.
1157 ///
1158 /// The user message and response (including any tool calls) are
1159 /// automatically added to conversation history.
1160 pub async fn send_message_with_tools(
1161 &self,
1162 user_message: impl Into<String>,
1163 tools: &[ToolDefinition],
1164 ) -> Result<ToolAwareResult, InferenceError> {
1165 let user_message = user_message.into();
1166
1167 let (model, request) = {
1168 let guard = self.inner.lock().await;
1169 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
1170 let request = build_request_with_tools(loaded, &user_message, tools);
1171 (loaded.model.clone(), request)
1172 };
1173
1174 log::info!(
1175 "ChatEngine: tool inference START — message: \"{}\"",
1176 truncate_for_log(&user_message, 100)
1177 );
1178
1179 let start = std::time::Instant::now();
1180 let response =
1181 model
1182 .send_chat_request(request)
1183 .await
1184 .map_err(|e| InferenceError::Inference {
1185 reason: e.to_string(),
1186 })?;
1187 let elapsed = start.elapsed();
1188
1189 let choice = &response.choices[0];
1190 let reply = choice
1191 .message
1192 .content
1193 .as_ref()
1194 .map(|c| c.trim().to_string())
1195 .unwrap_or_default();
1196 let finish_reason = choice.finish_reason.clone();
1197 let tool_calls = parse_tool_calls(choice);
1198
1199 log::info!(
1200 "ChatEngine: tool inference END — {} — tool_calls: {} — reply: \"{}\"",
1201 format_duration(elapsed),
1202 tool_calls.len(),
1203 truncate_for_log(&reply, 100)
1204 );
1205
1206 // Persist to history.
1207 {
1208 let mut guard = self.inner.lock().await;
1209 if let Some(loaded) = guard.as_mut() {
1210 loaded
1211 .history
1212 .push(HistoryEntry::Text(ChatMessage::user(&user_message)));
1213 if tool_calls.is_empty() {
1214 loaded
1215 .history
1216 .push(HistoryEntry::Text(ChatMessage::assistant(&reply)));
1217 } else {
1218 loaded.history.push(HistoryEntry::AssistantToolCall {
1219 content: reply.clone(),
1220 tool_calls: tool_calls.clone(),
1221 });
1222 }
1223 }
1224 }
1225
1226 Ok(ToolAwareResult {
1227 text: reply,
1228 tool_calls,
1229 duration_secs: elapsed.as_secs_f64(),
1230 duration_display: format_duration(elapsed),
1231 finish_reason,
1232 })
1233 }
1234
1235 /// Send tool execution results back to the model. Non-streaming.
1236 ///
1237 /// Call this after executing the tools from a previous
1238 /// [`send_message_with_tools`](Self::send_message_with_tools) response.
1239 /// Pass the same (or updated) tool definitions if the model should be
1240 /// allowed to make further tool calls.
1241 pub async fn send_tool_results(
1242 &self,
1243 results: Vec<ToolResult>,
1244 tools: Option<&[ToolDefinition]>,
1245 ) -> Result<ToolAwareResult, InferenceError> {
1246 let (model, request) = {
1247 let guard = self.inner.lock().await;
1248 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
1249 let request = build_tool_results_request(loaded, &results, tools);
1250 (loaded.model.clone(), request)
1251 };
1252
1253 log::info!(
1254 "ChatEngine: tool results inference START — {} results",
1255 results.len()
1256 );
1257
1258 let start = std::time::Instant::now();
1259 let response =
1260 model
1261 .send_chat_request(request)
1262 .await
1263 .map_err(|e| InferenceError::Inference {
1264 reason: e.to_string(),
1265 })?;
1266 let elapsed = start.elapsed();
1267
1268 let choice = &response.choices[0];
1269 let reply = choice
1270 .message
1271 .content
1272 .as_ref()
1273 .map(|c| c.trim().to_string())
1274 .unwrap_or_default();
1275 let finish_reason = choice.finish_reason.clone();
1276 let tool_calls = parse_tool_calls(choice);
1277
1278 log::info!(
1279 "ChatEngine: tool results inference END — {} — tool_calls: {} — reply: \"{}\"",
1280 format_duration(elapsed),
1281 tool_calls.len(),
1282 truncate_for_log(&reply, 100)
1283 );
1284
1285 // Persist to history.
1286 {
1287 let mut guard = self.inner.lock().await;
1288 if let Some(loaded) = guard.as_mut() {
1289 for result in &results {
1290 loaded.history.push(HistoryEntry::ToolResult {
1291 tool_call_id: result.tool_call_id.clone(),
1292 content: result.content.clone(),
1293 });
1294 }
1295 if tool_calls.is_empty() {
1296 loaded
1297 .history
1298 .push(HistoryEntry::Text(ChatMessage::assistant(&reply)));
1299 } else {
1300 loaded.history.push(HistoryEntry::AssistantToolCall {
1301 content: reply.clone(),
1302 tool_calls: tool_calls.clone(),
1303 });
1304 }
1305 }
1306 }
1307
1308 Ok(ToolAwareResult {
1309 text: reply,
1310 tool_calls,
1311 duration_secs: elapsed.as_secs_f64(),
1312 duration_display: format_duration(elapsed),
1313 finish_reason,
1314 })
1315 }
1316
1317 /// Stream tool execution results back to the model.
1318 ///
1319 /// Similar to [`send_tool_results`](Self::send_tool_results) but returns
1320 /// a streaming receiver. Tool calls in the streaming response are NOT
1321 /// parsed — this is intended for the final text response after all tool
1322 /// rounds complete.
1323 pub async fn stream_tool_results(
1324 &self,
1325 results: Vec<ToolResult>,
1326 tools: Option<Vec<ToolDefinition>>,
1327 ) -> Result<tokio::sync::mpsc::Receiver<StreamChunk>, InferenceError> {
1328 let (model, request) = {
1329 let guard = self.inner.lock().await;
1330 let loaded = guard.as_ref().ok_or(InferenceError::NoModelLoaded)?;
1331 let request = build_tool_results_request(loaded, &results, tools.as_deref());
1332 (loaded.model.clone(), request)
1333 };
1334
1335 let (tx, rx) = tokio::sync::mpsc::channel::<StreamChunk>(64);
1336
1337 let inner_ptr = &self.inner as *const Mutex<Option<LoadedModel>>;
1338 // SAFETY: ChatEngine is stored in a Lazy static or equivalent with
1339 // 'static lifetime. The spawned task cannot outlive the engine.
1340 let inner_ref: &'static Mutex<Option<LoadedModel>> = unsafe { &*inner_ptr };
1341
1342 let results_clone = results;
1343
1344 tokio::task::spawn(async move {
1345 let stream_result = model.stream_chat_request(request).await;
1346
1347 match stream_result {
1348 Ok(mut stream) => {
1349 let mut assembled = String::new();
1350 let mut last_finish_reason = None;
1351
1352 while let Some(response) = stream.next().await {
1353 match response {
1354 mistralrs::Response::Chunk(chunk) => {
1355 if let Some(choice) = chunk.choices.first() {
1356 if let Some(ref text) = choice.delta.content {
1357 assembled.push_str(text);
1358 let _ = tx
1359 .send(StreamChunk {
1360 delta: text.clone(),
1361 done: false,
1362 finish_reason: None,
1363 })
1364 .await;
1365 }
1366 if let Some(ref reason) = choice.finish_reason {
1367 last_finish_reason = Some(reason.clone());
1368 }
1369 }
1370 }
1371 mistralrs::Response::Done(_) => break,
1372 mistralrs::Response::InternalError(e) => {
1373 log::error!("stream_tool_results internal error: {}", e);
1374 break;
1375 }
1376 mistralrs::Response::ValidationError(e) => {
1377 log::error!("stream_tool_results validation error: {}", e);
1378 break;
1379 }
1380 mistralrs::Response::ModelError(msg, _) => {
1381 log::error!("stream_tool_results model error: {}", msg);
1382 break;
1383 }
1384 _ => break,
1385 }
1386 }
1387
1388 // Persist to history.
1389 {
1390 let mut guard = inner_ref.lock().await;
1391 if let Some(loaded) = guard.as_mut() {
1392 for result in &results_clone {
1393 loaded.history.push(HistoryEntry::ToolResult {
1394 tool_call_id: result.tool_call_id.clone(),
1395 content: result.content.clone(),
1396 });
1397 }
1398 loaded
1399 .history
1400 .push(HistoryEntry::Text(ChatMessage::assistant(assembled.trim())));
1401 }
1402 }
1403
1404 let _ = tx
1405 .send(StreamChunk {
1406 delta: String::new(),
1407 done: true,
1408 finish_reason: last_finish_reason,
1409 })
1410 .await;
1411 }
1412 Err(e) => {
1413 log::error!("stream_tool_results failed: {}", e);
1414 let _ = tx
1415 .send(StreamChunk {
1416 delta: String::new(),
1417 done: true,
1418 finish_reason: Some(format!("error: {}", e)),
1419 })
1420 .await;
1421 }
1422 }
1423 });
1424
1425 Ok(rx)
1426 }
1427 }
1428
1429 // ═════════════════════════════════════════════════════════════════════════════
1430 // Default impl
1431 // ═════════════════════════════════════════════════════════════════════════════
1432
1433 #[cfg(any(
1434 target_os = "macos",
1435 target_os = "ios",
1436 target_os = "tvos",
1437 target_os = "visionos",
1438 target_os = "watchos",
1439 target_os = "windows",
1440 target_os = "linux",
1441 target_os = "android"
1442 ))]
1443 impl Default for ChatEngine {
1444 fn default() -> Self {
1445 Self::new()
1446 }
1447 }
1448
1449 // ═════════════════════════════════════════════════════════════════════════════
1450 // Internal helpers (platform-gated)
1451 // ═════════════════════════════════════════════════════════════════════════════
1452
1453 /// Build a `RequestBuilder` from the current engine state and a new user message.
1454 #[cfg(any(
1455 target_os = "macos",
1456 target_os = "ios",
1457 target_os = "tvos",
1458 target_os = "visionos",
1459 target_os = "watchos",
1460 target_os = "windows",
1461 target_os = "linux",
1462 target_os = "android"
1463 ))]
1464 fn build_request(loaded: &LoadedModel, user_message: &str) -> RequestBuilder {
1465 let mut req = RequestBuilder::new();
1466
1467 // Apply sampling.
1468 req = apply_sampling(req, &loaded.sampling);
1469
1470 // System prompt.
1471 if let Some(ref sp) = loaded.system_prompt {
1472 req = req.add_message(TextMessageRole::System, sp);
1473 }
1474
1475 // Replay conversation history for multi-turn context.
1476 for entry in &loaded.history {
1477 match entry {
1478 HistoryEntry::Text(msg) => {
1479 req = req.add_message(chat_role_to_mistral(&msg.role), &msg.content);
1480 }
1481 HistoryEntry::AssistantToolCall { content, .. } => {
1482 // When replaying without tools, just use the text content.
1483 req = req.add_message(TextMessageRole::Assistant, content);
1484 }
1485 HistoryEntry::ToolResult { .. } => {
1486 // Skip tool results when not in tool mode.
1487 }
1488 }
1489 }
1490
1491 // Append the current user message.
1492 req = req.add_message(TextMessageRole::User, user_message);
1493
1494 req
1495 }
1496
1497 /// Build a request with tool support, properly replaying tool-related history.
1498 #[cfg(any(
1499 target_os = "macos",
1500 target_os = "ios",
1501 target_os = "tvos",
1502 target_os = "visionos",
1503 target_os = "watchos",
1504 target_os = "windows",
1505 target_os = "linux",
1506 target_os = "android"
1507 ))]
1508 fn build_request_with_tools(
1509 loaded: &LoadedModel,
1510 user_message: &str,
1511 tools: &[ToolDefinition],
1512 ) -> RequestBuilder {
1513 let mut req = RequestBuilder::new();
1514 req = apply_sampling(req, &loaded.sampling);
1515
1516 if let Some(ref sp) = loaded.system_prompt {
1517 req = req.add_message(TextMessageRole::System, sp);
1518 }
1519
1520 req = replay_history_with_tools(req, &loaded.history);
1521
1522 req = req.add_message(TextMessageRole::User, user_message);
1523
1524 // Add tool definitions.
1525 req = attach_tools(req, tools);
1526
1527 req
1528 }
1529
1530 /// Build a request for sending tool results back to the model.
1531 #[cfg(any(
1532 target_os = "macos",
1533 target_os = "ios",
1534 target_os = "tvos",
1535 target_os = "visionos",
1536 target_os = "watchos",
1537 target_os = "windows",
1538 target_os = "linux",
1539 target_os = "android"
1540 ))]
1541 fn build_tool_results_request(
1542 loaded: &LoadedModel,
1543 results: &[ToolResult],
1544 tools: Option<&[ToolDefinition]>,
1545 ) -> RequestBuilder {
1546 let mut req = RequestBuilder::new();
1547 req = apply_sampling(req, &loaded.sampling);
1548
1549 if let Some(ref sp) = loaded.system_prompt {
1550 req = req.add_message(TextMessageRole::System, sp);
1551 }
1552
1553 // Replay full history (including tool call entries).
1554 req = replay_history_with_tools(req, &loaded.history);
1555
1556 // Add new tool results.
1557 for result in results {
1558 req = req.add_tool_message(&result.content, &result.tool_call_id);
1559 }
1560
1561 // Optionally add tools for further rounds.
1562 if let Some(tool_defs) = tools {
1563 req = attach_tools(req, tool_defs);
1564 }
1565
1566 req
1567 }
1568
1569 /// Replay conversation history entries onto a [`RequestBuilder`], preserving
1570 /// tool call and tool result messages.
1571 #[cfg(any(
1572 target_os = "macos",
1573 target_os = "ios",
1574 target_os = "tvos",
1575 target_os = "visionos",
1576 target_os = "watchos",
1577 target_os = "windows",
1578 target_os = "linux",
1579 target_os = "android"
1580 ))]
1581 fn replay_history_with_tools(mut req: RequestBuilder, history: &[HistoryEntry]) -> RequestBuilder {
1582 for entry in history {
1583 match entry {
1584 HistoryEntry::Text(msg) => {
1585 req = req.add_message(chat_role_to_mistral(&msg.role), &msg.content);
1586 }
1587 HistoryEntry::AssistantToolCall {
1588 content,
1589 tool_calls,
1590 } => {
1591 let mistral_tcs: Vec<ToolCallResponse> = tool_calls
1592 .iter()
1593 .enumerate()
1594 .map(|(i, tc)| ToolCallResponse {
1595 index: i,
1596 id: tc.id.clone(),
1597 tp: ToolCallType::Function,
1598 function: CalledFunction {
1599 name: tc.function_name.clone(),
1600 arguments: tc.arguments.clone(),
1601 },
1602 })
1603 .collect();
1604 req = req.add_message_with_tool_call(
1605 TextMessageRole::Assistant,
1606 content.clone(),
1607 mistral_tcs,
1608 );
1609 }
1610 HistoryEntry::ToolResult {
1611 tool_call_id,
1612 content,
1613 } => {
1614 req = req.add_tool_message(content, tool_call_id);
1615 }
1616 }
1617 }
1618 req
1619 }
1620
1621 /// Attach [`ToolDefinition`] slices as mistralrs [`Tool`] values and set
1622 /// [`ToolChoice::Auto`].
1623 #[cfg(any(
1624 target_os = "macos",
1625 target_os = "ios",
1626 target_os = "tvos",
1627 target_os = "visionos",
1628 target_os = "watchos",
1629 target_os = "windows",
1630 target_os = "linux",
1631 target_os = "android"
1632 ))]
1633 fn attach_tools(mut req: RequestBuilder, tools: &[ToolDefinition]) -> RequestBuilder {
1634 if !tools.is_empty() {
1635 let mistral_tools: Vec<MistralTool> = tools
1636 .iter()
1637 .map(|td| {
1638 let params: HashMap<String, serde_json::Value> =
1639 match serde_json::from_str(&td.parameters_schema) {
1640 Ok(p) => p,
1641 Err(err) => {
1642 log::warn!(
1643 "tool '{}': malformed parameters_schema JSON ({}), using empty params",
1644 td.name,
1645 err
1646 );
1647 HashMap::new()
1648 }
1649 };
1650 MistralTool {
1651 tp: ToolType::Function,
1652 function: Function {
1653 description: Some(td.description.clone()),
1654 name: td.name.clone(),
1655 parameters: Some(params),
1656 strict: Some(true),
1657 },
1658 }
1659 })
1660 .collect();
1661 req = req.set_tools(mistral_tools);
1662 req = req.set_tool_choice(ToolChoice::Auto);
1663 }
1664 req
1665 }
1666
1667 /// Parse tool calls from a mistralrs response choice.
1668 #[cfg(any(
1669 target_os = "macos",
1670 target_os = "ios",
1671 target_os = "tvos",
1672 target_os = "visionos",
1673 target_os = "watchos",
1674 target_os = "windows",
1675 target_os = "linux",
1676 target_os = "android"
1677 ))]
1678 fn parse_tool_calls(choice: &mistralrs::Choice) -> Vec<ToolCallRequest> {
1679 choice
1680 .message
1681 .tool_calls
1682 .as_ref()
1683 .map(|tcs| {
1684 tcs.iter()
1685 .enumerate()
1686 .map(|(i, tc)| {
1687 let id = if tc.id.is_empty() {
1688 log::warn!("tool call at index {i} has empty id — generating fallback");
1689 format!("call_{i}")
1690 } else {
1691 tc.id.clone()
1692 };
1693 if serde_json::from_str::<serde_json::Value>(&tc.function.arguments).is_err() {
1694 log::warn!(
1695 "tool call '{}' has malformed arguments JSON: {}",
1696 tc.function.name,
1697 tc.function.arguments.chars().take(200).collect::<String>()
1698 );
1699 }
1700 ToolCallRequest {
1701 id,
1702 function_name: tc.function.name.clone(),
1703 arguments: tc.function.arguments.clone(),
1704 }
1705 })
1706 .collect()
1707 })
1708 .unwrap_or_default()
1709 }
1710
1711 /// Apply a [`SamplingConfig`] to a [`RequestBuilder`].
1712 #[cfg(any(
1713 target_os = "macos",
1714 target_os = "ios",
1715 target_os = "tvos",
1716 target_os = "visionos",
1717 target_os = "watchos",
1718 target_os = "windows",
1719 target_os = "linux",
1720 target_os = "android"
1721 ))]
1722 fn apply_sampling(mut req: RequestBuilder, sampling: &SamplingConfig) -> RequestBuilder {
1723 if let Some(temp) = sampling.temperature {
1724 req = req.set_sampler_temperature(temp);
1725 }
1726 if let Some(top_p) = sampling.top_p {
1727 req = req.set_sampler_topp(top_p);
1728 }
1729 if let Some(top_k) = sampling.top_k {
1730 req = req.set_sampler_topk(top_k as usize);
1731 }
1732 if let Some(min_p) = sampling.min_p {
1733 req = req.set_sampler_minp(min_p);
1734 }
1735 if let Some(max_tokens) = sampling.max_tokens {
1736 req = req.set_sampler_max_len(max_tokens as usize);
1737 }
1738 if let Some(freq) = sampling.frequency_penalty {
1739 req = req.set_sampler_frequency_penalty(freq);
1740 }
1741 if let Some(pres) = sampling.presence_penalty {
1742 req = req.set_sampler_presence_penalty(pres);
1743 }
1744 req
1745 }
1746
1747 /// Convert a [`ChatRole`] to a [`TextMessageRole`].
1748 #[cfg(any(
1749 target_os = "macos",
1750 target_os = "ios",
1751 target_os = "tvos",
1752 target_os = "visionos",
1753 target_os = "watchos",
1754 target_os = "windows",
1755 target_os = "linux",
1756 target_os = "android"
1757 ))]
1758 fn chat_role_to_mistral(role: &ChatRole) -> TextMessageRole {
1759 match role {
1760 ChatRole::System => TextMessageRole::System,
1761 ChatRole::User => TextMessageRole::User,
1762 ChatRole::Assistant => TextMessageRole::Assistant,
1763 }
1764 }
1765
1766 /// Truncate a string for log output, appending `"..."` if truncated.
1767 fn truncate_for_log(s: &str, max_len: usize) -> String {
1768 if s.len() > max_len {
1769 format!("{}...", &s[..max_len])
1770 } else {
1771 s.to_string()
1772 }
1773 }
1774
1775 // ═════════════════════════════════════════════════════════════════════════════
1776 // Unsupported-platform stub
1777 // ═════════════════════════════════════════════════════════════════════════════
1778
1779 /// Stub `ChatEngine` for platforms where mistral.rs is not available.
1780 ///
1781 /// All methods return errors or empty values so that code using the engine
1782 /// can compile on any target without `#[cfg]` at every call site.
1783 #[cfg(not(any(
1784 target_os = "macos",
1785 target_os = "ios",
1786 target_os = "tvos",
1787 target_os = "visionos",
1788 target_os = "watchos",
1789 target_os = "windows",
1790 target_os = "linux",
1791 target_os = "android"
1792 )))]
1793 pub struct ChatEngine;
1794
1795 #[cfg(not(any(
1796 target_os = "macos",
1797 target_os = "ios",
1798 target_os = "tvos",
1799 target_os = "visionos",
1800 target_os = "watchos",
1801 target_os = "windows",
1802 target_os = "linux",
1803 target_os = "android"
1804 )))]
1805 impl ChatEngine {
1806 pub fn new() -> Self {
1807 Self
1808 }
1809
1810 /// Stub mirror of the real `with_app_id`. The unsupported-platform engine
1811 /// ignores telemetry, so the app id is dropped.
1812 pub fn with_app_id(_onde_app_id: Option<String>) -> Self {
1813 Self
1814 }
1815
1816 pub async fn load_gguf_model(
1817 &self,
1818 _config: GgufModelConfig,
1819 _system_prompt: Option<String>,
1820 _sampling: Option<SamplingConfig>,
1821 ) -> Result<std::time::Duration, InferenceError> {
1822 Err(InferenceError::Other {
1823 reason: "LLM inference is not supported on this platform.".into(),
1824 })
1825 }
1826
1827 pub async fn unload_model(&self) -> Option<String> {
1828 None
1829 }
1830
1831 pub async fn is_loaded(&self) -> bool {
1832 false
1833 }
1834
1835 pub async fn info(&self) -> EngineInfo {
1836 EngineInfo {
1837 status: EngineStatus::Unloaded,
1838 model_name: None,
1839 approx_memory: None,
1840 history_length: 0u64,
1841 }
1842 }
1843
1844 pub async fn set_system_prompt(&self, _prompt: impl Into<String>) {}
1845
1846 pub async fn clear_system_prompt(&self) {}
1847
1848 pub async fn set_sampling(&self, _sampling: SamplingConfig) {}
1849
1850 pub async fn history(&self) -> Vec<ChatMessage> {
1851 Vec::new()
1852 }
1853
1854 pub async fn clear_history(&self) -> usize {
1855 0
1856 }
1857
1858 pub async fn push_history(&self, _message: ChatMessage) {}
1859
1860 pub async fn send_message(
1861 &self,
1862 _user_message: impl Into<String>,
1863 ) -> Result<InferenceResult, InferenceError> {
1864 Err(InferenceError::Other {
1865 reason: "LLM inference is not supported on this platform.".into(),
1866 })
1867 }
1868
1869 pub async fn generate(
1870 &self,
1871 _messages: Vec<ChatMessage>,
1872 _sampling: Option<SamplingConfig>,
1873 ) -> Result<InferenceResult, InferenceError> {
1874 Err(InferenceError::Other {
1875 reason: "LLM inference is not supported on this platform.".into(),
1876 })
1877 }
1878
1879 pub async fn stream_message(
1880 &self,
1881 _user_message: impl Into<String>,
1882 ) -> Result<tokio::sync::mpsc::Receiver<StreamChunk>, InferenceError> {
1883 Err(InferenceError::Other {
1884 reason: "LLM inference is not supported on this platform.".into(),
1885 })
1886 }
1887
1888 pub async fn send_message_with_tools(
1889 &self,
1890 _user_message: &str,
1891 _tools: &[ToolDefinition],
1892 ) -> Result<InferenceResult, InferenceError> {
1893 Err(InferenceError::Other {
1894 reason: "LLM inference is not supported on this platform.".into(),
1895 })
1896 }
1897
1898 pub async fn send_tool_results(
1899 &self,
1900 _results: Vec<ToolResult>,
1901 _next_tools: Option<&[ToolDefinition]>,
1902 ) -> Result<InferenceResult, InferenceError> {
1903 Err(InferenceError::Other {
1904 reason: "LLM inference is not supported on this platform.".into(),
1905 })
1906 }
1907 }
1908
1909 #[cfg(not(any(
1910 target_os = "macos",
1911 target_os = "ios",
1912 target_os = "tvos",
1913 target_os = "visionos",
1914 target_os = "watchos",
1915 target_os = "windows",
1916 target_os = "linux",
1917 target_os = "android"
1918 )))]
1919 impl Default for ChatEngine {
1920 fn default() -> Self {
1921 Self::new()
1922 }
1923 }
1924
1925 // ═════════════════════════════════════════════════════════════════════════════
1926 // Prebuilt model configs (convenience constructors)
1927 // ═════════════════════════════════════════════════════════════════════════════
1928
1929 /// DeepSeek Coder v1 Jinja chat template.
1930 ///
1931 /// TheBloke's GGUF files do not embed a chat template — without an explicit
1932 /// template, mistral.rs rejects chat requests with "model does not have a
1933 /// chat template". This template reproduces the official DeepSeek Coder v1
1934 /// `### Instruction: / ### Response:` format.
1935 const DEEPSEEK_CODER_CHAT_TEMPLATE: &str = include_str!("deepseek_coder_chat_template.txt");
1936
1937 /// Convenience constructors for common GGUF model configurations.
1938 ///
1939 /// These mirror the constants in [`super::models`] but return a ready-to-use
1940 /// [`GgufModelConfig`].
1941 impl GgufModelConfig {
1942 /// Whether the model's default chat template emits an extended reasoning
1943 /// block that needs a larger output budget than the general presets.
1944 fn uses_extended_thinking(&self) -> bool {
1945 matches!(
1946 self.model_id.as_str(),
1947 super::models::BARTOWSKI_QWEN3_0_6B_GGUF
1948 | super::models::BARTOWSKI_QWEN3_1_7B_GGUF
1949 | super::models::BARTOWSKI_QWEN3_4B_GGUF
1950 | super::models::BARTOWSKI_QWEN3_8B_GGUF
1951 | super::models::BARTOWSKI_QWEN3_14B_GGUF
1952 | super::models::BARTOWSKI_QWEN3_32B_GGUF
1953 | super::models::BARTOWSKI_QWEN3_4B_THINKING_2507_GGUF
1954 )
1955 }
1956
1957 /// Resolve a model-management ID to the same configuration used by the
1958 /// inference engine. Keeping this mapping in one place prevents download
1959 /// and load support from drifting apart.
1960 pub(crate) fn from_supported_model_id(model_id: &str) -> Option<Self> {
1961 match model_id {
1962 super::models::BARTOWSKI_QWEN25_0_5B_INSTRUCT_GGUF => Some(Self::qwen25_0_5b()),
1963 super::models::BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => Some(Self::qwen25_1_5b()),
1964 super::models::BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => Some(Self::qwen25_3b()),
1965 super::models::BARTOWSKI_QWEN25_CODER_1_5B_INSTRUCT_GGUF => {
1966 Some(Self::qwen25_coder_1_5b())
1967 }
1968 super::models::BARTOWSKI_QWEN25_CODER_3B_INSTRUCT_GGUF => Some(Self::qwen25_coder_3b()),
1969 super::models::BARTOWSKI_QWEN3_0_6B_GGUF => Some(Self::qwen3_0_6b()),
1970 super::models::BARTOWSKI_QWEN3_1_7B_GGUF => Some(Self::qwen3_1_7b()),
1971 super::models::BARTOWSKI_QWEN3_4B_GGUF => Some(Self::qwen3_4b()),
1972 super::models::BARTOWSKI_QWEN3_8B_GGUF => Some(Self::qwen3_8b()),
1973 super::models::BARTOWSKI_QWEN3_14B_GGUF => Some(Self::qwen3_14b()),
1974 super::models::BARTOWSKI_QWEN3_32B_GGUF => Some(Self::qwen3_32b()),
1975 super::models::BARTOWSKI_QWEN3_4B_INSTRUCT_2507_GGUF => {
1976 Some(Self::qwen3_4b_instruct_2507())
1977 }
1978 super::models::BARTOWSKI_QWEN3_4B_THINKING_2507_GGUF => {
1979 Some(Self::qwen3_4b_thinking_2507())
1980 }
1981 super::models::BARTOWSKI_QWEN3_30B_A3B_INSTRUCT_2507_GGUF => {
1982 Some(Self::qwen3_30b_a3b_instruct_2507())
1983 }
1984 super::models::BARTOWSKI_QWEN25_CODER_7B_INSTRUCT_GGUF => Some(Self::qwen25_coder_7b()),
1985 super::models::THEBLOKE_DEEPSEEK_CODER_6_7B_INSTRUCT_GGUF => {
1986 Some(Self::deepseek_coder_6_7b())
1987 }
1988 _ => None,
1989 }
1990 }
1991
1992 /// Qwen 2.5 0.5B Instruct (GGUF Q4_K_M), ~380 MB.
1993 ///
1994 /// Smallest option. Used on tvOS (Apple TV) where the memory limit is 2 GB.
1995 pub fn qwen25_0_5b() -> Self {
1996 Self {
1997 model_id: super::models::BARTOWSKI_QWEN25_0_5B_INSTRUCT_GGUF.into(),
1998 files: vec![super::models::QWEN25_0_5B_GGUF_FILE.into()],
1999 tok_model_id: if cfg!(target_os = "android") {
2000 Some(super::models::QWEN25_0_5B_TOK_MODEL_ID.into())
2001 } else {
2002 None
2003 },
2004 display_name: "Qwen 2.5 0.5B".into(),
2005 approx_memory: "~380 MB (GGUF Q4_K_M)".into(),
2006 chat_template: None,
2007 }
2008 }
2009
2010 /// Qwen 2.5 1.5B Instruct (GGUF Q4_K_M) — ~941 MB.
2011 ///
2012 /// Lightest option, ideal for iOS and memory-constrained Android devices.
2013 pub fn qwen25_1_5b() -> Self {
2014 Self {
2015 model_id: super::models::BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF.into(),
2016 files: vec![super::models::QWEN25_1_5B_GGUF_FILE.into()],
2017 tok_model_id: if cfg!(target_os = "android") {
2018 Some(super::models::QWEN25_1_5B_TOK_MODEL_ID.into())
2019 } else {
2020 None
2021 },
2022 display_name: "Qwen 2.5 1.5B".into(),
2023 approx_memory: "~941 MB (GGUF Q4_K_M)".into(),
2024 chat_template: None,
2025 }
2026 }
2027
2028 /// Qwen 2.5 3B Instruct (GGUF Q4_K_M) — ~1.93 GB.
2029 ///
2030 /// Balanced quality/size for macOS desktops and high-RAM Android devices.
2031 /// Not recommended for iOS (OOM on 8 GB devices).
2032 pub fn qwen25_3b() -> Self {
2033 Self {
2034 model_id: super::models::BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF.into(),
2035 files: vec![super::models::QWEN25_3B_GGUF_FILE.into()],
2036 tok_model_id: if cfg!(target_os = "android") {
2037 Some(super::models::QWEN25_3B_TOK_MODEL_ID.into())
2038 } else {
2039 None
2040 },
2041 display_name: "Qwen 2.5 3B".into(),
2042 approx_memory: "~1.93 GB (GGUF Q4_K_M)".into(),
2043 chat_template: None,
2044 }
2045 }
2046
2047 /// Qwen 2.5 Coder 1.5B Instruct (GGUF Q4_K_M) — ~941 MB.
2048 ///
2049 /// Dedicated coding model fine-tuned on 5.5T tokens of code and math data.
2050 /// Same `qwen2` GGUF architecture as the general 1.5B — loads via the
2051 /// existing `quantized_qwen.rs` path without any mistral.rs changes.
2052 /// Same memory footprint as the general 1.5B but dramatically better at
2053 /// code generation, explanation, and debugging.
2054 ///
2055 /// Preferred default for iOS, Android, and any coding-focused deployment.
2056 pub fn qwen25_coder_1_5b() -> Self {
2057 Self {
2058 model_id: super::models::BARTOWSKI_QWEN25_CODER_1_5B_INSTRUCT_GGUF.into(),
2059 files: vec![super::models::QWEN25_CODER_1_5B_GGUF_FILE.into()],
2060 tok_model_id: if cfg!(target_os = "android") {
2061 Some(super::models::QWEN25_CODER_1_5B_TOK_MODEL_ID.into())
2062 } else {
2063 None
2064 },
2065 display_name: "Qwen 2.5 Coder 1.5B".into(),
2066 approx_memory: "~941 MB (GGUF Q4_K_M)".into(),
2067 chat_template: None,
2068 }
2069 }
2070
2071 /// Qwen 2.5 Coder 3B Instruct (GGUF Q4_K_M) — ~1.93 GB.
2072 ///
2073 /// Best on-device coding quality for macOS desktop. Fine-tuned on 5.5T
2074 /// tokens of code and math. Same `qwen2` architecture as the general 3B.
2075 /// Not recommended for iOS due to memory constraints.
2076 pub fn qwen25_coder_3b() -> Self {
2077 Self {
2078 model_id: super::models::BARTOWSKI_QWEN25_CODER_3B_INSTRUCT_GGUF.into(),
2079 files: vec![super::models::QWEN25_CODER_3B_GGUF_FILE.into()],
2080 tok_model_id: if cfg!(target_os = "android") {
2081 Some(super::models::QWEN25_CODER_3B_TOK_MODEL_ID.into())
2082 } else {
2083 None
2084 },
2085 display_name: "Qwen 2.5 Coder 3B".into(),
2086 approx_memory: "~1.93 GB (GGUF Q4_K_M)".into(),
2087 chat_template: None,
2088 }
2089 }
2090
2091 /// Qwen 2.5 Coder 7B Instruct (GGUF Q4_K_M) — ~4.4 GB.
2092 ///
2093 /// Strong coding model with tool/function calling support via the Qwen2.5
2094 /// chat template. Uses the `qwen2` architecture. Requires 8+ GB RAM.
2095 pub fn qwen25_coder_7b() -> Self {
2096 Self {
2097 model_id: super::models::BARTOWSKI_QWEN25_CODER_7B_INSTRUCT_GGUF.into(),
2098 files: vec![super::models::QWEN25_CODER_7B_GGUF_FILE.into()],
2099 tok_model_id: if cfg!(target_os = "android") {
2100 Some(super::models::QWEN25_CODER_7B_TOK_MODEL_ID.into())
2101 } else {
2102 None
2103 },
2104 display_name: "Qwen 2.5 Coder 7B (Q4_K_M)".into(),
2105 approx_memory: "~4.4 GB".into(),
2106 chat_template: None,
2107 }
2108 }
2109
2110 /// Qwen 3 0.6B (GGUF Q4_K_M) — smallest Qwen 3 variant, ~0.5 GB.
2111 ///
2112 /// Lightest tool-capable Qwen 3 model; fits tvOS and the most
2113 /// memory-constrained mobile devices. Extended thinking mode — load with
2114 /// `max_tokens ≥ 4096`.
2115 pub fn qwen3_0_6b() -> Self {
2116 Self {
2117 model_id: super::models::BARTOWSKI_QWEN3_0_6B_GGUF.into(),
2118 files: vec![super::models::QWEN3_0_6B_GGUF_FILE.into()],
2119 tok_model_id: if cfg!(target_os = "android") {
2120 Some(super::models::QWEN3_0_6B_TOK_MODEL_ID.into())
2121 } else {
2122 None
2123 },
2124 display_name: "Qwen 3 0.6B (Q4_K_M)".into(),
2125 approx_memory: "~0.5 GB".into(),
2126 chat_template: None,
2127 }
2128 }
2129
2130 /// Qwen 3 1.7B (GGUF Q4_K_M) — lightweight tool-calling, ~1.3 GB.
2131 pub fn qwen3_1_7b() -> Self {
2132 Self {
2133 model_id: super::models::BARTOWSKI_QWEN3_1_7B_GGUF.into(),
2134 files: vec![super::models::QWEN3_1_7B_GGUF_FILE.into()],
2135 tok_model_id: if cfg!(target_os = "android") {
2136 Some(super::models::QWEN3_1_7B_TOK_MODEL_ID.into())
2137 } else {
2138 None
2139 },
2140 display_name: "Qwen 3 1.7B (Q4_K_M)".into(),
2141 approx_memory: "~1.3 GB".into(),
2142 chat_template: None,
2143 }
2144 }
2145
2146 /// Qwen 3 4B (GGUF Q4_K_M) — ~2.7 GB.
2147 ///
2148 /// Full OpenAI-compatible tool calling with extended thinking mode.
2149 /// Always load with `max_tokens ≥ 4096`; the `<think>…</think>` block can
2150 /// consume 300–400 tokens before the real response begins.
2151 pub fn qwen3_4b() -> Self {
2152 Self {
2153 model_id: super::models::BARTOWSKI_QWEN3_4B_GGUF.into(),
2154 files: vec![super::models::QWEN3_4B_GGUF_FILE.into()],
2155 tok_model_id: if cfg!(target_os = "android") {
2156 Some(super::models::QWEN3_4B_TOK_MODEL_ID.into())
2157 } else {
2158 None
2159 },
2160 display_name: "Qwen 3 4B (Q4_K_M)".into(),
2161 approx_memory: "~2.7 GB".into(),
2162 chat_template: None,
2163 }
2164 }
2165
2166 /// Strong tool-calling model with extended thinking. Best balance of
2167 /// quality and memory for macOS with 24+ GB RAM.
2168 pub fn qwen3_8b() -> Self {
2169 Self {
2170 model_id: super::models::BARTOWSKI_QWEN3_8B_GGUF.into(),
2171 files: vec![super::models::QWEN3_8B_GGUF_FILE.into()],
2172 tok_model_id: if cfg!(target_os = "android") {
2173 Some(super::models::QWEN3_8B_TOK_MODEL_ID.into())
2174 } else {
2175 None
2176 },
2177 display_name: "Qwen 3 8B (Q4_K_M)".into(),
2178 approx_memory: "~5 GB".into(),
2179 chat_template: None,
2180 }
2181 }
2182
2183 /// Qwen 3 14B (GGUF Q4_K_M) — ~8.4 GB.
2184 ///
2185 /// Strong reasoning and tool-calling model with extended thinking.
2186 /// Best all-around model for macOS with 16+ GB RAM.
2187 pub fn qwen3_14b() -> Self {
2188 Self {
2189 model_id: super::models::BARTOWSKI_QWEN3_14B_GGUF.into(),
2190 files: vec![super::models::QWEN3_14B_GGUF_FILE.into()],
2191 tok_model_id: if cfg!(target_os = "android") {
2192 Some(super::models::QWEN3_14B_TOK_MODEL_ID.into())
2193 } else {
2194 None
2195 },
2196 display_name: "Qwen 3 14B (Q4_K_M)".into(),
2197 approx_memory: "~8.4 GB".into(),
2198 chat_template: None,
2199 }
2200 }
2201
2202 /// Qwen 3 32B (GGUF Q4_K_M) — largest dense Qwen 3, ~19.8 GB.
2203 ///
2204 /// Highest-quality dense Qwen 3 variant. Requires a high-memory desktop
2205 /// (32+ GB RAM / unified memory). Extended thinking and full tool calling.
2206 pub fn qwen3_32b() -> Self {
2207 Self {
2208 model_id: super::models::BARTOWSKI_QWEN3_32B_GGUF.into(),
2209 files: vec![super::models::QWEN3_32B_GGUF_FILE.into()],
2210 tok_model_id: if cfg!(target_os = "android") {
2211 Some(super::models::QWEN3_32B_TOK_MODEL_ID.into())
2212 } else {
2213 None
2214 },
2215 display_name: "Qwen 3 32B (Q4_K_M)".into(),
2216 approx_memory: "~19.8 GB".into(),
2217 chat_template: None,
2218 }
2219 }
2220
2221 /// Qwen 3 4B Instruct 2507 (GGUF Q4_K_M) — latest non-thinking 4B, ~2.5 GB.
2222 ///
2223 /// Updated instruction-tuned checkpoint that does *not* emit a `<think>`
2224 /// block, making it faster and more predictable for chat and tool use.
2225 /// Recommended general-purpose Qwen 3 model.
2226 pub fn qwen3_4b_instruct_2507() -> Self {
2227 Self {
2228 model_id: super::models::BARTOWSKI_QWEN3_4B_INSTRUCT_2507_GGUF.into(),
2229 files: vec![super::models::QWEN3_4B_INSTRUCT_2507_GGUF_FILE.into()],
2230 tok_model_id: if cfg!(target_os = "android") {
2231 Some(super::models::QWEN3_4B_INSTRUCT_2507_TOK_MODEL_ID.into())
2232 } else {
2233 None
2234 },
2235 display_name: "Qwen 3 4B Instruct 2507 (Q4_K_M)".into(),
2236 approx_memory: "~2.5 GB".into(),
2237 chat_template: None,
2238 }
2239 }
2240
2241 /// Qwen 3 4B Thinking 2507 (GGUF Q4_K_M) — latest reasoning 4B, ~2.5 GB.
2242 ///
2243 /// Updated thinking-only checkpoint with stronger reasoning and tool-use
2244 /// accuracy. Always emits a `<think>` block — load with `max_tokens ≥ 4096`.
2245 pub fn qwen3_4b_thinking_2507() -> Self {
2246 Self {
2247 model_id: super::models::BARTOWSKI_QWEN3_4B_THINKING_2507_GGUF.into(),
2248 files: vec![super::models::QWEN3_4B_THINKING_2507_GGUF_FILE.into()],
2249 tok_model_id: if cfg!(target_os = "android") {
2250 Some(super::models::QWEN3_4B_THINKING_2507_TOK_MODEL_ID.into())
2251 } else {
2252 None
2253 },
2254 display_name: "Qwen 3 4B Thinking 2507 (Q4_K_M)".into(),
2255 approx_memory: "~2.5 GB".into(),
2256 chat_template: None,
2257 }
2258 }
2259
2260 /// Qwen 3 30B-A3B Instruct 2507 (GGUF Q4_K_M) — flagship MoE, ~18.6 GB.
2261 ///
2262 /// Mixture-of-experts model: 30B total parameters, ~3B active per token, so
2263 /// inference is far cheaper than a 30B dense model while quality rivals it.
2264 /// Loads via the `qwen3moe` GGUF architecture. Requires 32+ GB RAM.
2265 pub fn qwen3_30b_a3b_instruct_2507() -> Self {
2266 Self {
2267 model_id: super::models::BARTOWSKI_QWEN3_30B_A3B_INSTRUCT_2507_GGUF.into(),
2268 files: vec![super::models::QWEN3_30B_A3B_INSTRUCT_2507_GGUF_FILE.into()],
2269 tok_model_id: if cfg!(target_os = "android") {
2270 Some(super::models::QWEN3_30B_A3B_INSTRUCT_2507_TOK_MODEL_ID.into())
2271 } else {
2272 None
2273 },
2274 display_name: "Qwen 3 30B-A3B Instruct 2507 (Q4_K_M)".into(),
2275 approx_memory: "~18.6 GB".into(),
2276 chat_template: None,
2277 }
2278 }
2279
2280 /// DeepSeek Coder v1 6.7B Instruct (GGUF Q4_K_M) — ~3.8 GB.
2281 ///
2282 /// Strong code generation model using the `llama` GGUF architecture.
2283 /// Requires 8+ GB RAM; recommended for macOS desktops and Linux.
2284 pub fn deepseek_coder_6_7b() -> Self {
2285 Self {
2286 model_id: super::models::THEBLOKE_DEEPSEEK_CODER_6_7B_INSTRUCT_GGUF.into(),
2287 files: vec![super::models::DEEPSEEK_CODER_6_7B_GGUF_FILE.into()],
2288 tok_model_id: if cfg!(target_os = "android") {
2289 Some(super::models::DEEPSEEK_CODER_6_7B_TOK_MODEL_ID.into())
2290 } else {
2291 None
2292 },
2293 display_name: "DeepSeek Coder 6.7B (Q4_K_M)".into(),
2294 approx_memory: "~3.8 GB".into(),
2295 chat_template: Some(DEEPSEEK_CODER_CHAT_TEMPLATE.into()),
2296 }
2297 }
2298
2299 /// Return the platform-appropriate default **coding** model config.
2300 ///
2301 /// - iOS / Android → Qwen 2.5 Coder 1.5B (~941 MB, fits mobile memory budgets)
2302 /// - macOS / Windows / Linux → Qwen 2.5 Coder 3B (~1.93 GB, best desktop quality)
2303 ///
2304 /// Both are dedicated coding models (trained on 5.5T code + math tokens)
2305 /// and use the `qwen2` GGUF architecture supported by mistral.rs.
2306 pub fn platform_default() -> Self {
2307 if cfg!(target_os = "tvos") {
2308 Self::qwen25_0_5b()
2309 } else if cfg!(any(
2310 target_os = "ios",
2311 target_os = "visionos",
2312 target_os = "watchos",
2313 target_os = "android"
2314 )) {
2315 Self::qwen25_coder_1_5b()
2316 } else {
2317 Self::qwen25_coder_3b()
2318 }
2319 }
2320 }
2321
2322 // ═════════════════════════════════════════════════════════════════════════════
2323 // Tests
2324 // ═════════════════════════════════════════════════════════════════════════════
2325
2326 #[cfg(test)]
2327 mod tests {
2328 use super::*;
2329
2330 #[test]
2331 fn truncate_for_log_short() {
2332 assert_eq!(truncate_for_log("hello", 10), "hello");
2333 }
2334
2335 #[test]
2336 fn truncate_for_log_long() {
2337 let long = "a".repeat(200);
2338 let result = truncate_for_log(&long, 50);
2339 assert!(result.ends_with("..."));
2340 assert_eq!(result.len(), 53); // 50 chars + "..."
2341 }
2342
2343 #[test]
2344 fn gguf_model_config_qwen25_1_5b() {
2345 let cfg = GgufModelConfig::qwen25_1_5b();
2346 assert!(cfg.model_id.contains("1.5B"));
2347 assert_eq!(cfg.files.len(), 1);
2348 }
2349
2350 #[test]
2351 fn gguf_model_config_qwen25_3b() {
2352 let cfg = GgufModelConfig::qwen25_3b();
2353 assert!(cfg.model_id.contains("3B"));
2354 assert_eq!(cfg.files.len(), 1);
2355 }
2356
2357 #[test]
2358 fn gguf_model_config_qwen3_variants() {
2359 // Every Qwen 3 constructor must produce a valid, single-file config
2360 // whose repo ID is registered in the supported-models list.
2361 let configs = [
2362 GgufModelConfig::qwen3_0_6b(),
2363 GgufModelConfig::qwen3_1_7b(),
2364 GgufModelConfig::qwen3_4b(),
2365 GgufModelConfig::qwen3_8b(),
2366 GgufModelConfig::qwen3_14b(),
2367 GgufModelConfig::qwen3_32b(),
2368 GgufModelConfig::qwen3_4b_instruct_2507(),
2369 GgufModelConfig::qwen3_4b_thinking_2507(),
2370 GgufModelConfig::qwen3_30b_a3b_instruct_2507(),
2371 ];
2372 for cfg in &configs {
2373 assert!(
2374 cfg.model_id.contains("Qwen3"),
2375 "unexpected id: {}",
2376 cfg.model_id
2377 );
2378 assert_eq!(cfg.files.len(), 1);
2379 assert!(cfg.files[0].ends_with("-Q4_K_M.gguf"));
2380 assert!(
2381 super::super::models::SUPPORTED_MODELS.contains(&cfg.model_id.as_str()),
2382 "{} missing from SUPPORTED_MODELS",
2383 cfg.model_id
2384 );
2385 }
2386 }
2387
2388 #[test]
2389 fn every_supported_model_resolves_to_download_config() {
2390 for model_id in super::super::models::SUPPORTED_MODELS {
2391 let config = GgufModelConfig::from_supported_model_id(model_id)
2392 .unwrap_or_else(|| panic!("missing config for {model_id}"));
2393 assert_eq!(config.model_id, *model_id);
2394 assert_eq!(config.files.len(), 1);
2395 }
2396 }
2397
2398 #[test]
2399 fn qwen3_thinking_models_request_extended_sampling() {
2400 assert!(GgufModelConfig::qwen3_4b_thinking_2507().uses_extended_thinking());
2401 assert!(!GgufModelConfig::qwen3_4b_instruct_2507().uses_extended_thinking());
2402 }
2403
2404 #[test]
2405 fn gguf_model_config_platform_default() {
2406 let cfg = GgufModelConfig::platform_default();
2407 // On the test host (macOS or Linux), this should be the 3B model.
2408 // On iOS/Android CI it would be 1.5B. We just check it's valid.
2409 assert!(!cfg.model_id.is_empty());
2410 assert!(!cfg.files.is_empty());
2411 }
2412
2413 #[tokio::test]
2414 async fn engine_new_is_unloaded() {
2415 let engine = ChatEngine::new();
2416 assert!(!engine.is_loaded().await);
2417 let info = engine.info().await;
2418 assert_eq!(info.status, EngineStatus::Unloaded);
2419 assert_eq!(info.history_length, 0);
2420 }
2421
2422 #[tokio::test]
2423 async fn engine_send_without_model_errors() {
2424 let engine = ChatEngine::new();
2425 let result = engine.send_message("hello").await;
2426 assert!(result.is_err());
2427 match result.unwrap_err() {
2428 InferenceError::NoModelLoaded => {} // expected
2429 other => panic!("Expected NoModelLoaded, got: {:?}", other),
2430 }
2431 }
2432
2433 #[tokio::test]
2434 async fn engine_history_empty_when_no_model() {
2435 let engine = ChatEngine::new();
2436 assert!(engine.history().await.is_empty());
2437 assert_eq!(engine.clear_history().await, 0);
2438 }
2439
2440 #[tokio::test]
2441 async fn engine_unload_when_none() {
2442 let engine = ChatEngine::new();
2443 assert!(engine.unload_model().await.is_none());
2444 }
2445 }