aria_ffi/
lib.rs

1//! C ABI for Aria inference (chat / embeddings / ASR / tools).
2
3#![allow(clippy::not_unsafe_ptr_arg_deref)] // C ABI: pointers are caller-owned
4#![allow(clippy::too_many_arguments)]
5
6use aria_inference::{ChatTurn, GenerateOpts, Session, SessionBuilder};
7use serde_json::{json, Value};
8use std::cell::RefCell;
9use std::env;
10use std::ffi::{CStr, CString};
11use std::os::raw::{c_char, c_int, c_uchar, c_void};
12use std::path::PathBuf;
13use std::ptr;
14use std::slice;
15
16thread_local! {
17    static LAST_ERROR: RefCell<Option<CString>> = const { RefCell::new(None) };
18    // Cache for strings returned to the caller as `const char *`. The pointer is
19    // valid until the next call to a returning function on this thread.
20    static LAST_STRING: RefCell<Option<CString>> = const { RefCell::new(None) };
21}
22
23/// Resolve the user-level Aria home (`~/.ariacompute`), overridable via
24/// `ARIA_COMPUTE_HOME` (mirrors `openai/src/config.rs::aria_home`).
25fn aria_home() -> Result<PathBuf, String> {
26    if let Ok(override_home) = env::var("ARIA_COMPUTE_HOME") {
27        if !override_home.is_empty() {
28            return Ok(PathBuf::from(override_home));
29        }
30    }
31    let home = if cfg!(windows) {
32        env::var("USERPROFILE").map_err(|_| "could not resolve home directory".to_string())?
33    } else {
34        env::var("HOME").map_err(|_| "could not resolve home directory".to_string())?
35    };
36    Ok(PathBuf::from(home).join(".ariacompute"))
37}
38
39/// Return the on-disk cache directory for a model name:
40/// `~/.ariacompute/models/{model}`. Overridable via `ARIA_COMPUTE_HOME`.
41///
42/// The returned `const char *` is owned by a thread-local cache and is valid
43/// until the next call to a returning FFI function on this thread. Returns
44/// `NULL` on error (inspect `aria_last_error`).
45#[no_mangle]
46pub extern "C" fn aria_model_cache_dir(model: *const c_char) -> *const c_char {
47    clear_error();
48    let model = match cstr_to_str(model) {
49        Ok(m) => m,
50        Err(e) => {
51            set_error(e);
52            return ptr::null();
53        }
54    };
55    match aria_home() {
56        Ok(home) => {
57            let dir = home.join("models").join(model);
58            let s = match CString::new(dir.to_string_lossy().as_bytes()) {
59                Ok(s) => s,
60                Err(e) => {
61                    set_error(e.to_string());
62                    return ptr::null();
63                }
64            };
65            let ptr = s.as_ptr();
66            LAST_STRING.with(|c| *c.borrow_mut() = Some(s));
67            ptr
68        }
69        Err(e) => {
70            set_error(e);
71            ptr::null()
72        }
73    }
74}
75
76/// Whether `ref_` should be treated as a local bundle path rather than a model
77/// name. A value containing a path separator or one that already exists on disk
78/// is a local path.
79#[no_mangle]
80pub extern "C" fn aria_is_local_path(ref_: *const c_char) -> c_int {
81    clear_error();
82    let s = match cstr_to_str(ref_) {
83        Ok(s) => s,
84        Err(e) => {
85            set_error(e);
86            return -1;
87        }
88    };
89    let is_local = s.contains('/') || s.contains('\\') || std::path::Path::new(s).exists();
90    if is_local {
91        1
92    } else {
93        0
94    }
95}
96
97pub struct AriaModel {
98    session: Session,
99}
100
101fn set_error(msg: impl Into<String>) {
102    let s = CString::new(msg.into()).unwrap_or_else(|_| CString::new("error").unwrap());
103    LAST_ERROR.with(|e| *e.borrow_mut() = Some(s));
104}
105
106fn clear_error() {
107    LAST_ERROR.with(|e| *e.borrow_mut() = None);
108}
109
110fn cstr_to_str<'a>(p: *const c_char) -> Result<&'a str, String> {
111    if p.is_null() {
112        return Err("null string".into());
113    }
114    unsafe { CStr::from_ptr(p) }
115        .to_str()
116        .map_err(|e| e.to_string())
117}
118
119fn write_out(out: *mut c_char, out_len: usize, s: &str) -> c_int {
120    if out.is_null() || out_len == 0 {
121        set_error("null output buffer");
122        return -1;
123    }
124    let bytes = s.as_bytes();
125    if bytes.len() + 1 > out_len {
126        set_error(format!(
127            "output buffer too small: need {}, have {}",
128            bytes.len() + 1,
129            out_len
130        ));
131        return -1;
132    }
133    unsafe {
134        ptr::copy_nonoverlapping(bytes.as_ptr(), out.cast::<u8>(), bytes.len());
135        *out.add(bytes.len()) = 0;
136    }
137    0
138}
139
140fn parse_messages(messages_json: &str) -> Result<Vec<ChatTurn>, String> {
141    let v: Value = serde_json::from_str(messages_json).map_err(|e| e.to_string())?;
142    let arr = v
143        .as_array()
144        .ok_or_else(|| "messages must be a JSON array".to_string())?;
145    let mut turns = Vec::new();
146    for m in arr {
147        let role = m
148            .get("role")
149            .and_then(|x| x.as_str())
150            .unwrap_or("user")
151            .to_string();
152        let content = m
153            .get("content")
154            .and_then(|x| x.as_str())
155            .unwrap_or("")
156            .to_string();
157        turns.push(ChatTurn { role, content });
158    }
159    Ok(turns)
160}
161
162fn parse_options(options_json: Option<&str>) -> GenerateOpts {
163    let mut opts = GenerateOpts::default();
164    if let Some(raw) = options_json {
165        if let Ok(v) = serde_json::from_str::<Value>(raw) {
166            if let Some(n) = v.get("max_tokens").and_then(|x| x.as_u64()) {
167                opts.max_tokens = n as usize;
168            }
169            if let Some(t) = v.get("temperature").and_then(|x| x.as_f64()) {
170                opts.temperature = t as f32;
171            }
172        }
173    }
174    if opts.max_tokens == 0 {
175        opts.max_tokens = 16;
176    }
177    opts
178}
179
180fn parse_tools(tools_json: Option<&str>) -> Result<Value, String> {
181    match tools_json {
182        None | Some("") => Ok(json!([])),
183        Some(raw) => {
184            let v: Value = serde_json::from_str(raw).map_err(|e| e.to_string())?;
185            if !v.is_array() && !v.is_null() {
186                return Err("tools must be a JSON array or null".into());
187            }
188            Ok(if v.is_null() { json!([]) } else { v })
189        }
190    }
191}
192
193/// Opaque model handle.
194pub type AriaModelHandle = *mut AriaModel;
195
196/// Last error message (thread-local). Valid until next call on this thread.
197#[no_mangle]
198pub extern "C" fn aria_last_error() -> *const c_char {
199    LAST_ERROR.with(|e| match e.borrow().as_ref() {
200        Some(s) => s.as_ptr(),
201        None => ptr::null(),
202    })
203}
204
205/// Load an Aria quant bundle from `bundle_path`. Returns null on error.
206#[no_mangle]
207pub extern "C" fn aria_model_init(bundle_path: *const c_char) -> AriaModelHandle {
208    clear_error();
209    let path = match cstr_to_str(bundle_path) {
210        Ok(p) => p,
211        Err(e) => {
212            set_error(e);
213            return ptr::null_mut();
214        }
215    };
216    match SessionBuilder::new().model(path).build() {
217        Ok(session) => Box::into_raw(Box::new(AriaModel { session })),
218        Err(e) => {
219            set_error(e.to_string());
220            ptr::null_mut()
221        }
222    }
223}
224
225/// Destroy a model handle. Safe on null. Double-destroy is undefined if caller reuses the pointer.
226#[no_mangle]
227pub extern "C" fn aria_model_destroy(model: AriaModelHandle) {
228    clear_error();
229    if model.is_null() {
230        return;
231    }
232    unsafe {
233        drop(Box::from_raw(model));
234    }
235}
236
237fn complete_inner(
238    model: AriaModelHandle,
239    messages_json: *const c_char,
240    options_json: *const c_char,
241    tools_json: *const c_char,
242    out: *mut c_char,
243    out_len: usize,
244    stream_cb: Option<unsafe extern "C" fn(*const c_char, *mut c_void)>,
245    user_data: *mut c_void,
246) -> c_int {
247    clear_error();
248    if model.is_null() {
249        set_error("null model");
250        return -1;
251    }
252    let messages = match cstr_to_str(messages_json) {
253        Ok(s) => s,
254        Err(e) => {
255            set_error(e);
256            return -1;
257        }
258    };
259    let options = if options_json.is_null() {
260        None
261    } else {
262        match cstr_to_str(options_json) {
263            Ok(s) => Some(s),
264            Err(e) => {
265                set_error(e);
266                return -1;
267            }
268        }
269    };
270    let tools_raw = if tools_json.is_null() {
271        None
272    } else {
273        match cstr_to_str(tools_json) {
274            Ok(s) => Some(s),
275            Err(e) => {
276                set_error(e);
277                return -1;
278            }
279        }
280    };
281
282    let turns = match parse_messages(messages) {
283        Ok(p) => p,
284        Err(e) => {
285            set_error(e);
286            return -1;
287        }
288    };
289    let tools = match parse_tools(tools_raw) {
290        Ok(t) => t,
291        Err(e) => {
292            set_error(e);
293            return -1;
294        }
295    };
296    let opts = parse_options(options);
297    let m = unsafe { &mut *model };
298    let tokens = m.session.encode_chat(&turns);
299    let gen = match m.session.generate(&tokens, &opts) {
300        Ok(g) => g,
301        Err(e) => {
302            set_error(e.to_string());
303            return -1;
304        }
305    };
306
307    if let Some(cb) = stream_cb {
308        // Stream decoded text (same as gen.text), not raw `<id>` placeholders.
309        if let Ok(c) = CString::new(gen.text.as_str()) {
310            unsafe { cb(c.as_ptr(), user_data) };
311        }
312    }
313
314    let body = json!({
315        "success": true,
316        "error": null,
317        "response": gen.text,
318        "function_calls": json!([]),
319        "segments": [],
320        "cloud_handoff": false,
321        "total_tokens": gen.tokens.len(),
322    });
323    // tools accepted for OpenAI parity; real tool routing is stage C.
324    let _ = tools;
325    write_out(out, out_len, &body.to_string())
326}
327
328/// Non-streaming chat completion. `tools_json` may be null.
329#[no_mangle]
330pub extern "C" fn aria_complete(
331    model: AriaModelHandle,
332    messages_json: *const c_char,
333    options_json: *const c_char,
334    tools_json: *const c_char,
335    out: *mut c_char,
336    out_len: usize,
337) -> c_int {
338    complete_inner(
339        model,
340        messages_json,
341        options_json,
342        tools_json,
343        out,
344        out_len,
345        None,
346        ptr::null_mut(),
347    )
348}
349
350/// Streaming chat; `callback` receives each chunk as a C string.
351#[no_mangle]
352pub extern "C" fn aria_complete_stream(
353    model: AriaModelHandle,
354    messages_json: *const c_char,
355    options_json: *const c_char,
356    tools_json: *const c_char,
357    out: *mut c_char,
358    out_len: usize,
359    callback: Option<unsafe extern "C" fn(*const c_char, *mut c_void)>,
360    user_data: *mut c_void,
361) -> c_int {
362    complete_inner(
363        model,
364        messages_json,
365        options_json,
366        tools_json,
367        out,
368        out_len,
369        callback,
370        user_data,
371    )
372}
373
374/// Embeddings. `input_json` is a string or `{"input":"..."}` / `{"input":[...]}`.
375#[no_mangle]
376pub extern "C" fn aria_embed(
377    model: AriaModelHandle,
378    input_json: *const c_char,
379    out: *mut c_char,
380    out_len: usize,
381) -> c_int {
382    clear_error();
383    if model.is_null() {
384        set_error("null model");
385        return -1;
386    }
387    let raw = match cstr_to_str(input_json) {
388        Ok(s) => s,
389        Err(e) => {
390            set_error(e);
391            return -1;
392        }
393    };
394    let text = match serde_json::from_str::<Value>(raw) {
395        Ok(Value::String(s)) => s,
396        Ok(v) => v
397            .get("input")
398            .and_then(|x| {
399                x.as_str()
400                    .map(|s| s.to_string())
401                    .or_else(|| x.as_array().and_then(|a| a.first()).and_then(|x| x.as_str()).map(|s| s.to_string()))
402            })
403            .unwrap_or_default(),
404        Err(_) => raw.to_string(),
405    };
406    if text.is_empty() {
407        set_error("empty embedding input");
408        return -1;
409    }
410    let m = unsafe { &*model };
411    let emb = match m.session.embed_text(&text) {
412        Ok(e) => e,
413        Err(e) => {
414            set_error(e.to_string());
415            return -1;
416        }
417    };
418    let body = json!({
419        "object": "list",
420        "data": [{
421            "object": "embedding",
422            "embedding": emb,
423            "index": 0
424        }]
425    });
426    write_out(out, out_len, &body.to_string())
427}
428
429/// Transcribe PCM16 LE bytes.
430#[no_mangle]
431pub extern "C" fn aria_transcribe(
432    model: AriaModelHandle,
433    pcm: *const c_uchar,
434    pcm_len: usize,
435    _options_json: *const c_char,
436    out: *mut c_char,
437    out_len: usize,
438) -> c_int {
439    clear_error();
440    if model.is_null() {
441        set_error("null model");
442        return -1;
443    }
444    if pcm.is_null() || pcm_len == 0 {
445        set_error("empty pcm");
446        return -1;
447    }
448    let bytes = unsafe { slice::from_raw_parts(pcm, pcm_len) };
449    let m = unsafe { &*model };
450    let text = match m.session.transcribe_pcm16le(bytes) {
451        Ok(t) => t,
452        Err(e) => {
453            set_error(e.to_string());
454            return -1;
455        }
456    };
457    let body = json!({
458        "text": text,
459        "segments": []
460    });
461    write_out(out, out_len, &body.to_string())
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use aria_inference::fixture::write_tiny_q4_bundle;
468    use std::sync::atomic::{AtomicUsize, Ordering};
469
470    fn out_buf() -> Vec<u8> {
471        vec![0u8; 64 * 1024]
472    }
473
474    #[test]
475    fn init_complete_embed_transcribe_destroy() {
476        let dir = tempfile::tempdir().unwrap();
477        write_tiny_q4_bundle(dir.path()).unwrap();
478        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
479        let model = aria_model_init(path.as_ptr());
480        assert!(!model.is_null(), "{:?}", unsafe {
481            CStr::from_ptr(aria_last_error()).to_string_lossy()
482        });
483
484        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
485        let options = CString::new(r#"{"max_tokens":2}"#).unwrap();
486        let tools = CString::new("[]").unwrap();
487        let mut buf = out_buf();
488        assert_eq!(
489            aria_complete(
490                model,
491                messages.as_ptr(),
492                options.as_ptr(),
493                tools.as_ptr(),
494                buf.as_mut_ptr() as *mut c_char,
495                buf.len(),
496            ),
497            0
498        );
499        let s = unsafe { CStr::from_ptr(buf.as_ptr() as *const c_char) }
500            .to_str()
501            .unwrap();
502        let v: Value = serde_json::from_str(s).unwrap();
503        assert_eq!(v["success"], true);
504        assert!(!v["response"].as_str().unwrap().is_empty());
505
506        let input = CString::new(r#"{"input":"hello"}"#).unwrap();
507        buf.fill(0);
508        assert_eq!(
509            aria_embed(
510                model,
511                input.as_ptr(),
512                buf.as_mut_ptr() as *mut c_char,
513                buf.len()
514            ),
515            0
516        );
517
518        let pcm = [0u8, 1, 2, 3, 4, 5];
519        buf.fill(0);
520        assert_eq!(
521            aria_transcribe(
522                model,
523                pcm.as_ptr(),
524                pcm.len(),
525                ptr::null(),
526                buf.as_mut_ptr() as *mut c_char,
527                buf.len()
528            ),
529            0
530        );
531
532        aria_model_destroy(model);
533    }
534
535    #[test]
536    fn init_infers_family_from_hub_cache_dirname() {
537        let parent = tempfile::tempdir().unwrap();
538        let bundle = parent.path().join("gemma-3-1b-it_q326");
539        std::fs::create_dir(&bundle).unwrap();
540        write_tiny_q4_bundle(&bundle).unwrap();
541        let path = CString::new(bundle.to_str().unwrap()).unwrap();
542        let model = aria_model_init(path.as_ptr());
543        assert!(!model.is_null(), "{:?}", unsafe {
544            CStr::from_ptr(aria_last_error()).to_string_lossy()
545        });
546        aria_model_destroy(model);
547    }
548
549    #[test]
550    fn init_missing_path() {
551        let path = CString::new("/no/such/bundle").unwrap();
552        let model = aria_model_init(path.as_ptr());
553        assert!(model.is_null());
554        assert!(!aria_last_error().is_null());
555    }
556
557    #[test]
558    fn complete_bad_json() {
559        let dir = tempfile::tempdir().unwrap();
560        write_tiny_q4_bundle(dir.path()).unwrap();
561        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
562        let model = aria_model_init(path.as_ptr());
563        let bad = CString::new("not-json").unwrap();
564        let mut buf = out_buf();
565        assert_ne!(
566            aria_complete(
567                model,
568                bad.as_ptr(),
569                ptr::null(),
570                ptr::null(),
571                buf.as_mut_ptr() as *mut c_char,
572                buf.len()
573            ),
574            0
575        );
576        aria_model_destroy(model);
577    }
578
579    static CHUNKS: AtomicUsize = AtomicUsize::new(0);
580
581    unsafe extern "C" fn on_chunk(_s: *const c_char, _ud: *mut c_void) {
582        CHUNKS.fetch_add(1, Ordering::SeqCst);
583    }
584
585    #[test]
586    fn complete_stream_ok() {
587        let dir = tempfile::tempdir().unwrap();
588        write_tiny_q4_bundle(dir.path()).unwrap();
589        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
590        let model = aria_model_init(path.as_ptr());
591        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
592        let options = CString::new(r#"{"max_tokens":2}"#).unwrap();
593        let mut buf = out_buf();
594        CHUNKS.store(0, Ordering::SeqCst);
595        assert_eq!(
596            aria_complete_stream(
597                model,
598                messages.as_ptr(),
599                options.as_ptr(),
600                ptr::null(),
601                buf.as_mut_ptr() as *mut c_char,
602                buf.len(),
603                Some(on_chunk),
604                ptr::null_mut(),
605            ),
606            0
607        );
608        assert!(CHUNKS.load(Ordering::SeqCst) >= 1);
609        aria_model_destroy(model);
610    }
611
612    #[test]
613    fn destroy_null_and_use_after_destroy() {
614        aria_model_destroy(ptr::null_mut());
615        let dir = tempfile::tempdir().unwrap();
616        write_tiny_q4_bundle(dir.path()).unwrap();
617        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
618        let model = aria_model_init(path.as_ptr());
619        aria_model_destroy(model);
620        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
621        let mut buf = out_buf();
622        // After destroy the pointer must not be reused by callers; we only check null path.
623        assert_ne!(
624            aria_complete(
625                ptr::null_mut(),
626                messages.as_ptr(),
627                ptr::null(),
628                ptr::null(),
629                buf.as_mut_ptr() as *mut c_char,
630                buf.len()
631            ),
632            0
633        );
634    }
635
636    #[test]
637    fn cache_dir_uses_aria_compute_home() {
638        let tmp = tempfile::tempdir().unwrap();
639        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
640        let model = CString::new("gemma-4-e2b-it_q4").unwrap();
641        let dir = aria_model_cache_dir(model.as_ptr());
642        assert!(!dir.is_null());
643        let s = unsafe { CStr::from_ptr(dir) }.to_str().unwrap();
644        assert_eq!(
645            s,
646            tmp.path().join("models").join("gemma-4-e2b-it_q4").to_str().unwrap()
647        );
648        std::env::remove_var("ARIA_COMPUTE_HOME");
649    }
650
651    #[test]
652    fn is_local_path_detects_separator_and_existing() {
653        assert_eq!(aria_is_local_path(CString::new("/abs/path").unwrap().as_ptr()), 1);
654        assert_eq!(aria_is_local_path(CString::new("C:\\win\\path").unwrap().as_ptr()), 1);
655        assert_eq!(aria_is_local_path(CString::new("model_name").unwrap().as_ptr()), 0);
656        // an existing path on disk is treated as local
657        let tmp = tempfile::tempdir().unwrap();
658        assert_eq!(
659            aria_is_local_path(CString::new(tmp.path().to_str().unwrap()).unwrap().as_ptr()),
660            1
661        );
662    }
663}
aria_ffi/lib.rs

aria_ffi/
lib.rs
