#![cfg(feature = "needle")]
use core::ffi::{CStr, c_char, c_int, c_ulonglong};
use std::ffi::CString;
use cactus_sys::{
NEEDLE_TEXT, needle_complete, needle_embed, needle_init, needle_last_error, needle_load,
needle_models,
};
const OUT: usize = 256 * 1024;
fn archive(var: &str) -> Option<Vec<u8>> {
let Some(path) = std::env::var_os(var) else {
println!("skipping: set {var} to run this probe");
return None;
};
Some(std::fs::read(&path).unwrap_or_else(|error| panic!("{var}: {error}")))
}
unsafe fn last_error() -> String {
let ptr = unsafe { needle_last_error() };
if ptr.is_null() {
return String::new();
}
unsafe { CStr::from_ptr(ptr) }
.to_string_lossy()
.into_owned()
}
unsafe fn load(label: &str, bytes: &[u8]) -> (c_int, String, c_int) {
let rc = unsafe { needle_load(bytes.as_ptr(), bytes.len() as c_ulonglong) };
let error = if rc < 0 {
unsafe { last_error() }
} else {
String::new()
};
let models = unsafe { needle_models() };
println!("{label}: needle_load rc {rc}, last_error {error:?}, needle_models {models}");
(rc, error, models)
}
const TOOLS: &CStr = c"[{\"name\":\"control_lights\",\"description\":\"Turn lights on or off in a room.\",\"parameters\":{\"type\":\"object\",\"properties\":{\"room\":{\"type\":\"string\"},\"action\":{\"type\":\"string\",\"enum\":[\"on\",\"off\"]}},\"required\":[\"room\",\"action\"]}}]";
unsafe fn init() -> (c_int, String) {
let rc = unsafe {
needle_init(
c"You control a smart home.".as_ptr(),
TOOLS.as_ptr(),
core::ptr::null(),
)
};
let error = if rc < 0 {
unsafe { last_error() }
} else {
String::new()
};
println!("needle_init: rc {rc}, last_error {error:?}");
(rc, error)
}
unsafe fn complete(label: &str, input: &str) -> (c_int, String) {
let input = CString::new(input).unwrap();
let mut out = vec![0u8; OUT];
let rc = unsafe {
needle_complete(
input.as_ptr(),
core::ptr::null(),
0,
64,
out.as_mut_ptr().cast::<c_char>(),
OUT as c_int,
)
};
let error = if rc < 0 {
unsafe { last_error() }
} else {
String::new()
};
let envelope = CStr::from_bytes_until_nul(&out)
.unwrap()
.to_string_lossy()
.into_owned();
println!("{label}: needle_complete rc {rc}, last_error {error:?}, envelope {envelope}");
(rc, envelope)
}
unsafe fn embed(label: &str, input: &str) -> Vec<f32> {
let input = CString::new(input).unwrap();
let count = unsafe {
needle_embed(
input.as_ptr(),
core::ptr::null(),
0,
core::ptr::null_mut(),
0,
)
};
assert!(count > 0, "embed count {count}: {}", unsafe {
last_error()
});
let mut out = vec![0f32; count as usize];
let rc = unsafe {
needle_embed(
input.as_ptr(),
core::ptr::null(),
0,
out.as_mut_ptr(),
count,
)
};
assert_eq!(rc, count, "embed: {}", unsafe { last_error() });
println!(
"{label}: needle_embed({input:?}) dim {count}, first 8 {:?}",
&out[..8]
);
out
}
#[test]
fn reloading_the_same_text_archive_drops_init() {
let Some(text) = archive("CACTUS_NEEDLE_WEIGHTS") else {
return;
};
unsafe {
assert_eq!(load("text", &text), (0, String::new(), NEEDLE_TEXT));
assert!(init().0 > 0);
assert!(complete("before", "turn on the kitchen lights").0 >= 0);
let before = embed("before", "kitchen");
let (rc, error, models) = load("same bytes again", &text);
let (complete_rc, envelope) =
complete("after reload, no re-init", "turn on the kitchen lights");
let survived = complete_rc >= 0;
println!("init survived the identical reload: {survived}");
let after = embed("after reload", "kitchen");
println!("embedding changed: {}", before != after);
assert_eq!(rc, 0, "reload: {error}");
assert_eq!(models, NEEDLE_TEXT);
assert!(!survived);
assert!(envelope.contains("needle_init not called"), "{envelope}");
assert_eq!(before, after);
}
}