@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
development
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 593 lines 20.9 KB
Raw
1 //! UniFFI-exported wrapper around [`ChatEngine`] for Swift / Kotlin consumers.
2 //!
3 //! [`ChatEngine`] uses Rust idioms (generics, `impl Into<String>`,
4 //! `tokio::sync::mpsc::Receiver`) that cannot cross the FFI boundary.
5 //! This module provides [`OndeChatEngine`], a `#[derive(uniffi::Object)]`
6 //! wrapper with concrete, FFI-safe method signatures.
7 //!
8 //! ## Swift usage
9 //!
10 //! ```swift
11 //! let engine = OndeChatEngine()
12 //!
13 //! // Load the platform-appropriate default model.
14 //! let elapsed = try await engine.loadDefaultModel(
15 //! systemPrompt: "You are a helpful assistant.",
16 //! sampling: nil
17 //! )
18 //! print("Model loaded in \(elapsed)s")
19 //!
20 //! // Multi-turn chat (history managed automatically).
21 //! let result = try await engine.sendMessage(message: "Hello!")
22 //! print(result.text)
23 //!
24 //! // Streaming (callback-based for FFI compatibility).
25 //! try await streamChatMessage(engine: engine, message: "Tell me a story.", listener: myListener)
26 //!
27 //! // One-shot generation (does NOT modify conversation history).
28 //! let enhanced = try await engine.generate(
29 //! messages: [ChatMessage(role: .user, content: "Expand: a cat in space")],
30 //! sampling: SamplingConfig(temperature: 0.0, topP: nil, topK: nil, minP: nil, maxTokens: 512, frequencyPenalty: nil, presencePenalty: nil)
31 //! )
32 //! print(enhanced.text)
33 //!
34 //! // Status & history.
35 //! let info = await engine.info()
36 //! let history = await engine.history()
37 //!
38 //! // Cleanup.
39 //! await engine.unloadModel()
40 //! ```
41
42 use std::sync::Arc;
43
44 use super::engine::ChatEngine;
45 use super::types::*;
46
47 // ═══════════════════════════════════════════════════════════════════════════
48 // Callback interface for streaming (used by free function, not Object method)
49 // ═══════════════════════════════════════════════════════════════════════════
50
51 /// Callback interface for receiving streaming token chunks.
52 ///
53 /// Implement this in Swift / Kotlin and pass it to the free function
54 /// [`stream_chat_message`].
55 ///
56 /// ```swift
57 /// class MyStreamHandler: StreamChunkListener {
58 /// func onChunk(chunk: StreamChunk) -> Bool {
59 /// print(chunk.delta, terminator: "")
60 /// return !chunk.done // return false to stop early
61 /// }
62 /// }
63 /// ```
64 #[uniffi::export(callback_interface)]
65 pub trait StreamChunkListener: Send + Sync {
66 /// Called for each token chunk during streaming inference.
67 ///
68 /// Return `true` to continue receiving chunks, or `false` to cancel
69 /// the stream early (the engine will still persist partial history).
70 fn on_chunk(&self, chunk: StreamChunk) -> bool;
71 }
72
73 // ═══════════════════════════════════════════════════════════════════════════
74 // OndeChatEngine — UniFFI Object
75 // ═══════════════════════════════════════════════════════════════════════════
76
77 /// On-device LLM chat inference engine — UniFFI-exported wrapper.
78 ///
79 /// This is the primary entry point for Swift / Kotlin consumers. It wraps
80 /// [`ChatEngine`] with FFI-safe method signatures (no generics, no Rust
81 /// channels, concrete `String` parameters).
82 ///
83 /// Construct with [`OndeChatEngine::new()`]. The engine starts with no
84 /// model loaded — call [`load_gguf_model`] or [`load_default_model`] first.
85 #[derive(uniffi::Object)]
86 pub struct OndeChatEngine {
87 inner: ChatEngine,
88 }
89
90 #[uniffi::export(async_runtime = "tokio")]
91 impl OndeChatEngine {
92 // ── Construction ─────────────────────────────────────────────────────
93
94 /// Create a new engine with no model loaded.
95 #[uniffi::constructor]
96 pub fn new() -> Arc<Self> {
97 Arc::new(Self {
98 inner: ChatEngine::new(),
99 })
100 }
101
102 // ── Model lifecycle ──────────────────────────────────────────────────
103
104 /// Load a GGUF model into the engine.
105 ///
106 /// If a model is already loaded it will be unloaded first.
107 ///
108 /// Returns the wall-clock loading time in seconds.
109 pub async fn load_gguf_model(
110 &self,
111 config: GgufModelConfig,
112 system_prompt: Option<String>,
113 sampling: Option<SamplingConfig>,
114 ) -> Result<f64, InferenceError> {
115 let elapsed = self
116 .inner
117 .load_gguf_model(config, system_prompt, sampling)
118 .await?;
119 Ok(elapsed.as_secs_f64())
120 }
121
122 /// Load a UQFF model into the engine.
123 ///
124 /// UQFF files are pre-quantised and load directly without runtime
125 /// conversion. `config.model_id` names the UQFF export (not the original
126 /// unquantised repo), and `config.files` should contain the first UQFF
127 /// shard or explicit shard filenames.
128 ///
129 /// Returns the wall-clock loading time in seconds.
130 pub async fn load_uqff_model(
131 &self,
132 config: UqffModelConfig,
133 system_prompt: Option<String>,
134 sampling: Option<SamplingConfig>,
135 ) -> Result<f64, InferenceError> {
136 let elapsed = self
137 .inner
138 .load_uqff_model(config, system_prompt, sampling)
139 .await?;
140 Ok(elapsed.as_secs_f64())
141 }
142
143 /// Load the platform-appropriate default model.
144 ///
145 /// - tvOS / iOS / Android → Qwen 2.5 1.5B (~941 MB)
146 /// - macOS / Windows / Linux → Qwen 2.5 3B (~1.93 GB)
147 ///
148 /// Returns the wall-clock loading time in seconds.
149 pub async fn load_default_model(
150 &self,
151 system_prompt: Option<String>,
152 sampling: Option<SamplingConfig>,
153 ) -> Result<f64, InferenceError> {
154 let config = GgufModelConfig::platform_default();
155 let elapsed = self
156 .inner
157 .load_gguf_model(config, system_prompt, sampling)
158 .await?;
159 Ok(elapsed.as_secs_f64())
160 }
161
162 /// Load the model assigned to this app via the Onde dashboard.
163 ///
164 /// The `app_id` and `app_secret` identify your app on
165 /// [ondeinference.com](https://ondeinference.com). The backend returns
166 /// the GGUF model configuration assigned to your app; if no model is
167 /// assigned (HTTP 404) the platform default is loaded instead.
168 ///
169 /// Returns the wall-clock loading time in seconds.
170 pub async fn load_assigned_model(
171 &self,
172 app_id: String,
173 app_secret: String,
174 system_prompt: Option<String>,
175 sampling: Option<SamplingConfig>,
176 ) -> Result<f64, InferenceError> {
177 let elapsed = self
178 .inner
179 .load_assigned_model(
180 smbcloud_gresiq_sdk::Environment::Production,
181 &app_id,
182 &app_secret,
183 system_prompt,
184 sampling,
185 )
186 .await?;
187 Ok(elapsed.as_secs_f64())
188 }
189
190 /// Unload the current model, freeing all memory.
191 ///
192 /// Returns the display name of the model that was unloaded, or `nil`
193 /// if no model was loaded.
194 pub async fn unload_model(&self) -> Option<String> {
195 self.inner.unload_model().await
196 }
197
198 /// Check whether a model is currently loaded.
199 pub async fn is_loaded(&self) -> bool {
200 self.inner.is_loaded().await
201 }
202
203 /// Get a snapshot of the engine's current state.
204 pub async fn info(&self) -> EngineInfo {
205 self.inner.info().await
206 }
207
208 // ── System prompt ────────────────────────────────────────────────────
209
210 /// Set or replace the system prompt.
211 pub async fn set_system_prompt(&self, prompt: String) {
212 self.inner.set_system_prompt(prompt).await;
213 }
214
215 /// Clear the system prompt.
216 pub async fn clear_system_prompt(&self) {
217 self.inner.clear_system_prompt().await;
218 }
219
220 // ── Sampling configuration ───────────────────────────────────────────
221
222 /// Replace the sampling configuration.
223 pub async fn set_sampling(&self, sampling: SamplingConfig) {
224 self.inner.set_sampling(sampling).await;
225 }
226
227 // ── Conversation history ─────────────────────────────────────────────
228
229 /// Get a copy of the full conversation history.
230 pub async fn history(&self) -> Vec<ChatMessage> {
231 self.inner.history().await
232 }
233
234 /// Clear the conversation history but keep the model loaded.
235 ///
236 /// Returns the number of turns that were removed.
237 pub async fn clear_history(&self) -> u64 {
238 self.inner.clear_history().await as u64
239 }
240
241 /// Append a message to history without running inference.
242 pub async fn push_history(&self, message: ChatMessage) {
243 self.inner.push_history(message).await;
244 }
245
246 // ── Non-streaming inference ──────────────────────────────────────────
247
248 /// Send a user message and receive a complete assistant reply.
249 ///
250 /// The user message and assistant reply are automatically appended to
251 /// the conversation history on success.
252 pub async fn send_message(&self, message: String) -> Result<InferenceResult, InferenceError> {
253 self.inner.send_message(message).await
254 }
255
256 /// Run inference on an explicit list of messages WITHOUT modifying the
257 /// engine's internal history.
258 ///
259 /// Useful for one-shot prompts (e.g. prompt enhancement).
260 pub async fn generate(
261 &self,
262 messages: Vec<ChatMessage>,
263 sampling: Option<SamplingConfig>,
264 ) -> Result<InferenceResult, InferenceError> {
265 self.inner.generate(messages, sampling).await
266 }
267 }
268
269 // ═══════════════════════════════════════════════════════════════════════════
270 // Free function for streaming — UniFFI callback_interface works in free fns
271 // ═══════════════════════════════════════════════════════════════════════════
272
273 /// Stream a chat message through the engine, delivering token chunks to
274 /// the `listener` callback.
275 ///
276 /// This is a **free function** (not an Object method) because UniFFI 0.31
277 /// requires `callback_interface` parameters to be passed to free functions
278 /// rather than Object methods.
279 ///
280 /// The `listener.on_chunk()` callback is called for each token delta.
281 /// Return `true` to continue, or `false` to cancel early.
282 ///
283 /// The user message and assembled reply are automatically appended to the
284 /// engine's conversation history when streaming completes.
285 ///
286 /// ```swift
287 /// class MyHandler: StreamChunkListener {
288 /// func onChunk(chunk: StreamChunk) -> Bool {
289 /// print(chunk.delta, terminator: "")
290 /// return !chunk.done
291 /// }
292 /// }
293 ///
294 /// try await streamChatMessage(
295 /// engine: engine,
296 /// message: "Tell me a story.",
297 /// listener: MyHandler()
298 /// )
299 /// ```
300 #[uniffi::export(async_runtime = "tokio")]
301 pub async fn stream_chat_message(
302 engine: Arc<OndeChatEngine>,
303 message: String,
304 listener: Box<dyn StreamChunkListener>,
305 ) -> Result<(), InferenceError> {
306 let mut rx = engine.inner.stream_message(message).await?;
307
308 while let Some(chunk) = rx.recv().await {
309 let done = chunk.done;
310 let should_continue = listener.on_chunk(chunk);
311 if done || !should_continue {
312 break;
313 }
314 }
315
316 Ok(())
317 }
318
319 // ═══════════════════════════════════════════════════════════════════════════
320 // Free functions — convenience exports for Swift / Kotlin
321 // ═══════════════════════════════════════════════════════════════════════════
322
323 /// Return the platform-appropriate default GGUF model configuration.
324 ///
325 /// - tvOS / iOS / Android → Qwen 2.5 1.5B (~941 MB)
326 /// - macOS / Windows / Linux → Qwen 2.5 3B (~1.93 GB)
327 #[uniffi::export]
328 pub fn default_model_config() -> GgufModelConfig {
329 GgufModelConfig::platform_default()
330 }
331
332 /// Create a generic UQFF model configuration.
333 ///
334 /// `model_id` is the HuggingFace repo or local directory of the UQFF export
335 /// (the shards plus `residual.safetensors`, `config.json`, and the tokenizer),
336 /// not the original unquantised repo. `files` accepts a first shard
337 /// (`q4k-0.uqff`), explicit shards, or mistral.rs shorthands such as `q4k`
338 /// and `4`.
339 #[uniffi::export]
340 pub fn uqff_model_config(
341 model_id: String,
342 files: Vec<String>,
343 display_name: String,
344 approx_memory: String,
345 chat_template: Option<String>,
346 ) -> UqffModelConfig {
347 UqffModelConfig {
348 model_id,
349 files,
350 display_name,
351 approx_memory,
352 chat_template,
353 }
354 }
355
356 /// Return the Qwen 2.5 0.5B GGUF model configuration (~380 MB).
357 #[uniffi::export]
358 pub fn qwen25_0_5b_config() -> GgufModelConfig {
359 GgufModelConfig::qwen25_0_5b()
360 }
361
362 /// Return the Qwen 2.5 1.5B GGUF model configuration (~941 MB).
363 #[uniffi::export]
364 pub fn qwen25_1_5b_config() -> GgufModelConfig {
365 GgufModelConfig::qwen25_1_5b()
366 }
367
368 /// Return the Qwen 2.5 3B GGUF model configuration (~1.93 GB).
369 #[uniffi::export]
370 pub fn qwen25_3b_config() -> GgufModelConfig {
371 GgufModelConfig::qwen25_3b()
372 }
373
374 // ── Qwen 3 family ────────────────────────────────────────────────────────────
375
376 /// Return the Qwen 3 0.6B GGUF model configuration (~0.5 GB).
377 #[uniffi::export]
378 pub fn qwen3_0_6b_config() -> GgufModelConfig {
379 GgufModelConfig::qwen3_0_6b()
380 }
381
382 /// Return the Qwen 3 1.7B GGUF model configuration (~1.3 GB).
383 #[uniffi::export]
384 pub fn qwen3_1_7b_config() -> GgufModelConfig {
385 GgufModelConfig::qwen3_1_7b()
386 }
387
388 /// Return the Qwen 3 4B GGUF model configuration (~2.7 GB).
389 #[uniffi::export]
390 pub fn qwen3_4b_config() -> GgufModelConfig {
391 GgufModelConfig::qwen3_4b()
392 }
393
394 /// Return the Qwen 3 8B GGUF model configuration (~5 GB).
395 #[uniffi::export]
396 pub fn qwen3_8b_config() -> GgufModelConfig {
397 GgufModelConfig::qwen3_8b()
398 }
399
400 /// Return the Qwen 3 14B GGUF model configuration (~8.4 GB).
401 #[uniffi::export]
402 pub fn qwen3_14b_config() -> GgufModelConfig {
403 GgufModelConfig::qwen3_14b()
404 }
405
406 /// Return the Qwen 3 32B GGUF model configuration (~19.8 GB).
407 #[uniffi::export]
408 pub fn qwen3_32b_config() -> GgufModelConfig {
409 GgufModelConfig::qwen3_32b()
410 }
411
412 /// Return the Qwen 3 4B Instruct 2507 GGUF model configuration (~2.5 GB).
413 ///
414 /// Latest non-thinking 4B checkpoint — recommended general-purpose Qwen 3 model.
415 #[uniffi::export]
416 pub fn qwen3_4b_instruct_2507_config() -> GgufModelConfig {
417 GgufModelConfig::qwen3_4b_instruct_2507()
418 }
419
420 /// Return the Qwen 3 4B Thinking 2507 GGUF model configuration (~2.5 GB).
421 ///
422 /// Latest reasoning-focused 4B checkpoint; load with `max_tokens ≥ 4096`.
423 #[uniffi::export]
424 pub fn qwen3_4b_thinking_2507_config() -> GgufModelConfig {
425 GgufModelConfig::qwen3_4b_thinking_2507()
426 }
427
428 /// Return the Qwen 3 30B-A3B Instruct 2507 (MoE) GGUF model configuration (~18.6 GB).
429 #[uniffi::export]
430 pub fn qwen3_30b_a3b_instruct_2507_config() -> GgufModelConfig {
431 GgufModelConfig::qwen3_30b_a3b_instruct_2507()
432 }
433
434 /// Return default sampling parameters for creative chat.
435 #[uniffi::export]
436 pub fn default_sampling_config() -> SamplingConfig {
437 SamplingConfig::default()
438 }
439
440 /// Return deterministic (greedy) sampling parameters.
441 #[uniffi::export]
442 pub fn deterministic_sampling_config() -> SamplingConfig {
443 SamplingConfig::deterministic()
444 }
445
446 /// Return conservative sampling parameters for mobile devices.
447 #[uniffi::export]
448 pub fn mobile_sampling_config() -> SamplingConfig {
449 SamplingConfig::mobile()
450 }
451
452 /// Create a system message.
453 #[uniffi::export]
454 pub fn system_message(content: String) -> ChatMessage {
455 ChatMessage::system(content)
456 }
457
458 /// Create a user message.
459 #[uniffi::export]
460 pub fn user_message(content: String) -> ChatMessage {
461 ChatMessage::user(content)
462 }
463
464 /// Create an assistant message.
465 #[uniffi::export]
466 pub fn assistant_message(content: String) -> ChatMessage {
467 ChatMessage::assistant(content)
468 }
469
470 /// Configure the HF cache directory for sandboxed platforms.
471 ///
472 /// Call this before any model-loading method (`load_gguf_model`,
473 /// `load_default_model`, etc.) on iOS, tvOS, and Android where the
474 /// default `~/.cache/huggingface` path is inaccessible.
475 ///
476 /// `path` should be a writable app data directory (e.g. the App Group
477 /// container or Application Support directory on Apple platforms).
478 ///
479 /// This sets the `HF_HOME` and `HF_HUB_CACHE` environment variables
480 /// and creates the directory structure.
481 #[uniffi::export]
482 pub fn configure_cache_dir(path: String) {
483 let base = std::path::PathBuf::from(&path);
484 let hf_home = base.join("models");
485 let hf_hub_cache = hf_home.join("hub");
486
487 if let Err(e) = std::fs::create_dir_all(&hf_hub_cache) {
488 log::warn!(
489 "configure_cache_dir: could not create {}: {}",
490 hf_hub_cache.display(),
491 e
492 );
493 }
494
495 std::env::set_var("HF_HOME", &hf_home);
496 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
497
498 log::info!(
499 "configure_cache_dir: HF_HOME={}, HF_HUB_CACHE={}",
500 hf_home.display(),
501 hf_hub_cache.display()
502 );
503 }
504
505 // ═══════════════════════════════════════════════════════════════════════════
506 // Tests
507 // ═══════════════════════════════════════════════════════════════════════════
508
509 #[cfg(test)]
510 mod tests {
511 use super::*;
512
513 #[test]
514 fn free_functions_return_valid_configs() {
515 let cfg = default_model_config();
516 assert!(!cfg.model_id.is_empty());
517
518 let cfg = qwen25_1_5b_config();
519 assert!(cfg.model_id.contains("1.5B"));
520
521 let cfg = qwen25_3b_config();
522 assert!(cfg.model_id.contains("3B"));
523 }
524
525 #[test]
526 fn free_functions_return_valid_sampling() {
527 let s = default_sampling_config();
528 assert_eq!(s.temperature, Some(0.7));
529
530 let s = deterministic_sampling_config();
531 assert_eq!(s.temperature, Some(0.0));
532
533 let s = mobile_sampling_config();
534 assert_eq!(s.max_tokens, Some(128));
535 }
536
537 #[test]
538 fn message_helpers() {
539 let m = system_message("You are helpful.".into());
540 assert_eq!(m.role, ChatRole::System);
541 assert_eq!(m.content, "You are helpful.");
542
543 let m = user_message("Hello".into());
544 assert_eq!(m.role, ChatRole::User);
545
546 let m = assistant_message("Hi!".into());
547 assert_eq!(m.role, ChatRole::Assistant);
548 }
549
550 #[test]
551 fn configure_cache_dir_sets_env() {
552 let tmp = std::env::temp_dir().join("onde-test-cache-dir");
553 configure_cache_dir(tmp.to_string_lossy().to_string());
554
555 let hf_home = std::env::var("HF_HOME").unwrap();
556 assert!(hf_home.contains("models"));
557
558 let hf_hub = std::env::var("HF_HUB_CACHE").unwrap();
559 assert!(hf_hub.contains("hub"));
560
561 // Cleanup
562 let _ = std::fs::remove_dir_all(&tmp);
563 }
564
565 #[tokio::test]
566 async fn onde_chat_engine_new_is_unloaded() {
567 let engine = OndeChatEngine::new();
568 assert!(!engine.is_loaded().await);
569 let info = engine.info().await;
570 assert_eq!(info.status, EngineStatus::Unloaded);
571 assert_eq!(info.history_length, 0);
572 }
573
574 #[tokio::test]
575 async fn onde_chat_engine_send_without_model_errors() {
576 let engine = OndeChatEngine::new();
577 let result = engine.send_message("hello".into()).await;
578 assert!(result.is_err());
579 }
580
581 #[tokio::test]
582 async fn onde_chat_engine_history_empty() {
583 let engine = OndeChatEngine::new();
584 assert!(engine.history().await.is_empty());
585 assert_eq!(engine.clear_history().await, 0);
586 }
587
588 #[tokio::test]
589 async fn onde_chat_engine_unload_when_none() {
590 let engine = OndeChatEngine::new();
591 assert!(engine.unload_model().await.is_none());
592 }
593 }