@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
copilot/assess-inference-engine-complexity
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 432 lines 15.8 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 1.5B GGUF model configuration (~941 MB).
312 #[uniffi::export]
313 pub fn qwen25_1_5b_config() -> GgufModelConfig {
314 GgufModelConfig::qwen25_1_5b()
315 }
316
317 /// Return the Qwen 2.5 3B GGUF model configuration (~1.93 GB).
318 #[uniffi::export]
319 pub fn qwen25_3b_config() -> GgufModelConfig {
320 GgufModelConfig::qwen25_3b()
321 }
322
323 /// Return default sampling parameters for creative chat.
324 #[uniffi::export]
325 pub fn default_sampling_config() -> SamplingConfig {
326 SamplingConfig::default()
327 }
328
329 /// Return deterministic (greedy) sampling parameters.
330 #[uniffi::export]
331 pub fn deterministic_sampling_config() -> SamplingConfig {
332 SamplingConfig::deterministic()
333 }
334
335 /// Return conservative sampling parameters for mobile devices.
336 #[uniffi::export]
337 pub fn mobile_sampling_config() -> SamplingConfig {
338 SamplingConfig::mobile()
339 }
340
341 /// Create a system message.
342 #[uniffi::export]
343 pub fn system_message(content: String) -> ChatMessage {
344 ChatMessage::system(content)
345 }
346
347 /// Create a user message.
348 #[uniffi::export]
349 pub fn user_message(content: String) -> ChatMessage {
350 ChatMessage::user(content)
351 }
352
353 /// Create an assistant message.
354 #[uniffi::export]
355 pub fn assistant_message(content: String) -> ChatMessage {
356 ChatMessage::assistant(content)
357 }
358
359 // ═══════════════════════════════════════════════════════════════════════════
360 // Tests
361 // ═══════════════════════════════════════════════════════════════════════════
362
363 #[cfg(test)]
364 mod tests {
365 use super::*;
366
367 #[test]
368 fn free_functions_return_valid_configs() {
369 let cfg = default_model_config();
370 assert!(!cfg.model_id.is_empty());
371
372 let cfg = qwen25_1_5b_config();
373 assert!(cfg.model_id.contains("1.5B"));
374
375 let cfg = qwen25_3b_config();
376 assert!(cfg.model_id.contains("3B"));
377 }
378
379 #[test]
380 fn free_functions_return_valid_sampling() {
381 let s = default_sampling_config();
382 assert_eq!(s.temperature, Some(0.7));
383
384 let s = deterministic_sampling_config();
385 assert_eq!(s.temperature, Some(0.0));
386
387 let s = mobile_sampling_config();
388 assert_eq!(s.max_tokens, Some(128));
389 }
390
391 #[test]
392 fn message_helpers() {
393 let m = system_message("You are helpful.".into());
394 assert_eq!(m.role, ChatRole::System);
395 assert_eq!(m.content, "You are helpful.");
396
397 let m = user_message("Hello".into());
398 assert_eq!(m.role, ChatRole::User);
399
400 let m = assistant_message("Hi!".into());
401 assert_eq!(m.role, ChatRole::Assistant);
402 }
403
404 #[tokio::test]
405 async fn onde_chat_engine_new_is_unloaded() {
406 let engine = OndeChatEngine::new();
407 assert!(!engine.is_loaded().await);
408 let info = engine.info().await;
409 assert_eq!(info.status, EngineStatus::Unloaded);
410 assert_eq!(info.history_length, 0);
411 }
412
413 #[tokio::test]
414 async fn onde_chat_engine_send_without_model_errors() {
415 let engine = OndeChatEngine::new();
416 let result = engine.send_message("hello".into()).await;
417 assert!(result.is_err());
418 }
419
420 #[tokio::test]
421 async fn onde_chat_engine_history_empty() {
422 let engine = OndeChatEngine::new();
423 assert!(engine.history().await.is_empty());
424 assert_eq!(engine.clear_history().await, 0);
425 }
426
427 #[tokio::test]
428 async fn onde_chat_engine_unload_when_none() {
429 let engine = OndeChatEngine::new();
430 assert!(engine.unload_model().await.is_none());
431 }
432 }