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