Skip to main content

fm_rs/
model.rs

1//! `SystemLanguageModel`, `TokenUsage`, and `ModelAvailability` types.
2
3use std::ffi::{CStr, CString};
4use std::ptr::{self, NonNull};
5use std::sync::Arc;
6
7use crate::error::{Error, Result};
8use crate::ffi::{self, AvailabilityCode, SwiftPtr};
9use crate::tool::{Tool, tools_to_json};
10
11const TOKEN_USAGE_UNAVAILABLE_SENTINEL: i64 = -2;
12const TOKEN_ESTIMATE_CHARS_PER_TOKEN: usize = 4;
13
14/// Represents the availability status of a `FoundationModel`.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ModelAvailability {
17    /// Model is available and ready to use.
18    Available,
19    /// Device is not eligible for Apple Intelligence.
20    DeviceNotEligible,
21    /// Apple Intelligence is not enabled in system settings.
22    AppleIntelligenceNotEnabled,
23    /// Model is not ready (downloading or other system reasons).
24    ModelNotReady,
25    /// Unavailability for an unknown reason.
26    Unknown,
27}
28
29impl ModelAvailability {
30    /// Returns an error describing why the model is unavailable.
31    pub fn into_error(self) -> Option<Error> {
32        match self {
33            ModelAvailability::Available => None,
34            ModelAvailability::DeviceNotEligible => Some(Error::DeviceNotEligible),
35            ModelAvailability::AppleIntelligenceNotEnabled => {
36                Some(Error::AppleIntelligenceNotEnabled)
37            }
38            ModelAvailability::ModelNotReady => Some(Error::ModelNotReady),
39            ModelAvailability::Unknown => Some(Error::ModelNotAvailable),
40        }
41    }
42}
43
44impl From<AvailabilityCode> for ModelAvailability {
45    fn from(code: AvailabilityCode) -> Self {
46        match code {
47            AvailabilityCode::Available => ModelAvailability::Available,
48            AvailabilityCode::DeviceNotEligible => ModelAvailability::DeviceNotEligible,
49            AvailabilityCode::AppleIntelligenceNotEnabled => {
50                ModelAvailability::AppleIntelligenceNotEnabled
51            }
52            AvailabilityCode::ModelNotReady => ModelAvailability::ModelNotReady,
53            AvailabilityCode::Unknown => ModelAvailability::Unknown,
54        }
55    }
56}
57
58/// The system language model provided by Apple Intelligence.
59///
60/// This is the main entry point for using on-device AI capabilities.
61/// Use [`SystemLanguageModel::new()`] to get the default model.
62///
63/// # Example
64///
65/// ```rust,no_run
66/// use fm_rs::SystemLanguageModel;
67///
68/// let model = SystemLanguageModel::new()?;
69/// if model.is_available() {
70///     println!("Model is ready to use!");
71/// }
72/// # Ok::<(), fm_rs::Error>(())
73/// ```
74pub struct SystemLanguageModel {
75    ptr: NonNull<std::ffi::c_void>,
76}
77
78/// Token usage returned by `SystemLanguageModel` 26.4+ APIs.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct TokenUsage {
81    /// Number of tokens reported by the framework.
82    pub token_count: usize,
83}
84
85impl SystemLanguageModel {
86    /// Creates the default system language model.
87    ///
88    /// # Errors
89    ///
90    /// Returns an error if the model cannot be created or if `FoundationModels`
91    /// is not available on the device.
92    pub fn new() -> Result<Self> {
93        let mut error: SwiftPtr = ptr::null_mut();
94
95        let ptr = unsafe { ffi::fm_model_default(&raw mut error) };
96
97        if !error.is_null() {
98            return Err(error_from_swift(error));
99        }
100
101        NonNull::new(ptr).map(|ptr| Self { ptr }).ok_or_else(|| {
102            Error::InternalError(
103                "SystemLanguageModel creation returned null without error. \
104                 This may indicate FoundationModels.framework is unavailable."
105                    .to_string(),
106            )
107        })
108    }
109
110    /// Returns a raw pointer to the underlying Swift object.
111    ///
112    /// This is used internally for FFI calls.
113    pub(crate) fn as_ptr(&self) -> SwiftPtr {
114        self.ptr.as_ptr()
115    }
116
117    /// Checks if the model is available for use.
118    ///
119    /// Returns `true` if the model is available and ready to generate responses.
120    pub fn is_available(&self) -> bool {
121        unsafe { ffi::fm_model_is_available(self.ptr.as_ptr()) }
122    }
123
124    /// Gets the current availability status of the model.
125    ///
126    /// This provides more detailed information about why the model might not be available.
127    pub fn availability(&self) -> ModelAvailability {
128        let code = unsafe { ffi::fm_model_availability(self.ptr.as_ptr()) };
129        AvailabilityCode::from(code).into()
130    }
131
132    /// Returns a reason-specific error if the model is unavailable.
133    pub fn ensure_available(&self) -> Result<()> {
134        match self.availability().into_error() {
135            Some(err) => Err(err),
136            None => Ok(()),
137        }
138    }
139
140    /// Returns token usage for a prompt.
141    ///
142    /// Uses platform token-usage APIs when available in both the build SDK and runtime.
143    /// Otherwise returns a heuristic estimate.
144    pub fn token_usage_for(&self, prompt: &str) -> Result<TokenUsage> {
145        let prompt_c = CString::new(prompt)?;
146        let mut error: SwiftPtr = ptr::null_mut();
147
148        let token_count = unsafe {
149            ffi::fm_model_token_usage_for(self.ptr.as_ptr(), prompt_c.as_ptr(), &raw mut error)
150        };
151
152        if !error.is_null() {
153            return Err(error_from_swift(error));
154        }
155
156        if token_count == TOKEN_USAGE_UNAVAILABLE_SENTINEL {
157            return Ok(TokenUsage {
158                token_count: estimate_tokens(prompt, TOKEN_ESTIMATE_CHARS_PER_TOKEN),
159            });
160        }
161
162        token_usage_from_raw(token_count)
163    }
164
165    /// Returns token usage for session instructions and tool definitions.
166    ///
167    /// Tool definitions are serialized from the Rust [`Tool`] trait objects.
168    /// Uses platform token-usage APIs when available in both the build SDK and runtime.
169    /// Otherwise returns a heuristic estimate.
170    pub fn token_usage_for_tools(
171        &self,
172        instructions: &str,
173        tools: &[Arc<dyn Tool>],
174    ) -> Result<TokenUsage> {
175        let instructions_c = CString::new(instructions)?;
176        let tools_json = if tools.is_empty() {
177            None
178        } else {
179            let tool_refs: Vec<&dyn Tool> = tools.iter().map(std::convert::AsRef::as_ref).collect();
180            Some(CString::new(tools_to_json(&tool_refs)?)?)
181        };
182        let tools_ptr = tools_json.as_ref().map_or(ptr::null(), |s| s.as_ptr());
183
184        let mut error: SwiftPtr = ptr::null_mut();
185        let token_count = unsafe {
186            ffi::fm_model_token_usage_for_tools(
187                self.ptr.as_ptr(),
188                instructions_c.as_ptr(),
189                tools_ptr,
190                &raw mut error,
191            )
192        };
193
194        if !error.is_null() {
195            return Err(error_from_swift(error));
196        }
197
198        if token_count == TOKEN_USAGE_UNAVAILABLE_SENTINEL {
199            let fallback = estimate_tokens(instructions, TOKEN_ESTIMATE_CHARS_PER_TOKEN)
200                + tools_json.as_ref().map_or(0, |json| {
201                    estimate_tokens(&json.to_string_lossy(), TOKEN_ESTIMATE_CHARS_PER_TOKEN)
202                });
203            return Ok(TokenUsage {
204                token_count: fallback,
205            });
206        }
207
208        token_usage_from_raw(token_count)
209    }
210}
211
212impl std::fmt::Debug for SystemLanguageModel {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        f.debug_struct("SystemLanguageModel")
215            .field("availability", &self.availability())
216            .finish()
217    }
218}
219
220impl Drop for SystemLanguageModel {
221    fn drop(&mut self) {
222        unsafe {
223            ffi::fm_model_free(self.ptr.as_ptr());
224        }
225    }
226}
227
228// SAFETY: SystemLanguageModel is a wrapper around a Swift object that is
229// internally thread-safe (uses DispatchQueue for async operations).
230unsafe impl Send for SystemLanguageModel {}
231unsafe impl Sync for SystemLanguageModel {}
232
233fn token_usage_from_raw(token_count: i64) -> Result<TokenUsage> {
234    if token_count < 0 {
235        return Err(Error::InternalError(
236            "Token usage API returned a negative token count".to_string(),
237        ));
238    }
239
240    let token_count = usize::try_from(token_count)
241        .map_err(|_| Error::InternalError("Token usage value does not fit in usize".to_string()))?;
242
243    Ok(TokenUsage { token_count })
244}
245
246fn estimate_tokens(text: &str, chars_per_token: usize) -> usize {
247    let denom = chars_per_token.max(1);
248    let chars = text.chars().count();
249    chars.div_ceil(denom)
250}
251
252/// Converts a Swift error pointer to a Rust Error.
253pub(crate) fn error_from_swift(error: SwiftPtr) -> Error {
254    use crate::error::ToolCallError;
255
256    if error.is_null() {
257        return Error::InternalError(
258            "FFI error object was null; unable to retrieve error details".to_string(),
259        );
260    }
261
262    let code = unsafe { ffi::fm_error_code(error) };
263    let msg_ptr = unsafe { ffi::fm_error_message(error) };
264
265    let message = if msg_ptr.is_null() {
266        "Error message unavailable (null pointer from Swift)".to_string()
267    } else {
268        unsafe { CStr::from_ptr(msg_ptr).to_string_lossy().into_owned() }
269    };
270
271    // Extract tool context if this is a tool error
272    let tool_name = unsafe {
273        let ptr = ffi::fm_error_tool_name(error);
274        if ptr.is_null() {
275            None
276        } else {
277            Some(CStr::from_ptr(ptr).to_string_lossy().into_owned())
278        }
279    };
280
281    let tool_arguments = unsafe {
282        let ptr = ffi::fm_error_tool_arguments(error);
283        if ptr.is_null() {
284            None
285        } else {
286            let json_str = CStr::from_ptr(ptr).to_string_lossy().into_owned();
287            serde_json::from_str(&json_str).ok()
288        }
289    };
290
291    unsafe {
292        ffi::fm_error_free(error);
293    }
294
295    match ffi::ErrorCode::from(code) {
296        ffi::ErrorCode::ModelNotAvailable => Error::ModelNotAvailable,
297        ffi::ErrorCode::GenerationFailed => Error::GenerationError(message),
298        ffi::ErrorCode::Cancelled => Error::GenerationError("Operation cancelled".to_string()),
299        ffi::ErrorCode::Timeout => Error::Timeout(message),
300        ffi::ErrorCode::ToolError => {
301            // Construct ToolCallError with context if available
302            Error::ToolCall(ToolCallError {
303                tool_name: tool_name.unwrap_or_else(|| "unknown".to_string()),
304                arguments: tool_arguments.unwrap_or(serde_json::Value::Null),
305                inner_error: message,
306            })
307        }
308        ffi::ErrorCode::InvalidInput => Error::InvalidInput(message),
309        ffi::ErrorCode::Unknown => Error::InternalError(message),
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::{estimate_tokens, token_usage_from_raw};
316
317    #[test]
318    fn token_usage_should_convert_positive_values() {
319        let usage = token_usage_from_raw(42).expect("positive token count should convert");
320        assert_eq!(usage.token_count, 42);
321    }
322
323    #[test]
324    fn token_usage_should_reject_negative_values() {
325        let err = token_usage_from_raw(-1).expect_err("negative token count should fail");
326        assert!(err.to_string().contains("negative token count"));
327    }
328
329    #[test]
330    fn estimate_tokens_should_use_div_ceil() {
331        assert_eq!(estimate_tokens("abcd", 4), 1);
332        assert_eq!(estimate_tokens("abcde", 4), 2);
333    }
334}