llama_cpp_bindings/context/
session.rs

1use std::ffi::CString;
2use std::path::Path;
3
4use crate::context::LlamaContext;
5use crate::context::llama_state_seq_flags::LlamaStateSeqFlags;
6use crate::context::load_seq_state_error::LoadSeqStateError;
7use crate::context::load_session_error::LoadSessionError;
8use crate::context::save_seq_state_error::SaveSeqStateError;
9use crate::context::save_session_error::SaveSessionError;
10use crate::token::LlamaToken;
11
12fn process_session_load_result(
13    success: bool,
14    n_out: usize,
15    max_tokens: usize,
16    mut tokens: Vec<LlamaToken>,
17) -> Result<Vec<LlamaToken>, LoadSessionError> {
18    if !success {
19        return Err(LoadSessionError::FailedToLoad);
20    }
21
22    if n_out > max_tokens {
23        return Err(LoadSessionError::InsufficientMaxLength { n_out, max_tokens });
24    }
25
26    unsafe { tokens.set_len(n_out) };
27
28    Ok(tokens)
29}
30
31fn process_seq_load_result(
32    bytes_read: usize,
33    n_out: usize,
34    max_tokens: usize,
35    mut tokens: Vec<LlamaToken>,
36) -> Result<(Vec<LlamaToken>, usize), LoadSeqStateError> {
37    if bytes_read == 0 {
38        return Err(LoadSeqStateError::FailedToLoad);
39    }
40
41    if n_out > max_tokens {
42        return Err(LoadSeqStateError::InsufficientMaxLength { n_out, max_tokens });
43    }
44
45    unsafe { tokens.set_len(n_out) };
46
47    Ok((tokens, bytes_read))
48}
49
50impl LlamaContext<'_> {
51    /// # Errors
52    ///
53    /// Fails if the path is not a valid utf8 or llama.cpp fails to save the state file.
54    pub fn state_save_file(
55        &self,
56        path_session: impl AsRef<Path>,
57        tokens: &[LlamaToken],
58    ) -> Result<(), SaveSessionError> {
59        let path = path_session.as_ref();
60        let path = path
61            .to_str()
62            .ok_or_else(|| SaveSessionError::PathToStrError(path.to_path_buf()))?;
63
64        let cstr = CString::new(path)?;
65
66        if unsafe {
67            llama_cpp_bindings_sys::llama_state_save_file(
68                self.context.as_ptr(),
69                cstr.as_ptr(),
70                tokens
71                    .as_ptr()
72                    .cast::<llama_cpp_bindings_sys::llama_token>(),
73                tokens.len(),
74            )
75        } {
76            Ok(())
77        } else {
78            Err(SaveSessionError::FailedToSave)
79        }
80    }
81
82    /// # Errors
83    ///
84    /// Fails if the path is not a valid utf8 or llama.cpp fails to load the state file.
85    pub fn state_load_file(
86        &mut self,
87        path_session: impl AsRef<Path>,
88        max_tokens: usize,
89    ) -> Result<Vec<LlamaToken>, LoadSessionError> {
90        let path = path_session.as_ref();
91        let path = path
92            .to_str()
93            .ok_or_else(|| LoadSessionError::PathToStrError(path.to_path_buf()))?;
94
95        let cstr = CString::new(path)?;
96        let mut tokens: Vec<LlamaToken> = Vec::with_capacity(max_tokens);
97        let mut n_out = 0;
98
99        // SAFETY: cast is valid as LlamaToken is repr(transparent)
100        let tokens_out = tokens
101            .as_mut_ptr()
102            .cast::<llama_cpp_bindings_sys::llama_token>();
103
104        let success = unsafe {
105            llama_cpp_bindings_sys::llama_state_load_file(
106                self.context.as_ptr(),
107                cstr.as_ptr(),
108                tokens_out,
109                max_tokens,
110                &raw mut n_out,
111            )
112        };
113        process_session_load_result(success, n_out, max_tokens, tokens)
114    }
115
116    /// # Errors
117    ///
118    /// Fails if the path is not a valid utf8 or llama.cpp fails to save the sequence state file.
119    ///
120    pub fn state_seq_save_file(
121        &self,
122        filepath: impl AsRef<Path>,
123        seq_id: i32,
124        tokens: &[LlamaToken],
125    ) -> Result<usize, SaveSeqStateError> {
126        let path = filepath.as_ref();
127        let path = path
128            .to_str()
129            .ok_or_else(|| SaveSeqStateError::PathToStrError(path.to_path_buf()))?;
130
131        let cstr = CString::new(path)?;
132
133        let bytes_written = unsafe {
134            llama_cpp_bindings_sys::llama_state_seq_save_file(
135                self.context.as_ptr(),
136                cstr.as_ptr(),
137                seq_id,
138                tokens
139                    .as_ptr()
140                    .cast::<llama_cpp_bindings_sys::llama_token>(),
141                tokens.len(),
142            )
143        };
144
145        if bytes_written == 0 {
146            Err(SaveSeqStateError::FailedToSave)
147        } else {
148            Ok(bytes_written)
149        }
150    }
151
152    /// # Errors
153    ///
154    /// Fails if the path is not a valid utf8 or llama.cpp fails to load the sequence state file.
155    ///
156    pub fn state_seq_load_file(
157        &mut self,
158        filepath: impl AsRef<Path>,
159        dest_seq_id: i32,
160        max_tokens: usize,
161    ) -> Result<(Vec<LlamaToken>, usize), LoadSeqStateError> {
162        let path = filepath.as_ref();
163        let path = path
164            .to_str()
165            .ok_or_else(|| LoadSeqStateError::PathToStrError(path.to_path_buf()))?;
166
167        let cstr = CString::new(path)?;
168        let mut tokens: Vec<LlamaToken> = Vec::with_capacity(max_tokens);
169        let mut n_out = 0;
170
171        // SAFETY: cast is valid as LlamaToken is repr(transparent)
172        let tokens_out = tokens
173            .as_mut_ptr()
174            .cast::<llama_cpp_bindings_sys::llama_token>();
175
176        let bytes_read = unsafe {
177            llama_cpp_bindings_sys::llama_state_seq_load_file(
178                self.context.as_ptr(),
179                cstr.as_ptr(),
180                dest_seq_id,
181                tokens_out,
182                max_tokens,
183                &raw mut n_out,
184            )
185        };
186
187        process_seq_load_result(bytes_read, n_out, max_tokens, tokens)
188    }
189
190    #[must_use]
191    pub fn get_state_size(&self) -> usize {
192        unsafe { llama_cpp_bindings_sys::llama_state_get_size(self.context.as_ptr()) }
193    }
194
195    /// # Safety
196    ///
197    /// The `dest` buffer must be large enough to hold the complete state data.
198    pub unsafe fn copy_state_data(&self, dest: &mut [u8]) -> usize {
199        unsafe {
200            llama_cpp_bindings_sys::llama_state_get_data(
201                self.context.as_ptr(),
202                dest.as_mut_ptr(),
203                dest.len(),
204            )
205        }
206    }
207
208    /// # Safety
209    ///
210    /// The `src` buffer must contain data previously obtained from [`copy_state_data`](Self::copy_state_data)
211    /// on a compatible context (same model and parameters). Passing arbitrary or corrupted bytes
212    /// will lead to undefined behavior.
213    pub unsafe fn set_state_data(&mut self, src: &[u8]) -> usize {
214        unsafe {
215            llama_cpp_bindings_sys::llama_state_set_data(
216                self.context.as_ptr(),
217                src.as_ptr(),
218                src.len(),
219            )
220        }
221    }
222
223    #[must_use]
224    pub fn state_seq_get_size_ext(&self, seq_id: i32, flags: &LlamaStateSeqFlags) -> usize {
225        unsafe {
226            llama_cpp_bindings_sys::llama_state_seq_get_size_ext(
227                self.context.as_ptr(),
228                seq_id,
229                flags.bits(),
230            )
231        }
232    }
233
234    /// # Safety
235    ///
236    /// The `dest` buffer must be large enough to hold the complete state data.
237    pub unsafe fn state_seq_get_data_ext(
238        &self,
239        dest: &mut [u8],
240        seq_id: i32,
241        flags: &LlamaStateSeqFlags,
242    ) -> usize {
243        unsafe {
244            llama_cpp_bindings_sys::llama_state_seq_get_data_ext(
245                self.context.as_ptr(),
246                dest.as_mut_ptr(),
247                dest.len(),
248                seq_id,
249                flags.bits(),
250            )
251        }
252    }
253
254    /// # Safety
255    ///
256    /// The `src` buffer must contain data previously obtained from
257    /// [`state_seq_get_data_ext`](Self::state_seq_get_data_ext) on a compatible context.
258    pub unsafe fn state_seq_set_data_ext(
259        &mut self,
260        src: &[u8],
261        dest_seq_id: i32,
262        flags: &LlamaStateSeqFlags,
263    ) -> usize {
264        unsafe {
265            llama_cpp_bindings_sys::llama_state_seq_set_data_ext(
266                self.context.as_ptr(),
267                src.as_ptr(),
268                src.len(),
269                dest_seq_id,
270                flags.bits(),
271            )
272        }
273    }
274}
275
276#[cfg(test)]
277mod unit_tests {
278    use crate::token::LlamaToken;
279
280    use crate::context::load_seq_state_error::LoadSeqStateError;
281    use crate::context::load_session_error::LoadSessionError;
282
283    use super::{process_seq_load_result, process_session_load_result};
284
285    #[test]
286    fn session_load_success_within_bounds() {
287        let tokens = vec![LlamaToken::new(0); 100];
288        let result = process_session_load_result(true, 10, 100, tokens);
289
290        assert!(result.is_ok());
291        assert_eq!(result.unwrap().len(), 10);
292    }
293
294    #[test]
295    fn session_load_fails_when_not_successful() {
296        let tokens = vec![LlamaToken::new(0); 100];
297        let result = process_session_load_result(false, 0, 100, tokens);
298
299        assert_eq!(result, Err(LoadSessionError::FailedToLoad));
300    }
301
302    #[test]
303    fn session_load_fails_when_n_out_exceeds_max() {
304        let tokens = vec![LlamaToken::new(0); 100];
305        let result = process_session_load_result(true, 101, 100, tokens);
306
307        assert_eq!(
308            result,
309            Err(LoadSessionError::InsufficientMaxLength {
310                n_out: 101,
311                max_tokens: 100,
312            })
313        );
314    }
315
316    #[test]
317    fn seq_load_success_within_bounds() {
318        let tokens = vec![LlamaToken::new(0); 100];
319        let result = process_seq_load_result(42, 10, 100, tokens);
320
321        assert!(result.is_ok());
322        let (loaded, bytes) = result.unwrap();
323        assert_eq!(loaded.len(), 10);
324        assert_eq!(bytes, 42);
325    }
326
327    #[test]
328    fn seq_load_fails_when_zero_bytes_read() {
329        let tokens = vec![LlamaToken::new(0); 100];
330        let result = process_seq_load_result(0, 0, 100, tokens);
331
332        assert_eq!(result, Err(LoadSeqStateError::FailedToLoad));
333    }
334
335    #[test]
336    fn seq_load_fails_when_n_out_exceeds_max() {
337        let tokens = vec![LlamaToken::new(0); 100];
338        let result = process_seq_load_result(42, 101, 100, tokens);
339
340        assert_eq!(
341            result,
342            Err(LoadSeqStateError::InsufficientMaxLength {
343                n_out: 101,
344                max_tokens: 100,
345            })
346        );
347    }
348}
llama_cpp_bindings/context/session.rs

llama_cpp_bindings/context/
session.rs
