@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 552 lines 20.9 KB
Raw
1 //! Shared types for the on-device LLM inference engine.
2 //!
3 //! These types are intentionally framework-agnostic — they do **not** depend on
4 //! any UI framework. They carry UniFFI derive annotations so that
5 //! UniFFI bindgen can generate Swift / Kotlin bindings automatically.
6 //!
7 //! ## UniFFI compatibility notes
8 //!
9 //! - `usize` is **not** supported by UniFFI → we use `u64` / `u32` instead.
10 //! - Enums use `#[derive(uniffi::Enum)]`.
11 //! - Plain data structs use `#[derive(uniffi::Record)]`.
12 //! - Error enums use `#[derive(uniffi::Error)]`.
13
14 use serde::{Deserialize, Serialize};
15
16 // ── Message role ─────────────────────────────────────────────────────────────
17
18 /// Role of a message in a chat conversation.
19 ///
20 /// Mirrors the standard OpenAI / mistral.rs role taxonomy but is decoupled
21 /// from `mistralrs::TextMessageRole` so that callers don't need a direct
22 /// dependency on mistral.rs.
23 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, uniffi::Enum)]
24 #[serde(rename_all = "lowercase")]
25 pub enum ChatRole {
26 System,
27 User,
28 Assistant,
29 }
30
31 impl std::fmt::Display for ChatRole {
32 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 match self {
34 ChatRole::System => write!(f, "system"),
35 ChatRole::User => write!(f, "user"),
36 ChatRole::Assistant => write!(f, "assistant"),
37 }
38 }
39 }
40
41 // ── Chat message ─────────────────────────────────────────────────────────────
42
43 /// A single message in a conversation.
44 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
45 pub struct ChatMessage {
46 pub role: ChatRole,
47 pub content: String,
48 }
49
50 impl ChatMessage {
51 pub fn system(content: impl Into<String>) -> Self {
52 Self {
53 role: ChatRole::System,
54 content: content.into(),
55 }
56 }
57
58 pub fn user(content: impl Into<String>) -> Self {
59 Self {
60 role: ChatRole::User,
61 content: content.into(),
62 }
63 }
64
65 pub fn assistant(content: impl Into<String>) -> Self {
66 Self {
67 role: ChatRole::Assistant,
68 content: content.into(),
69 }
70 }
71 }
72
73 // ── Sampling configuration ───────────────────────────────────────────────────
74
75 /// Sampling parameters for text generation.
76 ///
77 /// All fields are optional — `None` means "use the engine default".
78 ///
79 /// Note: `usize` is replaced with `u64` for UniFFI compatibility.
80 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
81 pub struct SamplingConfig {
82 /// Sampling temperature (higher = more random). Typical range: 0.0–2.0.
83 pub temperature: Option<f64>,
84 /// Nucleus (top-p) sampling threshold. Typical value: 0.9–0.95.
85 pub top_p: Option<f64>,
86 /// Top-k sampling limit.
87 pub top_k: Option<u64>,
88 /// Min-p sampling threshold.
89 pub min_p: Option<f64>,
90 /// Maximum number of tokens to generate.
91 pub max_tokens: Option<u64>,
92 /// Frequency penalty (penalise tokens proportional to occurrence count).
93 pub frequency_penalty: Option<f32>,
94 /// Presence penalty (penalise tokens that appeared at all).
95 pub presence_penalty: Option<f32>,
96 }
97
98 impl Default for SamplingConfig {
99 /// Sensible creative-chat defaults.
100 fn default() -> Self {
101 Self {
102 temperature: Some(0.7),
103 top_p: Some(0.95),
104 top_k: None,
105 min_p: None,
106 max_tokens: Some(512),
107 frequency_penalty: None,
108 presence_penalty: None,
109 }
110 }
111 }
112
113 impl SamplingConfig {
114 /// Deterministic sampling (temperature = 0, greedy decoding).
115 pub fn deterministic() -> Self {
116 Self {
117 temperature: Some(0.0),
118 top_p: None,
119 top_k: None,
120 min_p: None,
121 max_tokens: Some(512),
122 frequency_penalty: None,
123 presence_penalty: None,
124 }
125 }
126
127 /// Conservative mobile defaults — lower max_tokens for faster response on
128 /// constrained devices (CPU-only ARM SoCs doing ~1-3 tok/sec).
129 pub fn mobile() -> Self {
130 Self {
131 temperature: Some(0.7),
132 top_p: Some(0.95),
133 top_k: None,
134 min_p: None,
135 max_tokens: Some(128),
136 frequency_penalty: None,
137 presence_penalty: None,
138 }
139 }
140
141 /// Sampling config tuned for coding assistants.
142 ///
143 /// Key differences from the general-purpose default:
144 ///
145 /// - **`frequency_penalty = 1.15`** — the primary fix for the repeated-token
146 /// loop you saw with small GGUF models. Penalises each token proportional
147 /// to how many times it has already appeared in the output, making it
148 /// progressively harder to emit the same word again. 1.15 is aggressive
149 /// enough to kill echo loops without hurting code token repetition
150 /// (identifiers, keywords) at typical response lengths.
151 ///
152 /// - **`presence_penalty = 0.1`** — a light nudge away from any token that
153 /// has appeared at all, encouraging lexical variety across the response.
154 ///
155 /// - **`temperature = 0.2`** — low randomness keeps generated code
156 /// deterministic and syntactically correct. High temperature causes
157 /// hallucinated variable names and broken syntax in code blocks.
158 ///
159 /// - **`top_p = 0.95`** — nucleus sampling retained to avoid degenerate
160 /// greedy outputs while still keeping the distribution tight.
161 ///
162 /// - **`max_tokens = 1024`** — doubled vs. the chat default to allow
163 /// complete function / class explanations without truncation.
164 pub fn coding() -> Self {
165 Self {
166 temperature: Some(0.2),
167 top_p: Some(0.95),
168 top_k: None,
169 min_p: None,
170 max_tokens: Some(1024),
171 frequency_penalty: Some(1.15),
172 presence_penalty: Some(0.1),
173 }
174 }
175
176 /// Like [`coding`](Self::coding) but capped at 256 tokens for mobile
177 /// devices where latency matters more than completeness.
178 pub fn coding_mobile() -> Self {
179 Self {
180 max_tokens: Some(256),
181 ..Self::coding()
182 }
183 }
184 }
185
186 // ── GGUF model configuration ─────────────────────────────────────────────────
187
188 /// Configuration for loading a GGUF model via mistral.rs.
189 ///
190 /// This is the primary model format for on-device mobile/desktop inference
191 /// because GGUF files are pre-quantized and load at their compressed size,
192 /// avoiding the large memory spike of in-situ quantization.
193 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
194 pub struct GgufModelConfig {
195 /// HuggingFace model repository ID, e.g. `"bartowski/Qwen2.5-1.5B-Instruct-GGUF"`.
196 pub model_id: String,
197 /// GGUF filename(s) within the repository, e.g. `["Qwen2.5-1.5B-Instruct-Q4_K_M.gguf"]`.
198 pub files: Vec<String>,
199 /// Optional: explicit tokenizer model ID (required on Android where GGUF
200 /// embedded tokenizers are not supported by the candle backend).
201 pub tok_model_id: Option<String>,
202 /// Human-friendly display name, e.g. `"Qwen 2.5 1.5B"`.
203 pub display_name: String,
204 /// Approximate memory footprint description, e.g. `"~941 MB (GGUF Q4_K_M)"`.
205 pub approx_memory: String,
206 /// Optional Jinja chat template string. When set, this overrides the
207 /// template embedded in the GGUF file (if any). Required for older GGUF
208 /// files that don't embed a chat template (e.g. TheBloke models).
209 pub chat_template: Option<String>,
210 }
211
212 // ── ISQ model configuration ──────────────────────────────────────────────────
213
214 /// Configuration for loading an ISQ (in-situ quantised) model via mistral.rs
215 /// `TextModelBuilder`.
216 ///
217 /// Unlike GGUF models, ISQ models are loaded directly from a HuggingFace
218 /// safetensors repo and quantised in-situ on the device. They require more
219 /// RAM during the load phase but support a wider range of architectures
220 /// (e.g. DeepSeek-Coder-V2-Lite, Qwen2.5-Coder-7B).
221 ///
222 /// **macOS-only for now** — ISQ with Metal acceleration requires the `metal`
223 /// feature on mistral.rs. The struct is still defined on all platforms so
224 /// that types.rs compiles everywhere without `cfg` gates on the record
225 /// (UniFFI requires all record types to be reachable without feature flags).
226 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
227 pub struct IsqModelConfig {
228 /// HuggingFace model repository ID, e.g. `"Qwen/Qwen2.5-Coder-7B-Instruct"`.
229 pub model_id: String,
230 /// Number of quantisation bits. `4` → Q4K (AFQ4 on Metal), `8` → Q8_0 (AFQ8 on Metal).
231 /// Any value outside `{4, 8}` defaults to 4-bit.
232 pub isq_bits: u8,
233 /// Human-friendly display name, e.g. `"Qwen 2.5 Coder 7B"`.
234 pub display_name: String,
235 /// Approximate memory footprint description, e.g. `"~4.5 GB (Q4K, Metal)"`.
236 pub approx_memory: String,
237 }
238
239 impl IsqModelConfig {
240 /// Qwen 2.5 Coder 7B Instruct — ISQ 4-bit, Metal-accelerated (~4.5 GB).
241 ///
242 /// Best macOS desktop coding quality. Loaded from the official HF repo
243 /// and quantised in-situ; requires ~8 GB RAM during the load phase.
244 pub fn qwen25_coder_7b_isq4() -> Self {
245 Self {
246 model_id: super::models::QWEN25_CODER_7B_INSTRUCT.into(),
247 isq_bits: 4,
248 display_name: "Qwen 2.5 Coder 7B (ISQ 4-bit)".into(),
249 approx_memory: "~4.5 GB (ISQ Q4K, Metal)".into(),
250 }
251 }
252
253 /// Qwen 2.5 Coder 7B Instruct — ISQ 8-bit, Metal-accelerated (~9 GB).
254 ///
255 /// Higher quality than the 4-bit variant; requires ~12 GB RAM on load.
256 /// Suitable for Macs with 16 GB+ unified memory.
257 pub fn qwen25_coder_7b_isq8() -> Self {
258 Self {
259 model_id: super::models::QWEN25_CODER_7B_INSTRUCT.into(),
260 isq_bits: 8,
261 display_name: "Qwen 2.5 Coder 7B (ISQ 8-bit)".into(),
262 approx_memory: "~9 GB (ISQ Q8K, Metal)".into(),
263 }
264 }
265 }
266
267 // ── Inference result ─────────────────────────────────────────────────────────
268
269 /// The result of a chat inference request.
270 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
271 pub struct InferenceResult {
272 /// The generated assistant reply text.
273 pub text: String,
274 /// Wall-clock duration of the inference in seconds.
275 pub duration_secs: f64,
276 /// Human-readable duration string (e.g. `"2m 3.1s"` or `"4.5s"`).
277 pub duration_display: String,
278 /// Finish reason reported by the model (e.g. `"stop"`, `"length"`).
279 pub finish_reason: String,
280 /// Tool calls requested by the model (empty when no tools were invoked).
281 pub tool_calls: Vec<ToolCallInfo>,
282 }
283
284 // ── Tool calling ─────────────────────────────────────────────────────────────
285
286 /// A tool definition passed to the model so it knows which tools are available.
287 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
288 pub struct ToolDefinition {
289 /// Tool name (e.g. `"read_file"`).
290 pub name: String,
291 /// Human-readable description of what the tool does.
292 pub description: String,
293 /// JSON Schema string describing the tool's parameters.
294 pub parameters_schema: String,
295 }
296
297 /// Information about a single tool call requested by the model.
298 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
299 pub struct ToolCallInfo {
300 /// Unique identifier for this tool call (used to correlate results).
301 pub id: String,
302 /// The name of the function the model wants to invoke.
303 pub function_name: String,
304 /// JSON-encoded arguments for the function.
305 pub arguments: String,
306 }
307
308 /// The result of executing a tool, sent back to the model.
309 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
310 pub struct ToolResult {
311 /// The tool call ID this result corresponds to.
312 pub tool_call_id: String,
313 /// The output produced by executing the tool.
314 pub content: String,
315 }
316
317 // ── Streaming chunk ──────────────────────────────────────────────────────────
318
319 /// A single streaming token chunk from the inference engine.
320 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
321 pub struct StreamChunk {
322 /// The token text delta (may be empty for the final chunk).
323 pub delta: String,
324 /// Whether this is the final chunk in the stream.
325 pub done: bool,
326 /// Finish reason (only set on the final chunk).
327 pub finish_reason: Option<String>,
328 }
329
330 // ── Engine status ────────────────────────────────────────────────────────────
331
332 /// Lifecycle status of the inference engine.
333 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, uniffi::Enum)]
334 #[serde(rename_all = "snake_case")]
335 pub enum EngineStatus {
336 /// No model loaded.
337 Unloaded,
338 /// Model is being downloaded / loaded into memory.
339 Loading,
340 /// Model is loaded and ready to accept requests.
341 Ready,
342 /// Model is currently running inference.
343 Generating,
344 /// An error occurred (model may or may not still be loaded).
345 Error,
346 }
347
348 impl std::fmt::Display for EngineStatus {
349 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
350 match self {
351 EngineStatus::Unloaded => write!(f, "unloaded"),
352 EngineStatus::Loading => write!(f, "loading"),
353 EngineStatus::Ready => write!(f, "ready"),
354 EngineStatus::Generating => write!(f, "generating"),
355 EngineStatus::Error => write!(f, "error"),
356 }
357 }
358 }
359
360 // ── Engine info snapshot ─────────────────────────────────────────────────────
361
362 /// A point-in-time snapshot of the engine's state, suitable for status UIs.
363 ///
364 /// Note: `history_length` is `u64` (not `usize`) for UniFFI compatibility.
365 #[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
366 pub struct EngineInfo {
367 pub status: EngineStatus,
368 /// Display name of the currently loaded model, if any.
369 pub model_name: Option<String>,
370 /// Approximate memory footprint description, if a model is loaded.
371 pub approx_memory: Option<String>,
372 /// Number of conversation turns in the current history.
373 pub history_length: u64,
374 }
375
376 // ── Error type ───────────────────────────────────────────────────────────────
377
378 /// Errors that can occur during inference engine operations.
379 ///
380 /// Exposed to Swift / Kotlin as a sealed class / enum via `uniffi::Error`.
381 #[derive(Debug, thiserror::Error, uniffi::Error)]
382 pub enum InferenceError {
383 #[error("No model loaded — call `load_model` first")]
384 NoModelLoaded,
385
386 #[error("Model is already loaded: {model_name}")]
387 AlreadyLoaded { model_name: String },
388
389 #[error("Failed to build model: {reason}")]
390 ModelBuild { reason: String },
391
392 #[error("Inference failed: {reason}")]
393 Inference { reason: String },
394
395 #[error("Model loading was cancelled")]
396 Cancelled,
397
398 #[error("{reason}")]
399 Other { reason: String },
400 }
401
402 // ── Duration formatting helper ───────────────────────────────────────────────
403
404 /// Format a `std::time::Duration` as `Xm Ys` or just `Ys` when under a minute.
405 pub fn format_duration(d: std::time::Duration) -> String {
406 let total_secs = d.as_secs_f64();
407 let mins = (total_secs / 60.0).floor() as u64;
408 let secs = total_secs - (mins as f64 * 60.0);
409 if mins > 0 {
410 format!("{}m {:.1}s", mins, secs)
411 } else {
412 format!("{:.1}s", secs)
413 }
414 }
415
416 // ── Tool calling types (Rust-only) ──────────────────────────────────────────
417 //
418 // These types are intentionally NOT annotated with `uniffi::Record` /
419 // `uniffi::Enum` — they are for Rust consumers only. The UniFFI surface
420 // area (Swift / Kotlin) remains unchanged.
421 // (`ToolDefinition`, `ToolCallInfo`, and `ToolResult` are defined above
422 // with UniFFI derives and shared across both layers.)
423
424 /// A tool call requested by the model.
425 #[derive(Debug, Clone)]
426 pub struct ToolCallRequest {
427 /// Unique identifier for this tool call (used to correlate results).
428 pub id: String,
429 /// Name of the function the model wants to invoke.
430 pub function_name: String,
431 /// JSON string of the function arguments.
432 pub arguments: String,
433 }
434
435 /// Result from a tool-aware inference call (Rust-only, not UniFFI-exported).
436 #[derive(Debug, Clone)]
437 pub struct ToolAwareResult {
438 /// Text content of the assistant reply (may be empty when the model
439 /// decides to call tools instead of responding with text).
440 pub text: String,
441 /// Tool calls requested by the model. Empty when the model responds
442 /// with text only.
443 pub tool_calls: Vec<ToolCallRequest>,
444 /// Wall-clock inference duration in seconds.
445 pub duration_secs: f64,
446 /// Human-readable duration string (e.g. `"1.3s"`).
447 pub duration_display: String,
448 /// Finish reason reported by the model (e.g. `"stop"`, `"tool_calls"`).
449 pub finish_reason: String,
450 }
451
452 #[cfg(test)]
453 mod tests {
454 use super::*;
455
456 #[test]
457 fn chat_message_constructors() {
458 let sys = ChatMessage::system("You are helpful.");
459 assert_eq!(sys.role, ChatRole::System);
460 assert_eq!(sys.content, "You are helpful.");
461
462 let user = ChatMessage::user("Hello");
463 assert_eq!(user.role, ChatRole::User);
464
465 let asst = ChatMessage::assistant("Hi there!");
466 assert_eq!(asst.role, ChatRole::Assistant);
467 }
468
469 #[test]
470 fn chat_role_display() {
471 assert_eq!(ChatRole::System.to_string(), "system");
472 assert_eq!(ChatRole::User.to_string(), "user");
473 assert_eq!(ChatRole::Assistant.to_string(), "assistant");
474 }
475
476 #[test]
477 fn chat_role_serde_rename() {
478 // Verify that #[serde(rename_all = "lowercase")] is applied by
479 // checking the Display impl matches the expected serialised form.
480 // (Full serde_json round-trip is tested in integration tests that
481 // already depend on serde_json.)
482 assert_eq!(format!("{}", ChatRole::System), "system");
483 assert_eq!(format!("{}", ChatRole::User), "user");
484 assert_eq!(format!("{}", ChatRole::Assistant), "assistant");
485 }
486
487 #[test]
488 fn sampling_config_defaults() {
489 let cfg = SamplingConfig::default();
490 assert_eq!(cfg.temperature, Some(0.7));
491 assert_eq!(cfg.top_p, Some(0.95));
492 assert_eq!(cfg.max_tokens, Some(512));
493 }
494
495 #[test]
496 fn sampling_config_deterministic() {
497 let cfg = SamplingConfig::deterministic();
498 assert_eq!(cfg.temperature, Some(0.0));
499 assert!(cfg.top_p.is_none());
500 }
501
502 #[test]
503 fn sampling_config_mobile() {
504 let cfg = SamplingConfig::mobile();
505 assert_eq!(cfg.max_tokens, Some(128));
506 }
507
508 #[test]
509 fn engine_status_display() {
510 assert_eq!(EngineStatus::Ready.to_string(), "ready");
511 assert_eq!(EngineStatus::Generating.to_string(), "generating");
512 }
513
514 #[test]
515 fn format_duration_under_minute() {
516 let d = std::time::Duration::from_secs_f64(4.567);
517 assert_eq!(format_duration(d), "4.6s");
518 }
519
520 #[test]
521 fn format_duration_over_minute() {
522 let d = std::time::Duration::from_secs_f64(125.3);
523 assert_eq!(format_duration(d), "2m 5.3s");
524 }
525
526 #[test]
527 fn inference_error_display() {
528 let err = InferenceError::NoModelLoaded;
529 assert_eq!(err.to_string(), "No model loaded — call `load_model` first");
530
531 let err = InferenceError::ModelBuild {
532 reason: "out of memory".into(),
533 };
534 assert_eq!(err.to_string(), "Failed to build model: out of memory");
535
536 let err = InferenceError::AlreadyLoaded {
537 model_name: "Qwen 2.5".into(),
538 };
539 assert_eq!(err.to_string(), "Model is already loaded: Qwen 2.5");
540 }
541
542 #[test]
543 fn engine_info_history_is_u64() {
544 let info = EngineInfo {
545 status: EngineStatus::Ready,
546 model_name: Some("Test".into()),
547 approx_memory: None,
548 history_length: 42,
549 };
550 assert_eq!(info.history_length, 42u64);
551 }
552 }