@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 488 lines 17.6 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 /// Return default sampling parameters for creative chat.
330 #[uniffi::export]
331 pub fn default_sampling_config() -> SamplingConfig {
332 SamplingConfig::default()
333 }
334
335 /// Return deterministic (greedy) sampling parameters.
336 #[uniffi::export]
337 pub fn deterministic_sampling_config() -> SamplingConfig {
338 SamplingConfig::deterministic()
339 }
340
341 /// Return conservative sampling parameters for mobile devices.
342 #[uniffi::export]
343 pub fn mobile_sampling_config() -> SamplingConfig {
344 SamplingConfig::mobile()
345 }
346
347 /// Create a system message.
348 #[uniffi::export]
349 pub fn system_message(content: String) -> ChatMessage {
350 ChatMessage::system(content)
351 }
352
353 /// Create a user message.
354 #[uniffi::export]
355 pub fn user_message(content: String) -> ChatMessage {
356 ChatMessage::user(content)
357 }
358
359 /// Create an assistant message.
360 #[uniffi::export]
361 pub fn assistant_message(content: String) -> ChatMessage {
362 ChatMessage::assistant(content)
363 }
364
365 /// Configure the HF cache directory for sandboxed platforms.
366 ///
367 /// Call this before any model-loading method (`load_gguf_model`,
368 /// `load_default_model`, etc.) on iOS, tvOS, and Android where the
369 /// default `~/.cache/huggingface` path is inaccessible.
370 ///
371 /// `path` should be a writable app data directory (e.g. the App Group
372 /// container or Application Support directory on Apple platforms).
373 ///
374 /// This sets the `HF_HOME` and `HF_HUB_CACHE` environment variables
375 /// and creates the directory structure.
376 #[uniffi::export]
377 pub fn configure_cache_dir(path: String) {
378 let base = std::path::PathBuf::from(&path);
379 let hf_home = base.join("models");
380 let hf_hub_cache = hf_home.join("hub");
381
382 if let Err(e) = std::fs::create_dir_all(&hf_hub_cache) {
383 log::warn!(
384 "configure_cache_dir: could not create {}: {}",
385 hf_hub_cache.display(),
386 e
387 );
388 }
389
390 std::env::set_var("HF_HOME", &hf_home);
391 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
392
393 log::info!(
394 "configure_cache_dir: HF_HOME={}, HF_HUB_CACHE={}",
395 hf_home.display(),
396 hf_hub_cache.display()
397 );
398 }
399
400 // ═══════════════════════════════════════════════════════════════════════════
401 // Tests
402 // ═══════════════════════════════════════════════════════════════════════════
403
404 #[cfg(test)]
405 mod tests {
406 use super::*;
407
408 #[test]
409 fn free_functions_return_valid_configs() {
410 let cfg = default_model_config();
411 assert!(!cfg.model_id.is_empty());
412
413 let cfg = qwen25_1_5b_config();
414 assert!(cfg.model_id.contains("1.5B"));
415
416 let cfg = qwen25_3b_config();
417 assert!(cfg.model_id.contains("3B"));
418 }
419
420 #[test]
421 fn free_functions_return_valid_sampling() {
422 let s = default_sampling_config();
423 assert_eq!(s.temperature, Some(0.7));
424
425 let s = deterministic_sampling_config();
426 assert_eq!(s.temperature, Some(0.0));
427
428 let s = mobile_sampling_config();
429 assert_eq!(s.max_tokens, Some(128));
430 }
431
432 #[test]
433 fn message_helpers() {
434 let m = system_message("You are helpful.".into());
435 assert_eq!(m.role, ChatRole::System);
436 assert_eq!(m.content, "You are helpful.");
437
438 let m = user_message("Hello".into());
439 assert_eq!(m.role, ChatRole::User);
440
441 let m = assistant_message("Hi!".into());
442 assert_eq!(m.role, ChatRole::Assistant);
443 }
444
445 #[test]
446 fn configure_cache_dir_sets_env() {
447 let tmp = std::env::temp_dir().join("onde-test-cache-dir");
448 configure_cache_dir(tmp.to_string_lossy().to_string());
449
450 let hf_home = std::env::var("HF_HOME").unwrap();
451 assert!(hf_home.contains("models"));
452
453 let hf_hub = std::env::var("HF_HUB_CACHE").unwrap();
454 assert!(hf_hub.contains("hub"));
455
456 // Cleanup
457 let _ = std::fs::remove_dir_all(&tmp);
458 }
459
460 #[tokio::test]
461 async fn onde_chat_engine_new_is_unloaded() {
462 let engine = OndeChatEngine::new();
463 assert!(!engine.is_loaded().await);
464 let info = engine.info().await;
465 assert_eq!(info.status, EngineStatus::Unloaded);
466 assert_eq!(info.history_length, 0);
467 }
468
469 #[tokio::test]
470 async fn onde_chat_engine_send_without_model_errors() {
471 let engine = OndeChatEngine::new();
472 let result = engine.send_message("hello".into()).await;
473 assert!(result.is_err());
474 }
475
476 #[tokio::test]
477 async fn onde_chat_engine_history_empty() {
478 let engine = OndeChatEngine::new();
479 assert!(engine.history().await.is_empty());
480 assert_eq!(engine.clear_history().await, 0);
481 }
482
483 #[tokio::test]
484 async fn onde_chat_engine_unload_when_none() {
485 let engine = OndeChatEngine::new();
486 assert!(engine.unload_model().await.is_none());
487 }
488 }