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, 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 path = if model_ref.contains('/') || model_ref.contains('\\') || std::path::Path::new(model_ref).exists() {
43            std::path::PathBuf::from(model_ref)
44        } else {
45            download_model_auth(
46                model_ref,
47                opts.token.as_deref().unwrap_or(""),
48                opts.site.as_deref(),
49                opts.hf_token.as_deref(),
50                opts.modelscope_api_token.as_deref(),
51            )
52            .map_err(OpenError::Download)?
53        };
54        let session = SessionBuilder::new()
55            .model(&path)
56            .build()
57            .map_err(OpenError::Engine)?;
58        Ok(Self { session })
59    }
60
61    pub fn complete(&mut self, prompt: &str, opts: &GenerateOpts) -> Result<Generation, EngineError> {
62        let turns = [aria_inference::ChatTurn::new("user", prompt)];
63        let tokens = self.session.encode_chat(&turns);
64        self.session.generate(&tokens, opts)
65    }
66
67    pub fn embed(&self, text: &str) -> Result<Vec<f32>, EngineError> {
68        self.session.embed_text(text)
69    }
70
71    pub fn transcribe(&self, pcm: &[u8]) -> Result<String, EngineError> {
72        self.session.transcribe_pcm16le(pcm)
73    }
74}
75
76/// Error returned by [`Engine::open_model`].
77#[derive(Debug, thiserror::Error)]
78pub enum OpenError {
79    #[error("download failed: {0}")]
80    Download(#[from] DownloadError),
81    #[error("engine load failed: {0}")]
82    Engine(#[from] EngineError),
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use aria_inference::fixture::write_tiny_q4_bundle;
89
90    #[test]
91    fn engine_complete_ok() {
92        let dir = tempfile::tempdir().unwrap();
93        write_tiny_q4_bundle(dir.path()).unwrap();
94        let mut eng = Engine::open(dir.path()).unwrap();
95        let g = eng
96            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
97            .unwrap();
98        assert!(!g.text.is_empty());
99        assert!(!eng.embed("x").unwrap().is_empty());
100        assert!(!eng.transcribe(&[0, 1, 2, 3]).unwrap().is_empty());
101    }
102
103    #[test]
104    fn open_model_local_path_no_token() {
105        let dir = tempfile::tempdir().unwrap();
106        write_tiny_q4_bundle(dir.path()).unwrap();
107        let mut eng = Engine::open_model(dir.path().to_str().unwrap(), &OpenOptions::default()).unwrap();
108        let g = eng
109            .complete("hi", &GenerateOpts { max_tokens: 2, temperature: 0.0 })
110            .unwrap();
111        assert!(!g.text.is_empty());
112    }
113}
aria_engine/lib.rs

aria_engine/
lib.rs
