aria_engine/
lib.rs

1//! Thin Rust SDK: prefer `aria_inference` for native use; re-exports FFI for embedding tests.
2
3pub use aria_ffi::{
4    aria_complete, aria_complete_stream, aria_embed, aria_last_error, aria_model_destroy,
5    aria_model_init, aria_transcribe, AriaModelHandle,
6};
7pub use aria_inference::{EngineError, GenerateOpts, Generation, Session, SessionBuilder};
8
9mod download;
10mod auth;
11pub use download::{download_model, download_model_auth, ensure_ffi_lib, DownloadError};
12pub use auth::{
13    apply_auth, fill_auth_urls, AuthConfig, AuthError, AuthUpdates, CN_CLOUD, CN_SITE, CN_UPGRADE,
14};
15
16/// Options controlling model auto-download from the regional public hub.
17#[derive(Default, Clone)]
18pub struct OpenOptions {
19    /// Legacy generic hub token. Dashboard `sk-` / `bfvk-` keys are ignored.
20    pub token: Option<String>,
21    /// Hugging Face hub token (`.com`). Same field as `aria-engine auth` `hf_token`.
22    pub hf_token: Option<String>,
23    /// ModelScope hub token (`.cn`). Same field as `aria-engine auth` `modelscope_api_token`.
24    pub modelscope_api_token: Option<String>,
25    /// Site used to pick the regional hub. Defaults to `https://ariacompute.com` (`.com` → HF, `.cn` → ModelScope).
26    pub site: Option<String>,
27}
28
29/// High-level convenience over `Session`.
30pub struct Engine {
31    session: Option<Session>,
32    auth: AuthConfig,
33    generic_token: Option<String>,
34}
35
36impl Engine {
37    /// Empty construct. Call [`auth`](Self::auth) then [`open`](Self::open) to download/load.
38    pub fn new() -> Self {
39        Self {
40            session: None,
41            auth: AuthConfig::default(),
42            generic_token: None,
43        }
44    }
45
46    /// Set Config / Run fields on this instance only. Does not write config.yml.
47    pub fn auth(&mut self, updates: &AuthUpdates) -> Result<&mut Self, AuthError> {
48        self.auth = apply_auth(&self.auth, updates)?;
49        Ok(self)
50    }
51
52    pub fn auth_status(&self) -> &AuthConfig {
53        &self.auth
54    }
55
56    /// Reset instance defaults. Does not delete ~/.ariacompute/config.yml.
57    pub fn auth_clear(&mut self) -> &mut Self {
58        self.auth = AuthConfig::default();
59        self
60    }
61
62    /// Open a local Aria quant bundle directory (path only; no download).
63    pub fn from_bundle(bundle_path: impl AsRef<std::path::Path>) -> Result<Self, EngineError> {
64        let session = SessionBuilder::new().model(bundle_path).build()?;
65        Ok(Self {
66            session: Some(session),
67            auth: AuthConfig::default(),
68            generic_token: None,
69        })
70    }
71
72    /// Open a local Aria quant bundle directory (path only; no download).
73    pub fn open(bundle_path: impl AsRef<std::path::Path>) -> Result<Self, EngineError> {
74        Self::from_bundle(bundle_path)
75    }
76
77    fn load_ref(&mut self, model_ref: &str, opts: &OpenOptions) -> Result<(), OpenError> {
78        let _ = ensure_ffi_lib(opts.site.as_deref())?;
79        let path = if model_ref.contains('/')
80            || model_ref.contains('\\')
81            || std::path::Path::new(model_ref).exists()
82        {
83            std::path::PathBuf::from(model_ref)
84        } else {
85            download_model_auth(
86                model_ref,
87                opts.token.as_deref().unwrap_or(""),
88                opts.site.as_deref(),
89                opts.hf_token.as_deref(),
90                opts.modelscope_api_token.as_deref(),
91            )
92            .map_err(OpenError::Download)?
93        };
94        let session = SessionBuilder::new()
95            .model(&path)
96            .build()
97            .map_err(OpenError::Engine)?;
98        self.session = Some(session);
99        Ok(())
100    }
101
102    /// Download (if needed) and load a model using instance auth.
103    pub fn open_named(&mut self, model_ref: &str) -> Result<&mut Self, OpenError> {
104        let opts = OpenOptions {
105            token: self.generic_token.clone(),
106            hf_token: if self.auth.hf_token.is_empty() {
107                None
108            } else {
109                Some(self.auth.hf_token.clone())
110            },
111            modelscope_api_token: if self.auth.modelscope_api_token.is_empty() {
112                None
113            } else {
114                Some(self.auth.modelscope_api_token.clone())
115            },
116            site: if self.auth.site_url.is_empty() {
117                None
118            } else {
119                Some(self.auth.site_url.clone())
120            },
121        };
122        self.load_ref(model_ref, &opts)?;
123        Ok(self)
124    }
125
126    /// Open a model by reference. If `model_ref` looks like a local path
127    /// (contains a separator or exists on disk) it is loaded directly;
128    /// otherwise it is treated as a model name and auto-downloaded from the
129    /// regional public hub into `~/.ariacompute/models/{model}` before loading.
130    pub fn open_model(model_ref: &str, opts: &OpenOptions) -> Result<Self, OpenError> {
131        let mut eng = Self::new();
132        let mut updates = AuthUpdates::default();
133        updates.site_url = opts.site.clone();
134        updates.hf_token = opts.hf_token.clone();
135        updates.modelscope_api_token = opts.modelscope_api_token.clone();
136        eng.generic_token = opts.token.clone();
137        let _ = eng.auth(&updates);
138        eng.load_ref(model_ref, opts)?;
139        Ok(eng)
140    }
141
142    fn session_mut(&mut self) -> Result<&mut Session, EngineError> {
143        self.session
144            .as_mut()
145            .ok_or_else(|| EngineError::InvalidParam("engine not opened".into()))
146    }
147
148    fn session_ref(&self) -> Result<&Session, EngineError> {
149        self.session
150            .as_ref()
151            .ok_or_else(|| EngineError::InvalidParam("engine not opened".into()))
152    }
153
154    pub fn complete(&mut self, prompt: &str, opts: &GenerateOpts) -> Result<Generation, EngineError> {
155        let turns = [aria_inference::ChatTurn::new("user", prompt)];
156        let session = self.session_mut()?;
157        let tokens = session.encode_chat(&turns);
158        session.generate(&tokens, opts)
159    }
160
161    pub fn embed(&self, text: &str) -> Result<Vec<f32>, EngineError> {
162        self.session_ref()?.embed_text(text)
163    }
164
165    pub fn transcribe(&self, pcm: &[u8]) -> Result<String, EngineError> {
166        self.session_ref()?.transcribe_pcm16le(pcm)
167    }
168}
169
170impl Default for Engine {
171    fn default() -> Self {
172        Self::new()
173    }
174}
175
176/// Error returned by [`Engine::open_model`].
177#[derive(Debug, thiserror::Error)]
178pub enum OpenError {
179    #[error("download failed: {0}")]
180    Download(#[from] DownloadError),
181    #[error("engine load failed: {0}")]
182    Engine(#[from] EngineError),
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use aria_inference::fixture::write_tiny_q4_bundle;
189
190    #[test]
191    fn engine_complete_ok() {
192        let dir = tempfile::tempdir().unwrap();
193        write_tiny_q4_bundle(dir.path()).unwrap();
194        let mut eng = Engine::open(dir.path()).unwrap();
195        let g = eng
196            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
197            .unwrap();
198        assert!(!g.text.is_empty());
199        assert!(!eng.embed("x").unwrap().is_empty());
200        assert!(!eng.transcribe(&[0, 1, 2, 3]).unwrap().is_empty());
201    }
202
203    #[test]
204    fn open_model_local_path_no_token() {
205        let _guard = crate::download::ENV_LOCK.lock().unwrap();
206        let home = tempfile::tempdir().unwrap();
207        std::env::set_var("ARIA_COMPUTE_HOME", home.path());
208        let libdir = home.path().join("lib");
209        std::fs::create_dir_all(&libdir).unwrap();
210        let name = if cfg!(windows) {
211            "aria_ffi.dll"
212        } else if cfg!(target_os = "macos") {
213            "libaria_ffi.dylib"
214        } else {
215            "libaria_ffi.so"
216        };
217        std::fs::write(libdir.join(name), b"x").unwrap();
218        let dir = tempfile::tempdir().unwrap();
219        write_tiny_q4_bundle(dir.path()).unwrap();
220        let mut eng = Engine::open_model(dir.path().to_str().unwrap(), &OpenOptions::default()).unwrap();
221        let g = eng
222            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
223            .unwrap();
224        assert!(!g.text.is_empty());
225        std::env::remove_var("ARIA_COMPUTE_HOME");
226    }
227
228    #[test]
229    fn auth_instance_all_fields() {
230        let mut eng = Engine::new();
231        eng.auth(&AuthUpdates {
232            cloud_api_key: Some("sk-test".into()),
233            cloud_url: Some(CN_CLOUD.into()),
234            site_url: Some(crate::auth::CN_SITE.into()),
235            upgrade_url: Some(crate::auth::CN_UPGRADE.into()),
236            hybrid_mode: Some("cost".into()),
237            hybrid_execution: Some("device".into()),
238            hybrid_semantic: Some(false),
239            hybrid_semantic_timeout_ms: Some(250),
240            hybrid_semantic_cache_size: Some(16),
241            compute: Some("cpu".into()),
242            hf_token: Some("hf_abc".into()),
243            modelscope_api_token: Some("ms_xyz".into()),
244        })
245        .unwrap();
246        let st = eng.auth_status();
247        assert_eq!(st.cloud_api_key, "sk-test");
248        assert_eq!(st.hybrid_mode, "cost");
249        assert_eq!(st.hybrid_execution, "device");
250        assert!(!st.hybrid_semantic);
251        assert_eq!(st.hybrid_semantic_timeout_ms, 250);
252        assert_eq!(st.compute, "cpu");
253        assert_eq!(st.hf_token, "hf_abc");
254        assert_eq!(st.modelscope_api_token, "ms_xyz");
255        assert_eq!(st.site_url, crate::auth::CN_SITE);
256    }
257
258    #[test]
259    fn auth_partial_merge() {
260        let mut eng = Engine::new();
261        eng.auth(&AuthUpdates {
262            hf_token: Some("hf_one".into()),
263            hybrid_mode: Some("intelligence".into()),
264            ..Default::default()
265        })
266        .unwrap();
267        eng.auth(&AuthUpdates {
268            compute: Some("cuda".into()),
269            ..Default::default()
270        })
271        .unwrap();
272        let st = eng.auth_status();
273        assert_eq!(st.hf_token, "hf_one");
274        assert_eq!(st.hybrid_mode, "intelligence");
275        assert_eq!(st.compute, "cuda");
276    }
277
278    #[test]
279    fn auth_invalid_enum_leaves_state() {
280        let mut eng = Engine::new();
281        eng.auth(&AuthUpdates {
282            hybrid_mode: Some("cost".into()),
283            ..Default::default()
284        })
285        .unwrap();
286        assert!(eng
287            .auth(&AuthUpdates {
288                hybrid_mode: Some("nope".into()),
289                ..Default::default()
290            })
291            .is_err());
292        assert_eq!(eng.auth_status().hybrid_mode, "cost");
293    }
294
295    #[test]
296    fn auth_clear_resets_instance() {
297        let mut eng = Engine::new();
298        eng.auth(&AuthUpdates {
299            hf_token: Some("hf_x".into()),
300            hybrid_mode: Some("cost".into()),
301            ..Default::default()
302        })
303        .unwrap();
304        eng.auth_clear();
305        let st = eng.auth_status();
306        assert_eq!(st.hf_token, "");
307        assert_eq!(st.hybrid_mode, "balance");
308    }
309
310    #[test]
311    fn auth_fills_urls_from_site_tld() {
312        let mut eng = Engine::new();
313        eng.auth(&AuthUpdates {
314            site_url: Some("https://ariacompute.cn".into()),
315            ..Default::default()
316        })
317        .unwrap();
318        let st = eng.auth_status();
319        assert_eq!(st.cloud_url, CN_CLOUD);
320        assert_eq!(st.upgrade_url, crate::auth::CN_UPGRADE);
321    }
322
323    #[test]
324    fn auth_does_not_write_config_yml() {
325        let _guard = crate::download::ENV_LOCK.lock().unwrap();
326        let home = tempfile::tempdir().unwrap();
327        std::env::set_var("ARIA_COMPUTE_HOME", home.path());
328        let mut eng = Engine::new();
329        eng.auth(&AuthUpdates {
330            cloud_api_key: Some("sk-test".into()),
331            site_url: Some("https://ariacompute.com".into()),
332            hf_token: Some("hf_x".into()),
333            ..Default::default()
334        })
335        .unwrap();
336        assert!(!home.path().join("config.yml").is_file());
337        std::env::remove_var("ARIA_COMPUTE_HOME");
338    }
339
340    #[test]
341    fn auth_detect_urls_from_key_mocked() {
342        crate::auth::set_probe_dashboard(|site, _key| site.contains("ariacompute.cn"));
343        let mut eng = Engine::new();
344        let result = eng.auth(&AuthUpdates {
345            cloud_api_key: Some("sk-region".into()),
346            ..Default::default()
347        });
348        crate::auth::reset_probe_dashboard();
349        result.unwrap();
350        let st = eng.auth_status();
351        assert_eq!(st.site_url, crate::auth::CN_SITE);
352        assert_eq!(st.cloud_url, CN_CLOUD);
353        assert_eq!(st.upgrade_url, crate::auth::CN_UPGRADE);
354    }
355}
aria_engine/lib.rs

aria_engine/
lib.rs
