ik_llama_cpp_2/model/chat.rs
1//! Chat templates (`llama_chat_apply_template`).
2//!
3//! Mirrors the `llama-cpp-2` anchor's chat-template API ([`LlamaChatMessage`],
4//! [`LlamaChatTemplate`], [`LlamaModel::chat_template`],
5//! [`LlamaModel::apply_chat_template`]), adapted to ik_llama.cpp's C signatures.
6//!
7//! [`LlamaModel`]: crate::model::LlamaModel
8
9use std::ffi::{CStr, CString, NulError};
10use std::os::raw::c_char;
11use std::str::Utf8Error;
12use std::string::FromUtf8Error;
13
14use ik_llama_cpp_sys as sys;
15
16/// A performance-friendly wrapper around [`LlamaModel::chat_template`] which is then
17/// fed into [`LlamaModel::apply_chat_template`] to convert a list of messages into an LLM
18/// prompt. Internally the template is stored as a `CString` to avoid round-trip conversions
19/// within the FFI.
20///
21/// [`LlamaModel::chat_template`]: crate::model::LlamaModel::chat_template
22/// [`LlamaModel::apply_chat_template`]: crate::model::LlamaModel::apply_chat_template
23#[derive(Eq, PartialEq, Clone, PartialOrd, Ord, Hash)]
24pub struct LlamaChatTemplate(CString);
25
26impl LlamaChatTemplate {
27 /// Create a new template from a string. This can either be the name of a built-in
28 /// chat template like "chatml" or "llama3" or an actual Jinja template for
29 /// ik_llama.cpp to interpret.
30 ///
31 /// # Errors
32 /// If `template` contains a null byte and thus cannot be converted to a C string.
33 pub fn new(template: &str) -> Result<Self, NulError> {
34 Ok(Self(CString::new(template)?))
35 }
36
37 /// Accesses the template as a c string reference.
38 #[must_use]
39 pub fn as_c_str(&self) -> &CStr {
40 &self.0
41 }
42
43 /// Attempts to convert the `CString` into a Rust str reference.
44 ///
45 /// # Errors
46 /// If the template is not valid UTF-8.
47 pub fn to_str(&self) -> Result<&str, Utf8Error> {
48 self.0.to_str()
49 }
50
51 /// Convenience method to create an owned String.
52 ///
53 /// # Errors
54 /// If the template is not valid UTF-8.
55 pub fn to_string(&self) -> Result<String, Utf8Error> {
56 self.to_str().map(str::to_string)
57 }
58}
59
60impl std::fmt::Debug for LlamaChatTemplate {
61 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62 self.0.fmt(f)
63 }
64}
65
66/// A safe wrapper around `llama_chat_message`.
67///
68/// Owns its `role` and `content` as `CString`s so the borrowed pointers handed to
69/// the FFI stay valid for the duration of the call.
70#[derive(Debug, Eq, PartialEq, Clone)]
71pub struct LlamaChatMessage {
72 role: CString,
73 content: CString,
74}
75
76impl LlamaChatMessage {
77 /// Create a new `LlamaChatMessage`.
78 ///
79 /// # Errors
80 /// If either of `role` or `content` contain null bytes.
81 pub fn new(role: String, content: String) -> Result<Self, NewLlamaChatMessageError> {
82 Ok(Self {
83 role: CString::new(role)?,
84 content: CString::new(content)?,
85 })
86 }
87}
88
89/// Failed to construct a [`LlamaChatMessage`].
90#[derive(Debug, thiserror::Error)]
91pub enum NewLlamaChatMessageError {
92 /// The string contained a null byte and thus could not be converted to a c string.
93 #[error("{0}")]
94 NulError(#[from] NulError),
95}
96
97/// Failed to apply a chat template (see [`LlamaModel::apply_chat_template`]).
98///
99/// [`LlamaModel::apply_chat_template`]: crate::model::LlamaModel::apply_chat_template
100#[derive(Debug, thiserror::Error)]
101pub enum ApplyChatTemplateError {
102 /// The string contained a null byte and thus could not be converted to a c string.
103 #[error("{0}")]
104 NulError(#[from] NulError),
105 /// The formatted prompt could not be converted to utf8.
106 #[error("{0}")]
107 FromUtf8Error(#[from] FromUtf8Error),
108 /// `llama_chat_apply_template` returned an error code.
109 #[error("ffi error {0}")]
110 FfiError(i32),
111}
112
113impl crate::model::LlamaModel {
114 /// Get the chat template from the model by name. If the `name` parameter is `None`,
115 /// the model's default chat template will be returned.
116 ///
117 /// You supply this into [`Self::apply_chat_template`] to get back a string with the
118 /// appropriate template substitution applied to convert a list of messages into a
119 /// prompt the LLM can use to complete the chat.
120 ///
121 /// Returns `None` if the model has no chat template by that name (or if `name`
122 /// contains an interior null byte).
123 ///
124 /// You could also use an external jinja parser, like [minijinja](https://github.com/mitsuhiko/minijinja),
125 /// to parse jinja templates not supported by the ik_llama.cpp template engine.
126 #[must_use]
127 pub fn chat_template(&self, name: Option<&str>) -> Option<LlamaChatTemplate> {
128 // Keep the CString alive for the duration of the FFI call.
129 let name_cstr = match name {
130 Some(name) => match CString::new(name) {
131 Ok(name) => Some(name),
132 Err(_) => return None,
133 },
134 None => None,
135 };
136 let name_ptr = name_cstr.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
137
138 // SAFETY: `self.model` is a valid model pointer; `name_ptr` is either null or a
139 // valid, NUL-terminated C string.
140 let result = unsafe { sys::llama_model_chat_template(self.model.as_ptr(), name_ptr) };
141 if result.is_null() {
142 return None;
143 }
144
145 // SAFETY: `result` is a valid, NUL-terminated C string owned by the model.
146 let bytes = unsafe { CStr::from_ptr(result) }.to_bytes();
147 // `bytes` has no interior NUL (it came from a C string), so this cannot fail;
148 // `.ok()` keeps this `unwrap`-free.
149 CString::new(bytes).ok().map(LlamaChatTemplate)
150 }
151
152 /// Apply a chat template to a list of messages, returning the formatted prompt.
153 ///
154 /// Inspired by hf `apply_chat_template()` on python. Use [`Self::chat_template`] to
155 /// retrieve the template baked into the model (this is the preferred path), or build
156 /// one directly with [`LlamaChatTemplate::new`].
157 ///
158 /// `add_assistant` controls whether the prompt ends with the token(s) that indicate
159 /// the start of an assistant message.
160 ///
161 /// # Errors
162 /// There are many ways this can fail. See [`ApplyChatTemplateError`] for more information.
163 pub fn apply_chat_template(
164 &self,
165 tmpl: &LlamaChatTemplate,
166 chat: &[LlamaChatMessage],
167 add_assistant: bool,
168 ) -> Result<String, ApplyChatTemplateError> {
169 // Build the ik_llama_cpp_sys chat messages. These borrow the `CString`s owned by
170 // `chat`, which outlive every FFI call below.
171 let c_chat: Vec<sys::llama_chat_message> = chat
172 .iter()
173 .map(|c| sys::llama_chat_message {
174 role: c.role.as_ptr(),
175 content: c.content.as_ptr(),
176 })
177 .collect();
178 let n_msg = c_chat.len();
179 let tmpl_ptr = tmpl.0.as_ptr();
180
181 // Recommended starting size is 2 * total message bytes (per the C API docs).
182 let hint = chat
183 .iter()
184 .fold(0usize, |acc, c| {
185 acc + c.role.as_bytes().len() + c.content.as_bytes().len()
186 })
187 .saturating_mul(2);
188 let mut cap = i32::try_from(hint).unwrap_or(i32::MAX);
189 let mut buf = vec![0u8; cap.max(0) as usize];
190
191 // SAFETY: `buf` holds `cap` bytes and `c_chat` holds `n_msg` messages, all
192 // borrowed pointers stay valid across the call.
193 let mut n = unsafe {
194 sys::llama_chat_apply_template(
195 tmpl_ptr,
196 c_chat.as_ptr(),
197 n_msg,
198 add_assistant,
199 buf.as_mut_ptr().cast::<c_char>(),
200 cap,
201 )
202 };
203 if n < 0 {
204 return Err(ApplyChatTemplateError::FfiError(n));
205 }
206
207 // ik returns the total required length; if it exceeds the buffer, re-alloc and
208 // re-apply (two-call buffer sizing, like tokenize).
209 if n > cap {
210 cap = n;
211 buf = vec![0u8; cap as usize];
212 // SAFETY: as above, now with a buffer sized to the required length.
213 n = unsafe {
214 sys::llama_chat_apply_template(
215 tmpl_ptr,
216 c_chat.as_ptr(),
217 n_msg,
218 add_assistant,
219 buf.as_mut_ptr().cast::<c_char>(),
220 cap,
221 )
222 };
223 if n < 0 {
224 return Err(ApplyChatTemplateError::FfiError(n));
225 }
226 }
227
228 buf.truncate(n as usize);
229 Ok(String::from_utf8(buf)?)
230 }
231}