@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 272 lines 10.1 KB
Raw
1 //! Ruby native extension for the Onde on-device inference engine.
2 //!
3 //! This crate uses [magnus](https://github.com/matsadler/magnus) to expose
4 //! Onde's HuggingFace cache management, supported model metadata, and
5 //! inference types to Ruby as the `Onde` module.
6 //!
7 //! ## Ruby usage
8 //!
9 //! ```ruby
10 //! require "onde"
11 //!
12 //! # List models cached locally on disk.
13 //! response = Onde.list_local_models
14 //! response["models"].each do |m|
15 //! puts "#{m["model_id"]} — #{m["size_display"]}"
16 //! end
17 //!
18 //! # List all supported models (with download status).
19 //! Onde.list_supported_models["models"].each do |m|
20 //! status = m["is_downloaded"] ? "✓" : "✗"
21 //! puts "[#{status}] #{m["name"]} (#{m["org"]}) — #{m["expected_size_display"]}"
22 //! end
23 //!
24 //! # Delete a cached model.
25 //! Onde.delete_model("bartowski/Qwen2.5-1.5B-Instruct-GGUF")
26 //!
27 //! # Inspect supported model IDs.
28 //! Onde::SUPPORTED_MODELS # => ["black-forest-labs/FLUX.1-schnell", ...]
29 //!
30 //! # Access model metadata.
31 //! Onde.model_info("bartowski/Qwen2.5-1.5B-Instruct-GGUF")
32 //! # => { "id" => "bartowski/…", "name" => "Qwen 2.5 1.5B (GGUF)", … }
33 //!
34 //! # Sampling config helpers.
35 //! Onde.default_sampling_config
36 //! Onde.deterministic_sampling_config
37 //! Onde.mobile_sampling_config
38 //! ```
39
40 use magnus::{function, prelude::*, Error, RHash, Ruby};
41
42 use onde::hf_cache;
43 use onde::inference::models::{SUPPORTED_MODELS, SUPPORTED_MODEL_INFO};
44 use onde::inference::types::SamplingConfig;
45
46 // ---------------------------------------------------------------------------
47 // Helpers — convert Rust structs to Ruby hashes via serde_json
48 // ---------------------------------------------------------------------------
49
50 /// Serialize any `serde::Serialize` value into a Ruby Hash (or Array of
51 /// Hashes) by round-tripping through serde_json. This keeps the binding
52 /// layer thin — we don't need to define Ruby classes for every Onde struct.
53 fn to_ruby_value<T: serde::Serialize>(ruby: &Ruby, value: &T) -> Result<magnus::Value, Error> {
54 let json = serde_json::to_value(value).map_err(|e| {
55 Error::new(
56 ruby.exception_runtime_error(),
57 format!("serialization error: {e}"),
58 )
59 })?;
60 json_to_ruby(ruby, &json)
61 }
62
63 fn json_to_ruby(ruby: &Ruby, value: &serde_json::Value) -> Result<magnus::Value, Error> {
64 match value {
65 serde_json::Value::Null => Ok(ruby.qnil().as_value()),
66 serde_json::Value::Bool(b) => {
67 if *b {
68 Ok(ruby.qtrue().as_value())
69 } else {
70 Ok(ruby.qfalse().as_value())
71 }
72 }
73 serde_json::Value::Number(n) => {
74 if let Some(i) = n.as_i64() {
75 Ok(ruby.integer_from_i64(i).as_value())
76 } else if let Some(f) = n.as_f64() {
77 Ok(ruby.float_from_f64(f).as_value())
78 } else {
79 Ok(ruby.qnil().as_value())
80 }
81 }
82 serde_json::Value::String(s) => Ok(ruby.str_new(s).as_value()),
83 serde_json::Value::Array(arr) => {
84 let ary = ruby.ary_new_capa(arr.len());
85 for item in arr {
86 ary.push(json_to_ruby(ruby, item)?)?;
87 }
88 Ok(ary.as_value())
89 }
90 serde_json::Value::Object(map) => {
91 let hash = ruby.hash_new();
92 for (k, v) in map {
93 hash.aset(ruby.str_new(k), json_to_ruby(ruby, v)?)?;
94 }
95 Ok(hash.as_value())
96 }
97 }
98 }
99
100 // ---------------------------------------------------------------------------
101 // Exported Ruby methods
102 // ---------------------------------------------------------------------------
103
104 /// `Onde.list_local_models` → Hash
105 ///
106 /// Scans the local HuggingFace hub cache and returns all downloaded models
107 /// that the inference engine supports.
108 ///
109 /// Returns a Hash with keys: `"models"`, `"cache_path"`,
110 /// `"total_size_bytes"`, `"total_size_display"`.
111 fn list_local_models() -> Result<magnus::Value, Error> {
112 let ruby = Ruby::get().expect("called outside Ruby");
113 let response = hf_cache::list_local_hf_models();
114 to_ruby_value(&ruby, &response)
115 }
116
117 /// `Onde.list_supported_models` → Hash
118 ///
119 /// Returns all models the engine supports, together with flags indicating
120 /// whether each one is fully downloaded, partially downloaded, or absent.
121 ///
122 /// Returns a Hash with key `"models"` containing an Array of model Hashes.
123 fn list_supported_models() -> Result<magnus::Value, Error> {
124 let ruby = Ruby::get().expect("called outside Ruby");
125 let response = hf_cache::list_supported_hf_models();
126 to_ruby_value(&ruby, &response)
127 }
128
129 /// `Onde.delete_model(model_id)` → nil
130 ///
131 /// Delete a locally cached HuggingFace model.
132 /// `model_id` is the full identifier, e.g. `"black-forest-labs/FLUX.1-schnell"`.
133 ///
134 /// Raises `RuntimeError` if the model is not found or deletion fails.
135 fn delete_model(model_id: String) -> Result<magnus::Value, Error> {
136 let ruby = Ruby::get().expect("called outside Ruby");
137 hf_cache::delete_local_hf_model(model_id)
138 .map_err(|e| Error::new(ruby.exception_runtime_error(), e))?;
139 Ok(ruby.qnil().as_value())
140 }
141
142 /// `Onde.model_info(model_id)` → Hash or nil
143 ///
144 /// Look up rich metadata for a supported model by its ID.
145 /// Returns `nil` if the model ID is not in the supported list.
146 fn model_info(model_id: String) -> Result<magnus::Value, Error> {
147 let ruby = Ruby::get().expect("called outside Ruby");
148
149 let info = SUPPORTED_MODEL_INFO.iter().find(|i| i.id == model_id);
150
151 match info {
152 None => Ok(ruby.qnil().as_value()),
153 Some(i) => {
154 let hash = ruby.hash_new();
155 hash.aset(ruby.str_new("id"), ruby.str_new(i.id))?;
156 hash.aset(ruby.str_new("name"), ruby.str_new(i.name))?;
157 hash.aset(ruby.str_new("org"), ruby.str_new(i.org))?;
158 hash.aset(ruby.str_new("description"), ruby.str_new(i.description))?;
159 hash.aset(
160 ruby.str_new("expected_size_bytes"),
161 ruby.integer_from_u64(i.expected_size_bytes),
162 )?;
163 Ok(hash.as_value())
164 }
165 }
166 }
167
168 /// `Onde.supported_model_ids` → Array of String
169 ///
170 /// Returns the list of all supported model IDs.
171 fn supported_model_ids() -> Result<magnus::Value, Error> {
172 let ruby = Ruby::get().expect("called outside Ruby");
173 let ary = ruby.ary_new_capa(SUPPORTED_MODELS.len());
174 for id in SUPPORTED_MODELS {
175 ary.push(ruby.str_new(id))?;
176 }
177 Ok(ary.as_value())
178 }
179
180 /// `Onde.default_sampling_config` → Hash
181 fn default_sampling_config() -> Result<magnus::Value, Error> {
182 let ruby = Ruby::get().expect("called outside Ruby");
183 to_ruby_value(&ruby, &SamplingConfig::default())
184 }
185
186 /// `Onde.deterministic_sampling_config` → Hash
187 fn deterministic_sampling_config() -> Result<magnus::Value, Error> {
188 let ruby = Ruby::get().expect("called outside Ruby");
189 to_ruby_value(&ruby, &SamplingConfig::deterministic())
190 }
191
192 /// `Onde.mobile_sampling_config` → Hash
193 fn mobile_sampling_config() -> Result<magnus::Value, Error> {
194 let ruby = Ruby::get().expect("called outside Ruby");
195 to_ruby_value(&ruby, &SamplingConfig::mobile())
196 }
197
198 /// `Onde.cache_path` → String or nil
199 ///
200 /// Returns the resolved HuggingFace cache directory path, or nil if it
201 /// cannot be determined (e.g. `$HOME` is unset).
202 fn cache_path() -> Result<magnus::Value, Error> {
203 let ruby = Ruby::get().expect("called outside Ruby");
204 let response = hf_cache::list_local_hf_models();
205 if response.cache_path.is_empty() {
206 Ok(ruby.qnil().as_value())
207 } else {
208 Ok(ruby.str_new(&response.cache_path).as_value())
209 }
210 }
211
212 // ---------------------------------------------------------------------------
213 // Init — called by Ruby when `require "onde/onde"` loads the shared library
214 // ---------------------------------------------------------------------------
215
216 #[magnus::init]
217 fn init(ruby: &Ruby) -> Result<(), Error> {
218 let module = ruby.define_module("Onde")?;
219
220 // -- Singleton methods ----------------------------------------------------
221 module.define_singleton_method("list_local_models", function!(list_local_models, 0))?;
222 module.define_singleton_method("list_supported_models", function!(list_supported_models, 0))?;
223 module.define_singleton_method("delete_model", function!(delete_model, 1))?;
224 module.define_singleton_method("model_info", function!(model_info, 1))?;
225 module.define_singleton_method("supported_model_ids", function!(supported_model_ids, 0))?;
226 module.define_singleton_method(
227 "default_sampling_config",
228 function!(default_sampling_config, 0),
229 )?;
230 module.define_singleton_method(
231 "deterministic_sampling_config",
232 function!(deterministic_sampling_config, 0),
233 )?;
234 module.define_singleton_method(
235 "mobile_sampling_config",
236 function!(mobile_sampling_config, 0),
237 )?;
238 module.define_singleton_method("cache_path", function!(cache_path, 0))?;
239
240 // -- Constants ------------------------------------------------------------
241
242 // Onde::NATIVE_VERSION (Rust crate version for parity checks).
243 module.const_set("NATIVE_VERSION", ruby.str_new(env!("CARGO_PKG_VERSION")))?;
244
245 // Onde::SUPPORTED_MODELS — frozen Array of model ID strings.
246 let model_ids = ruby.ary_new_capa(SUPPORTED_MODELS.len());
247 for id in SUPPORTED_MODELS {
248 model_ids.push(ruby.str_new(id))?;
249 }
250 model_ids.freeze();
251 module.const_set("SUPPORTED_MODELS", model_ids)?;
252
253 // Onde::SUPPORTED_MODEL_INFO — frozen Array of frozen Hashes.
254 let info_ary = ruby.ary_new_capa(SUPPORTED_MODEL_INFO.len());
255 for info in SUPPORTED_MODEL_INFO {
256 let hash: RHash = ruby.hash_new();
257 hash.aset(ruby.str_new("id"), ruby.str_new(info.id))?;
258 hash.aset(ruby.str_new("name"), ruby.str_new(info.name))?;
259 hash.aset(ruby.str_new("org"), ruby.str_new(info.org))?;
260 hash.aset(ruby.str_new("description"), ruby.str_new(info.description))?;
261 hash.aset(
262 ruby.str_new("expected_size_bytes"),
263 ruby.integer_from_u64(info.expected_size_bytes),
264 )?;
265 hash.freeze();
266 info_ary.push(hash)?;
267 }
268 info_ary.freeze();
269 module.const_set("SUPPORTED_MODEL_INFO", info_ary)?;
270
271 Ok(())
272 }