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::ffi::{CStr, CString};
10use std::os::raw::{c_char, c_int, c_uchar, c_void};
11use std::ptr;
12use std::slice;
13
14thread_local! {
15    static LAST_ERROR: RefCell<Option<CString>> = const { RefCell::new(None) };
16}
17
18pub struct AriaModel {
19    session: Session,
20}
21
22fn set_error(msg: impl Into<String>) {
23    let s = CString::new(msg.into()).unwrap_or_else(|_| CString::new("error").unwrap());
24    LAST_ERROR.with(|e| *e.borrow_mut() = Some(s));
25}
26
27fn clear_error() {
28    LAST_ERROR.with(|e| *e.borrow_mut() = None);
29}
30
31fn cstr_to_str<'a>(p: *const c_char) -> Result<&'a str, String> {
32    if p.is_null() {
33        return Err("null string".into());
34    }
35    unsafe { CStr::from_ptr(p) }
36        .to_str()
37        .map_err(|e| e.to_string())
38}
39
40fn write_out(out: *mut c_char, out_len: usize, s: &str) -> c_int {
41    if out.is_null() || out_len == 0 {
42        set_error("null output buffer");
43        return -1;
44    }
45    let bytes = s.as_bytes();
46    if bytes.len() + 1 > out_len {
47        set_error(format!(
48            "output buffer too small: need {}, have {}",
49            bytes.len() + 1,
50            out_len
51        ));
52        return -1;
53    }
54    unsafe {
55        ptr::copy_nonoverlapping(bytes.as_ptr(), out.cast::<u8>(), bytes.len());
56        *out.add(bytes.len()) = 0;
57    }
58    0
59}
60
61fn parse_messages(messages_json: &str) -> Result<Vec<ChatTurn>, String> {
62    let v: Value = serde_json::from_str(messages_json).map_err(|e| e.to_string())?;
63    let arr = v
64        .as_array()
65        .ok_or_else(|| "messages must be a JSON array".to_string())?;
66    let mut turns = Vec::new();
67    for m in arr {
68        let role = m
69            .get("role")
70            .and_then(|x| x.as_str())
71            .unwrap_or("user")
72            .to_string();
73        let content = m
74            .get("content")
75            .and_then(|x| x.as_str())
76            .unwrap_or("")
77            .to_string();
78        turns.push(ChatTurn { role, content });
79    }
80    Ok(turns)
81}
82
83fn parse_options(options_json: Option<&str>) -> GenerateOpts {
84    let mut opts = GenerateOpts::default();
85    if let Some(raw) = options_json {
86        if let Ok(v) = serde_json::from_str::<Value>(raw) {
87            if let Some(n) = v.get("max_tokens").and_then(|x| x.as_u64()) {
88                opts.max_tokens = n as usize;
89            }
90            if let Some(t) = v.get("temperature").and_then(|x| x.as_f64()) {
91                opts.temperature = t as f32;
92            }
93        }
94    }
95    if opts.max_tokens == 0 {
96        opts.max_tokens = 16;
97    }
98    opts
99}
100
101fn parse_tools(tools_json: Option<&str>) -> Result<Value, String> {
102    match tools_json {
103        None | Some("") => Ok(json!([])),
104        Some(raw) => {
105            let v: Value = serde_json::from_str(raw).map_err(|e| e.to_string())?;
106            if !v.is_array() && !v.is_null() {
107                return Err("tools must be a JSON array or null".into());
108            }
109            Ok(if v.is_null() { json!([]) } else { v })
110        }
111    }
112}
113
114/// Opaque model handle.
115pub type AriaModelHandle = *mut AriaModel;
116
117/// Last error message (thread-local). Valid until next call on this thread.
118#[no_mangle]
119pub extern "C" fn aria_last_error() -> *const c_char {
120    LAST_ERROR.with(|e| match e.borrow().as_ref() {
121        Some(s) => s.as_ptr(),
122        None => ptr::null(),
123    })
124}
125
126/// Load an Aria quant bundle from `bundle_path`. Returns null on error.
127#[no_mangle]
128pub extern "C" fn aria_model_init(bundle_path: *const c_char) -> AriaModelHandle {
129    clear_error();
130    let path = match cstr_to_str(bundle_path) {
131        Ok(p) => p,
132        Err(e) => {
133            set_error(e);
134            return ptr::null_mut();
135        }
136    };
137    match SessionBuilder::new().model(path).build() {
138        Ok(session) => Box::into_raw(Box::new(AriaModel { session })),
139        Err(e) => {
140            set_error(e.to_string());
141            ptr::null_mut()
142        }
143    }
144}
145
146/// Destroy a model handle. Safe on null. Double-destroy is undefined if caller reuses the pointer.
147#[no_mangle]
148pub extern "C" fn aria_model_destroy(model: AriaModelHandle) {
149    clear_error();
150    if model.is_null() {
151        return;
152    }
153    unsafe {
154        drop(Box::from_raw(model));
155    }
156}
157
158fn complete_inner(
159    model: AriaModelHandle,
160    messages_json: *const c_char,
161    options_json: *const c_char,
162    tools_json: *const c_char,
163    out: *mut c_char,
164    out_len: usize,
165    stream_cb: Option<unsafe extern "C" fn(*const c_char, *mut c_void)>,
166    user_data: *mut c_void,
167) -> c_int {
168    clear_error();
169    if model.is_null() {
170        set_error("null model");
171        return -1;
172    }
173    let messages = match cstr_to_str(messages_json) {
174        Ok(s) => s,
175        Err(e) => {
176            set_error(e);
177            return -1;
178        }
179    };
180    let options = if options_json.is_null() {
181        None
182    } else {
183        match cstr_to_str(options_json) {
184            Ok(s) => Some(s),
185            Err(e) => {
186                set_error(e);
187                return -1;
188            }
189        }
190    };
191    let tools_raw = if tools_json.is_null() {
192        None
193    } else {
194        match cstr_to_str(tools_json) {
195            Ok(s) => Some(s),
196            Err(e) => {
197                set_error(e);
198                return -1;
199            }
200        }
201    };
202
203    let turns = match parse_messages(messages) {
204        Ok(p) => p,
205        Err(e) => {
206            set_error(e);
207            return -1;
208        }
209    };
210    let tools = match parse_tools(tools_raw) {
211        Ok(t) => t,
212        Err(e) => {
213            set_error(e);
214            return -1;
215        }
216    };
217    let opts = parse_options(options);
218    let m = unsafe { &mut *model };
219    let tokens = m.session.encode_chat(&turns);
220    let gen = match m.session.generate(&tokens, &opts) {
221        Ok(g) => g,
222        Err(e) => {
223            set_error(e.to_string());
224            return -1;
225        }
226    };
227
228    if let Some(cb) = stream_cb {
229        // Stream decoded text (same as gen.text), not raw `<id>` placeholders.
230        if let Ok(c) = CString::new(gen.text.as_str()) {
231            unsafe { cb(c.as_ptr(), user_data) };
232        }
233    }
234
235    let body = json!({
236        "success": true,
237        "error": null,
238        "response": gen.text,
239        "function_calls": json!([]),
240        "segments": [],
241        "cloud_handoff": false,
242        "total_tokens": gen.tokens.len(),
243    });
244    // tools accepted for OpenAI parity; real tool routing is stage C.
245    let _ = tools;
246    write_out(out, out_len, &body.to_string())
247}
248
249/// Non-streaming chat completion. `tools_json` may be null.
250#[no_mangle]
251pub extern "C" fn aria_complete(
252    model: AriaModelHandle,
253    messages_json: *const c_char,
254    options_json: *const c_char,
255    tools_json: *const c_char,
256    out: *mut c_char,
257    out_len: usize,
258) -> c_int {
259    complete_inner(
260        model,
261        messages_json,
262        options_json,
263        tools_json,
264        out,
265        out_len,
266        None,
267        ptr::null_mut(),
268    )
269}
270
271/// Streaming chat; `callback` receives each chunk as a C string.
272#[no_mangle]
273pub extern "C" fn aria_complete_stream(
274    model: AriaModelHandle,
275    messages_json: *const c_char,
276    options_json: *const c_char,
277    tools_json: *const c_char,
278    out: *mut c_char,
279    out_len: usize,
280    callback: Option<unsafe extern "C" fn(*const c_char, *mut c_void)>,
281    user_data: *mut c_void,
282) -> c_int {
283    complete_inner(
284        model,
285        messages_json,
286        options_json,
287        tools_json,
288        out,
289        out_len,
290        callback,
291        user_data,
292    )
293}
294
295/// Embeddings. `input_json` is a string or `{"input":"..."}` / `{"input":[...]}`.
296#[no_mangle]
297pub extern "C" fn aria_embed(
298    model: AriaModelHandle,
299    input_json: *const c_char,
300    out: *mut c_char,
301    out_len: usize,
302) -> c_int {
303    clear_error();
304    if model.is_null() {
305        set_error("null model");
306        return -1;
307    }
308    let raw = match cstr_to_str(input_json) {
309        Ok(s) => s,
310        Err(e) => {
311            set_error(e);
312            return -1;
313        }
314    };
315    let text = match serde_json::from_str::<Value>(raw) {
316        Ok(Value::String(s)) => s,
317        Ok(v) => v
318            .get("input")
319            .and_then(|x| {
320                x.as_str()
321                    .map(|s| s.to_string())
322                    .or_else(|| x.as_array().and_then(|a| a.first()).and_then(|x| x.as_str()).map(|s| s.to_string()))
323            })
324            .unwrap_or_default(),
325        Err(_) => raw.to_string(),
326    };
327    if text.is_empty() {
328        set_error("empty embedding input");
329        return -1;
330    }
331    let m = unsafe { &*model };
332    let emb = match m.session.embed_text(&text) {
333        Ok(e) => e,
334        Err(e) => {
335            set_error(e.to_string());
336            return -1;
337        }
338    };
339    let body = json!({
340        "object": "list",
341        "data": [{
342            "object": "embedding",
343            "embedding": emb,
344            "index": 0
345        }]
346    });
347    write_out(out, out_len, &body.to_string())
348}
349
350/// Transcribe PCM16 LE bytes.
351#[no_mangle]
352pub extern "C" fn aria_transcribe(
353    model: AriaModelHandle,
354    pcm: *const c_uchar,
355    pcm_len: usize,
356    _options_json: *const c_char,
357    out: *mut c_char,
358    out_len: usize,
359) -> c_int {
360    clear_error();
361    if model.is_null() {
362        set_error("null model");
363        return -1;
364    }
365    if pcm.is_null() || pcm_len == 0 {
366        set_error("empty pcm");
367        return -1;
368    }
369    let bytes = unsafe { slice::from_raw_parts(pcm, pcm_len) };
370    let m = unsafe { &*model };
371    let text = match m.session.transcribe_pcm16le(bytes) {
372        Ok(t) => t,
373        Err(e) => {
374            set_error(e.to_string());
375            return -1;
376        }
377    };
378    let body = json!({
379        "text": text,
380        "segments": []
381    });
382    write_out(out, out_len, &body.to_string())
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use aria_inference::fixture::write_tiny_q4_bundle;
389    use std::sync::atomic::{AtomicUsize, Ordering};
390
391    fn out_buf() -> Vec<u8> {
392        vec![0u8; 64 * 1024]
393    }
394
395    #[test]
396    fn init_complete_embed_transcribe_destroy() {
397        let dir = tempfile::tempdir().unwrap();
398        write_tiny_q4_bundle(dir.path()).unwrap();
399        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
400        let model = aria_model_init(path.as_ptr());
401        assert!(!model.is_null(), "{:?}", unsafe {
402            CStr::from_ptr(aria_last_error()).to_string_lossy()
403        });
404
405        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
406        let options = CString::new(r#"{"max_tokens":2}"#).unwrap();
407        let tools = CString::new("[]").unwrap();
408        let mut buf = out_buf();
409        assert_eq!(
410            aria_complete(
411                model,
412                messages.as_ptr(),
413                options.as_ptr(),
414                tools.as_ptr(),
415                buf.as_mut_ptr() as *mut c_char,
416                buf.len(),
417            ),
418            0
419        );
420        let s = unsafe { CStr::from_ptr(buf.as_ptr() as *const c_char) }
421            .to_str()
422            .unwrap();
423        let v: Value = serde_json::from_str(s).unwrap();
424        assert_eq!(v["success"], true);
425        assert!(!v["response"].as_str().unwrap().is_empty());
426
427        let input = CString::new(r#"{"input":"hello"}"#).unwrap();
428        buf.fill(0);
429        assert_eq!(
430            aria_embed(
431                model,
432                input.as_ptr(),
433                buf.as_mut_ptr() as *mut c_char,
434                buf.len()
435            ),
436            0
437        );
438
439        let pcm = [0u8, 1, 2, 3, 4, 5];
440        buf.fill(0);
441        assert_eq!(
442            aria_transcribe(
443                model,
444                pcm.as_ptr(),
445                pcm.len(),
446                ptr::null(),
447                buf.as_mut_ptr() as *mut c_char,
448                buf.len()
449            ),
450            0
451        );
452
453        aria_model_destroy(model);
454    }
455
456    #[test]
457    fn init_missing_path() {
458        let path = CString::new("/no/such/bundle").unwrap();
459        let model = aria_model_init(path.as_ptr());
460        assert!(model.is_null());
461        assert!(!aria_last_error().is_null());
462    }
463
464    #[test]
465    fn complete_bad_json() {
466        let dir = tempfile::tempdir().unwrap();
467        write_tiny_q4_bundle(dir.path()).unwrap();
468        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
469        let model = aria_model_init(path.as_ptr());
470        let bad = CString::new("not-json").unwrap();
471        let mut buf = out_buf();
472        assert_ne!(
473            aria_complete(
474                model,
475                bad.as_ptr(),
476                ptr::null(),
477                ptr::null(),
478                buf.as_mut_ptr() as *mut c_char,
479                buf.len()
480            ),
481            0
482        );
483        aria_model_destroy(model);
484    }
485
486    static CHUNKS: AtomicUsize = AtomicUsize::new(0);
487
488    unsafe extern "C" fn on_chunk(_s: *const c_char, _ud: *mut c_void) {
489        CHUNKS.fetch_add(1, Ordering::SeqCst);
490    }
491
492    #[test]
493    fn complete_stream_ok() {
494        let dir = tempfile::tempdir().unwrap();
495        write_tiny_q4_bundle(dir.path()).unwrap();
496        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
497        let model = aria_model_init(path.as_ptr());
498        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
499        let options = CString::new(r#"{"max_tokens":2}"#).unwrap();
500        let mut buf = out_buf();
501        CHUNKS.store(0, Ordering::SeqCst);
502        assert_eq!(
503            aria_complete_stream(
504                model,
505                messages.as_ptr(),
506                options.as_ptr(),
507                ptr::null(),
508                buf.as_mut_ptr() as *mut c_char,
509                buf.len(),
510                Some(on_chunk),
511                ptr::null_mut(),
512            ),
513            0
514        );
515        assert!(CHUNKS.load(Ordering::SeqCst) >= 1);
516        aria_model_destroy(model);
517    }
518
519    #[test]
520    fn destroy_null_and_use_after_destroy() {
521        aria_model_destroy(ptr::null_mut());
522        let dir = tempfile::tempdir().unwrap();
523        write_tiny_q4_bundle(dir.path()).unwrap();
524        let path = CString::new(dir.path().to_str().unwrap()).unwrap();
525        let model = aria_model_init(path.as_ptr());
526        aria_model_destroy(model);
527        let messages = CString::new(r#"[{"role":"user","content":"hi"}]"#).unwrap();
528        let mut buf = out_buf();
529        // After destroy the pointer must not be reused by callers; we only check null path.
530        assert_ne!(
531            aria_complete(
532                ptr::null_mut(),
533                messages.as_ptr(),
534                ptr::null(),
535                ptr::null(),
536                buf.as_mut_ptr() as *mut c_char,
537                buf.len()
538            ),
539            0
540        );
541    }
542}
aria_ffi/lib.rs

aria_ffi/
lib.rs
