1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ModelAvailability {
17 Available,
19 DeviceNotEligible,
21 AppleIntelligenceNotEnabled,
23 ModelNotReady,
25 Unknown,
27}
28
29impl ModelAvailability {
30 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
58pub struct SystemLanguageModel {
75 ptr: NonNull<std::ffi::c_void>,
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct TokenUsage {
81 pub token_count: usize,
83}
84
85impl SystemLanguageModel {
86 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 pub(crate) fn as_ptr(&self) -> SwiftPtr {
114 self.ptr.as_ptr()
115 }
116
117 pub fn is_available(&self) -> bool {
121 unsafe { ffi::fm_model_is_available(self.ptr.as_ptr()) }
122 }
123
124 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 pub fn ensure_available(&self) -> Result<()> {
134 match self.availability().into_error() {
135 Some(err) => Err(err),
136 None => Ok(()),
137 }
138 }
139
140 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 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
228unsafe 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
252pub(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 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 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}