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;
10pub use download::{download_model, download_model_auth, ensure_ffi_lib, DownloadError};
11
12/// Options controlling model auto-download from the regional public hub.
13#[derive(Default, Clone)]
14pub struct OpenOptions {
15    /// Legacy generic hub token. Dashboard `sk-` / `bfvk-` keys are ignored.
16    pub token: Option<String>,
17    /// Hugging Face hub token (`.com`). Same field as `aria-engine auth` `hf_token`.
18    pub hf_token: Option<String>,
19    /// ModelScope hub token (`.cn`). Same field as `aria-engine auth` `modelscope_api_token`.
20    pub modelscope_api_token: Option<String>,
21    /// Site used to pick the regional hub. Defaults to `https://ariacompute.com` (`.com` → HF, `.cn` → ModelScope).
22    pub site: Option<String>,
23}
24
25/// High-level convenience over `Session`.
26pub struct Engine {
27    session: Session,
28}
29
30impl Engine {
31    /// Open a local Aria quant bundle directory (path only; no download).
32    pub fn open(bundle_path: impl AsRef<std::path::Path>) -> Result<Self, EngineError> {
33        let session = SessionBuilder::new().model(bundle_path).build()?;
34        Ok(Self { session })
35    }
36
37    /// Open a model by reference. If `model_ref` looks like a local path
38    /// (contains a separator or exists on disk) it is loaded directly;
39    /// otherwise it is treated as a model name and auto-downloaded from the
40    /// regional public hub into `~/.ariacompute/models/{model}` before loading.
41    pub fn open_model(model_ref: &str, opts: &OpenOptions) -> Result<Self, OpenError> {
42        let _ = ensure_ffi_lib(opts.site.as_deref())?;
43        let path = if model_ref.contains('/') || model_ref.contains('\\') || std::path::Path::new(model_ref).exists() {
44            std::path::PathBuf::from(model_ref)
45        } else {
46            download_model_auth(
47                model_ref,
48                opts.token.as_deref().unwrap_or(""),
49                opts.site.as_deref(),
50                opts.hf_token.as_deref(),
51                opts.modelscope_api_token.as_deref(),
52            )
53            .map_err(OpenError::Download)?
54        };
55        let session = SessionBuilder::new()
56            .model(&path)
57            .build()
58            .map_err(OpenError::Engine)?;
59        Ok(Self { session })
60    }
61
62    pub fn complete(&mut self, prompt: &str, opts: &GenerateOpts) -> Result<Generation, EngineError> {
63        let turns = [aria_inference::ChatTurn::new("user", prompt)];
64        let tokens = self.session.encode_chat(&turns);
65        self.session.generate(&tokens, opts)
66    }
67
68    pub fn embed(&self, text: &str) -> Result<Vec<f32>, EngineError> {
69        self.session.embed_text(text)
70    }
71
72    pub fn transcribe(&self, pcm: &[u8]) -> Result<String, EngineError> {
73        self.session.transcribe_pcm16le(pcm)
74    }
75}
76
77/// Error returned by [`Engine::open_model`].
78#[derive(Debug, thiserror::Error)]
79pub enum OpenError {
80    #[error("download failed: {0}")]
81    Download(#[from] DownloadError),
82    #[error("engine load failed: {0}")]
83    Engine(#[from] EngineError),
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use aria_inference::fixture::write_tiny_q4_bundle;
90
91    #[test]
92    fn engine_complete_ok() {
93        let dir = tempfile::tempdir().unwrap();
94        write_tiny_q4_bundle(dir.path()).unwrap();
95        let mut eng = Engine::open(dir.path()).unwrap();
96        let g = eng
97            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
98            .unwrap();
99        assert!(!g.text.is_empty());
100        assert!(!eng.embed("x").unwrap().is_empty());
101        assert!(!eng.transcribe(&[0, 1, 2, 3]).unwrap().is_empty());
102    }
103
104    #[test]
105    fn open_model_local_path_no_token() {
106        let _guard = crate::download::ENV_LOCK.lock().unwrap();
107        let home = tempfile::tempdir().unwrap();
108        std::env::set_var("ARIA_COMPUTE_HOME", home.path());
109        let libdir = home.path().join("lib");
110        std::fs::create_dir_all(&libdir).unwrap();
111        let name = if cfg!(windows) {
112            "aria_ffi.dll"
113        } else if cfg!(target_os = "macos") {
114            "libaria_ffi.dylib"
115        } else {
116            "libaria_ffi.so"
117        };
118        std::fs::write(libdir.join(name), b"x").unwrap();
119        let dir = tempfile::tempdir().unwrap();
120        write_tiny_q4_bundle(dir.path()).unwrap();
121        let mut eng = Engine::open_model(dir.path().to_str().unwrap(), &OpenOptions::default()).unwrap();
122        let g = eng
123            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
124            .unwrap();
125        assert!(!g.text.is_empty());
126        std::env::remove_var("ARIA_COMPUTE_HOME");
127    }
128}
aria_engine/lib.rs

aria_engine/
lib.rs
