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