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, DownloadError};
11
12/// Options controlling model auto-download from the Dashboard private source.
13#[derive(Default, Clone)]
14pub struct OpenOptions {
15    /// Dashboard bearer token. Required when `model_ref` is a model name.
16    pub token: Option<String>,
17    /// Dashboard base URL. Defaults to `https://ariacompute.com`.
18    pub site: Option<String>,
19}
20
21/// High-level convenience over `Session`.
22pub struct Engine {
23    session: Session,
24}
25
26impl Engine {
27    /// Open a local Aria quant bundle directory (path only; no download).
28    pub fn open(bundle_path: impl AsRef<std::path::Path>) -> Result<Self, EngineError> {
29        let session = SessionBuilder::new().model(bundle_path).build()?;
30        Ok(Self { session })
31    }
32
33    /// Open a model by reference. If `model_ref` looks like a local path
34    /// (contains a separator or exists on disk) it is loaded directly;
35    /// otherwise it is treated as a model name and auto-downloaded from the
36    /// Dashboard source into `~/.ariacompute/models/{model}` before loading.
37    pub fn open_model(model_ref: &str, opts: &OpenOptions) -> Result<Self, OpenError> {
38        let path = if model_ref.contains('/') || model_ref.contains('\\') || std::path::Path::new(model_ref).exists() {
39            std::path::PathBuf::from(model_ref)
40        } else {
41            let token = opts
42                .token
43                .as_ref()
44                .ok_or_else(|| OpenError::MissingToken(model_ref.to_string()))?;
45            download_model(model_ref, token, opts.site.as_deref())
46                .map_err(OpenError::Download)?
47        };
48        let session = SessionBuilder::new()
49            .model(&path)
50            .build()
51            .map_err(OpenError::Engine)?;
52        Ok(Self { session })
53    }
54
55    pub fn complete(&mut self, prompt: &str, opts: &GenerateOpts) -> Result<Generation, EngineError> {
56        let turns = [aria_inference::ChatTurn::new("user", prompt)];
57        let tokens = self.session.encode_chat(&turns);
58        self.session.generate(&tokens, opts)
59    }
60
61    pub fn embed(&self, text: &str) -> Result<Vec<f32>, EngineError> {
62        self.session.embed_text(text)
63    }
64
65    pub fn transcribe(&self, pcm: &[u8]) -> Result<String, EngineError> {
66        self.session.transcribe_pcm16le(pcm)
67    }
68}
69
70/// Error returned by [`Engine::open_model`].
71#[derive(Debug, thiserror::Error)]
72pub enum OpenError {
73    #[error("model name '{0}' requires an api token")]
74    MissingToken(String),
75    #[error("download failed: {0}")]
76    Download(#[from] DownloadError),
77    #[error("engine load failed: {0}")]
78    Engine(#[from] EngineError),
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use aria_inference::fixture::write_tiny_q4_bundle;
85
86    #[test]
87    fn engine_complete_ok() {
88        let dir = tempfile::tempdir().unwrap();
89        write_tiny_q4_bundle(dir.path()).unwrap();
90        let mut eng = Engine::open(dir.path()).unwrap();
91        let g = eng
92            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
93            .unwrap();
94        assert!(!g.text.is_empty());
95        assert!(!eng.embed("x").unwrap().is_empty());
96        assert!(!eng.transcribe(&[0, 1, 2, 3]).unwrap().is_empty());
97    }
98
99    #[test]
100    fn open_model_local_path_no_token() {
101        let dir = tempfile::tempdir().unwrap();
102        write_tiny_q4_bundle(dir.path()).unwrap();
103        let mut eng = Engine::open_model(dir.path().to_str().unwrap(), &OpenOptions::default()).unwrap();
104        let g = eng
105            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
106            .unwrap();
107        assert!(!g.text.is_empty());
108    }
109
110    #[test]
111    fn open_model_name_requires_token() {
112        let err = Engine::open_model("gemma-4-e2b-it_q4", &OpenOptions::default());
113        assert!(matches!(err, Err(OpenError::MissingToken(_))));
114    }
115}
aria_engine/lib.rs

aria_engine/
lib.rs
