@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 548 lines 19.5 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 the platform-appropriate default model.
123 ///
124 /// - tvOS / iOS / Android → Qwen 2.5 1.5B (~941 MB)
125 /// - macOS / Windows / Linux → Qwen 2.5 3B (~1.93 GB)
126 ///
127 /// Returns the wall-clock loading time in seconds.
128 pub async fn load_default_model(
129 &self,
130 system_prompt: Option<String>,
131 sampling: Option<SamplingConfig>,
132 ) -> Result<f64, InferenceError> {
133 let config = GgufModelConfig::platform_default();
134 let elapsed = self
135 .inner
136 .load_gguf_model(config, system_prompt, sampling)
137 .await?;
138 Ok(elapsed.as_secs_f64())
139 }
140
141 /// Load the model assigned to this app via the Onde dashboard.
142 ///
143 /// The `app_id` and `app_secret` identify your app on
144 /// [ondeinference.com](https://ondeinference.com). The backend returns
145 /// the GGUF model configuration assigned to your app; if no model is
146 /// assigned (HTTP 404) the platform default is loaded instead.
147 ///
148 /// Returns the wall-clock loading time in seconds.
149 pub async fn load_assigned_model(
150 &self,
151 app_id: String,
152 app_secret: String,
153 system_prompt: Option<String>,
154 sampling: Option<SamplingConfig>,
155 ) -> Result<f64, InferenceError> {
156 let elapsed = self
157 .inner
158 .load_assigned_model(
159 smbcloud_gresiq_sdk::Environment::Production,
160 &app_id,
161 &app_secret,
162 system_prompt,
163 sampling,
164 )
165 .await?;
166 Ok(elapsed.as_secs_f64())
167 }
168
169 /// Unload the current model, freeing all memory.
170 ///
171 /// Returns the display name of the model that was unloaded, or `nil`
172 /// if no model was loaded.
173 pub async fn unload_model(&self) -> Option<String> {
174 self.inner.unload_model().await
175 }
176
177 /// Check whether a model is currently loaded.
178 pub async fn is_loaded(&self) -> bool {
179 self.inner.is_loaded().await
180 }
181
182 /// Get a snapshot of the engine's current state.
183 pub async fn info(&self) -> EngineInfo {
184 self.inner.info().await
185 }
186
187 // ── System prompt ────────────────────────────────────────────────────
188
189 /// Set or replace the system prompt.
190 pub async fn set_system_prompt(&self, prompt: String) {
191 self.inner.set_system_prompt(prompt).await;
192 }
193
194 /// Clear the system prompt.
195 pub async fn clear_system_prompt(&self) {
196 self.inner.clear_system_prompt().await;
197 }
198
199 // ── Sampling configuration ───────────────────────────────────────────
200
201 /// Replace the sampling configuration.
202 pub async fn set_sampling(&self, sampling: SamplingConfig) {
203 self.inner.set_sampling(sampling).await;
204 }
205
206 // ── Conversation history ─────────────────────────────────────────────
207
208 /// Get a copy of the full conversation history.
209 pub async fn history(&self) -> Vec<ChatMessage> {
210 self.inner.history().await
211 }
212
213 /// Clear the conversation history but keep the model loaded.
214 ///
215 /// Returns the number of turns that were removed.
216 pub async fn clear_history(&self) -> u64 {
217 self.inner.clear_history().await as u64
218 }
219
220 /// Append a message to history without running inference.
221 pub async fn push_history(&self, message: ChatMessage) {
222 self.inner.push_history(message).await;
223 }
224
225 // ── Non-streaming inference ──────────────────────────────────────────
226
227 /// Send a user message and receive a complete assistant reply.
228 ///
229 /// The user message and assistant reply are automatically appended to
230 /// the conversation history on success.
231 pub async fn send_message(&self, message: String) -> Result<InferenceResult, InferenceError> {
232 self.inner.send_message(message).await
233 }
234
235 /// Run inference on an explicit list of messages WITHOUT modifying the
236 /// engine's internal history.
237 ///
238 /// Useful for one-shot prompts (e.g. prompt enhancement).
239 pub async fn generate(
240 &self,
241 messages: Vec<ChatMessage>,
242 sampling: Option<SamplingConfig>,
243 ) -> Result<InferenceResult, InferenceError> {
244 self.inner.generate(messages, sampling).await
245 }
246 }
247
248 // ═══════════════════════════════════════════════════════════════════════════
249 // Free function for streaming — UniFFI callback_interface works in free fns
250 // ═══════════════════════════════════════════════════════════════════════════
251
252 /// Stream a chat message through the engine, delivering token chunks to
253 /// the `listener` callback.
254 ///
255 /// This is a **free function** (not an Object method) because UniFFI 0.31
256 /// requires `callback_interface` parameters to be passed to free functions
257 /// rather than Object methods.
258 ///
259 /// The `listener.on_chunk()` callback is called for each token delta.
260 /// Return `true` to continue, or `false` to cancel early.
261 ///
262 /// The user message and assembled reply are automatically appended to the
263 /// engine's conversation history when streaming completes.
264 ///
265 /// ```swift
266 /// class MyHandler: StreamChunkListener {
267 /// func onChunk(chunk: StreamChunk) -> Bool {
268 /// print(chunk.delta, terminator: "")
269 /// return !chunk.done
270 /// }
271 /// }
272 ///
273 /// try await streamChatMessage(
274 /// engine: engine,
275 /// message: "Tell me a story.",
276 /// listener: MyHandler()
277 /// )
278 /// ```
279 #[uniffi::export(async_runtime = "tokio")]
280 pub async fn stream_chat_message(
281 engine: Arc<OndeChatEngine>,
282 message: String,
283 listener: Box<dyn StreamChunkListener>,
284 ) -> Result<(), InferenceError> {
285 let mut rx = engine.inner.stream_message(message).await?;
286
287 while let Some(chunk) = rx.recv().await {
288 let done = chunk.done;
289 let should_continue = listener.on_chunk(chunk);
290 if done || !should_continue {
291 break;
292 }
293 }
294
295 Ok(())
296 }
297
298 // ═══════════════════════════════════════════════════════════════════════════
299 // Free functions — convenience exports for Swift / Kotlin
300 // ═══════════════════════════════════════════════════════════════════════════
301
302 /// Return the platform-appropriate default GGUF model configuration.
303 ///
304 /// - tvOS / iOS / Android → Qwen 2.5 1.5B (~941 MB)
305 /// - macOS / Windows / Linux → Qwen 2.5 3B (~1.93 GB)
306 #[uniffi::export]
307 pub fn default_model_config() -> GgufModelConfig {
308 GgufModelConfig::platform_default()
309 }
310
311 /// Return the Qwen 2.5 0.5B GGUF model configuration (~380 MB).
312 #[uniffi::export]
313 pub fn qwen25_0_5b_config() -> GgufModelConfig {
314 GgufModelConfig::qwen25_0_5b()
315 }
316
317 /// Return the Qwen 2.5 1.5B GGUF model configuration (~941 MB).
318 #[uniffi::export]
319 pub fn qwen25_1_5b_config() -> GgufModelConfig {
320 GgufModelConfig::qwen25_1_5b()
321 }
322
323 /// Return the Qwen 2.5 3B GGUF model configuration (~1.93 GB).
324 #[uniffi::export]
325 pub fn qwen25_3b_config() -> GgufModelConfig {
326 GgufModelConfig::qwen25_3b()
327 }
328
329 // ── Qwen 3 family ────────────────────────────────────────────────────────────
330
331 /// Return the Qwen 3 0.6B GGUF model configuration (~0.5 GB).
332 #[uniffi::export]
333 pub fn qwen3_0_6b_config() -> GgufModelConfig {
334 GgufModelConfig::qwen3_0_6b()
335 }
336
337 /// Return the Qwen 3 1.7B GGUF model configuration (~1.3 GB).
338 #[uniffi::export]
339 pub fn qwen3_1_7b_config() -> GgufModelConfig {
340 GgufModelConfig::qwen3_1_7b()
341 }
342
343 /// Return the Qwen 3 4B GGUF model configuration (~2.7 GB).
344 #[uniffi::export]
345 pub fn qwen3_4b_config() -> GgufModelConfig {
346 GgufModelConfig::qwen3_4b()
347 }
348
349 /// Return the Qwen 3 8B GGUF model configuration (~5 GB).
350 #[uniffi::export]
351 pub fn qwen3_8b_config() -> GgufModelConfig {
352 GgufModelConfig::qwen3_8b()
353 }
354
355 /// Return the Qwen 3 14B GGUF model configuration (~8.4 GB).
356 #[uniffi::export]
357 pub fn qwen3_14b_config() -> GgufModelConfig {
358 GgufModelConfig::qwen3_14b()
359 }
360
361 /// Return the Qwen 3 32B GGUF model configuration (~19.8 GB).
362 #[uniffi::export]
363 pub fn qwen3_32b_config() -> GgufModelConfig {
364 GgufModelConfig::qwen3_32b()
365 }
366
367 /// Return the Qwen 3 4B Instruct 2507 GGUF model configuration (~2.5 GB).
368 ///
369 /// Latest non-thinking 4B checkpoint — recommended general-purpose Qwen 3 model.
370 #[uniffi::export]
371 pub fn qwen3_4b_instruct_2507_config() -> GgufModelConfig {
372 GgufModelConfig::qwen3_4b_instruct_2507()
373 }
374
375 /// Return the Qwen 3 4B Thinking 2507 GGUF model configuration (~2.5 GB).
376 ///
377 /// Latest reasoning-focused 4B checkpoint; load with `max_tokens ≥ 4096`.
378 #[uniffi::export]
379 pub fn qwen3_4b_thinking_2507_config() -> GgufModelConfig {
380 GgufModelConfig::qwen3_4b_thinking_2507()
381 }
382
383 /// Return the Qwen 3 30B-A3B Instruct 2507 (MoE) GGUF model configuration (~18.6 GB).
384 #[uniffi::export]
385 pub fn qwen3_30b_a3b_instruct_2507_config() -> GgufModelConfig {
386 GgufModelConfig::qwen3_30b_a3b_instruct_2507()
387 }
388
389 /// Return default sampling parameters for creative chat.
390 #[uniffi::export]
391 pub fn default_sampling_config() -> SamplingConfig {
392 SamplingConfig::default()
393 }
394
395 /// Return deterministic (greedy) sampling parameters.
396 #[uniffi::export]
397 pub fn deterministic_sampling_config() -> SamplingConfig {
398 SamplingConfig::deterministic()
399 }
400
401 /// Return conservative sampling parameters for mobile devices.
402 #[uniffi::export]
403 pub fn mobile_sampling_config() -> SamplingConfig {
404 SamplingConfig::mobile()
405 }
406
407 /// Create a system message.
408 #[uniffi::export]
409 pub fn system_message(content: String) -> ChatMessage {
410 ChatMessage::system(content)
411 }
412
413 /// Create a user message.
414 #[uniffi::export]
415 pub fn user_message(content: String) -> ChatMessage {
416 ChatMessage::user(content)
417 }
418
419 /// Create an assistant message.
420 #[uniffi::export]
421 pub fn assistant_message(content: String) -> ChatMessage {
422 ChatMessage::assistant(content)
423 }
424
425 /// Configure the HF cache directory for sandboxed platforms.
426 ///
427 /// Call this before any model-loading method (`load_gguf_model`,
428 /// `load_default_model`, etc.) on iOS, tvOS, and Android where the
429 /// default `~/.cache/huggingface` path is inaccessible.
430 ///
431 /// `path` should be a writable app data directory (e.g. the App Group
432 /// container or Application Support directory on Apple platforms).
433 ///
434 /// This sets the `HF_HOME` and `HF_HUB_CACHE` environment variables
435 /// and creates the directory structure.
436 #[uniffi::export]
437 pub fn configure_cache_dir(path: String) {
438 let base = std::path::PathBuf::from(&path);
439 let hf_home = base.join("models");
440 let hf_hub_cache = hf_home.join("hub");
441
442 if let Err(e) = std::fs::create_dir_all(&hf_hub_cache) {
443 log::warn!(
444 "configure_cache_dir: could not create {}: {}",
445 hf_hub_cache.display(),
446 e
447 );
448 }
449
450 std::env::set_var("HF_HOME", &hf_home);
451 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
452
453 log::info!(
454 "configure_cache_dir: HF_HOME={}, HF_HUB_CACHE={}",
455 hf_home.display(),
456 hf_hub_cache.display()
457 );
458 }
459
460 // ═══════════════════════════════════════════════════════════════════════════
461 // Tests
462 // ═══════════════════════════════════════════════════════════════════════════
463
464 #[cfg(test)]
465 mod tests {
466 use super::*;
467
468 #[test]
469 fn free_functions_return_valid_configs() {
470 let cfg = default_model_config();
471 assert!(!cfg.model_id.is_empty());
472
473 let cfg = qwen25_1_5b_config();
474 assert!(cfg.model_id.contains("1.5B"));
475
476 let cfg = qwen25_3b_config();
477 assert!(cfg.model_id.contains("3B"));
478 }
479
480 #[test]
481 fn free_functions_return_valid_sampling() {
482 let s = default_sampling_config();
483 assert_eq!(s.temperature, Some(0.7));
484
485 let s = deterministic_sampling_config();
486 assert_eq!(s.temperature, Some(0.0));
487
488 let s = mobile_sampling_config();
489 assert_eq!(s.max_tokens, Some(128));
490 }
491
492 #[test]
493 fn message_helpers() {
494 let m = system_message("You are helpful.".into());
495 assert_eq!(m.role, ChatRole::System);
496 assert_eq!(m.content, "You are helpful.");
497
498 let m = user_message("Hello".into());
499 assert_eq!(m.role, ChatRole::User);
500
501 let m = assistant_message("Hi!".into());
502 assert_eq!(m.role, ChatRole::Assistant);
503 }
504
505 #[test]
506 fn configure_cache_dir_sets_env() {
507 let tmp = std::env::temp_dir().join("onde-test-cache-dir");
508 configure_cache_dir(tmp.to_string_lossy().to_string());
509
510 let hf_home = std::env::var("HF_HOME").unwrap();
511 assert!(hf_home.contains("models"));
512
513 let hf_hub = std::env::var("HF_HUB_CACHE").unwrap();
514 assert!(hf_hub.contains("hub"));
515
516 // Cleanup
517 let _ = std::fs::remove_dir_all(&tmp);
518 }
519
520 #[tokio::test]
521 async fn onde_chat_engine_new_is_unloaded() {
522 let engine = OndeChatEngine::new();
523 assert!(!engine.is_loaded().await);
524 let info = engine.info().await;
525 assert_eq!(info.status, EngineStatus::Unloaded);
526 assert_eq!(info.history_length, 0);
527 }
528
529 #[tokio::test]
530 async fn onde_chat_engine_send_without_model_errors() {
531 let engine = OndeChatEngine::new();
532 let result = engine.send_message("hello".into()).await;
533 assert!(result.is_err());
534 }
535
536 #[tokio::test]
537 async fn onde_chat_engine_history_empty() {
538 let engine = OndeChatEngine::new();
539 assert!(engine.history().await.is_empty());
540 assert_eq!(engine.clear_history().await, 0);
541 }
542
543 #[tokio::test]
544 async fn onde_chat_engine_unload_when_none() {
545 let engine = OndeChatEngine::new();
546 assert!(engine.unload_model().await.is_none());
547 }
548 }