rlx_text/
chat.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Chat-template engine for RLX runners.
17//!
18//! Replaces `LlamaModel::apply_chat_template` (llama-cpp-4) end-to-end. Two
19//! sources: an inline Jinja2 string, or `tokenizer.chat_template` (and
20//! `tokenizer.ggml.chat_template`) read directly from a GGUF file's
21//! metadata. Rendering uses `minijinja`.
22//!
23//! BOS/EOS strings are looked up via `tokenizer.ggml.bos_token_id` /
24//! `eos_token_id` against the `tokenizer.ggml.tokens` array (the GGUF
25//! convention).
26
27use anyhow::{Context, Result, anyhow};
28use minijinja::value::Object;
29use minijinja::{Environment, Error as JinjaError, ErrorKind, State, Value};
30use rlx_gguf::{GgufFile, MetaValue};
31use serde_json::Value as JsonValue;
32use std::path::Path;
33use std::sync::Arc;
34
35/// Convenience for the M3 auto-dispatch family: load the chat template
36/// + BOS/EOS strings directly from a GGUF path.
37///
38/// Alias for [`ChatTemplate::from_gguf`]. Use `rlx_models::run::auto_chat_template(path)`
39/// next to `rlx_models::run::auto_runner(path)`.
40pub fn auto_chat_template(path: &Path) -> Result<ChatTemplate> {
41    ChatTemplate::from_gguf(path)
42}
43
44/// One chat turn. `role` is conventionally one of `system`, `user`,
45/// `assistant`, `tool` — but templates can accept anything.
46#[derive(Debug, Clone)]
47pub struct ChatMessage {
48    pub role: String,
49    pub content: String,
50}
51
52/// Extra Jinja variables for templates that need more than the ChatML
53/// baseline (Gemma 4 thinking channel, tool schemas, …).
54#[derive(Debug, Clone, Copy)]
55pub struct ChatRenderOptions {
56    pub add_generation_prompt: bool,
57    /// Gemma 4 unified templates gate the `<|think|>` prefix and thought
58    /// channel on this flag (HF `enable_thinking`, default true for IT).
59    pub enable_thinking: bool,
60}
61
62impl Default for ChatRenderOptions {
63    fn default() -> Self {
64        Self {
65            add_generation_prompt: true,
66            enable_thinking: false,
67        }
68    }
69}
70
71impl ChatRenderOptions {
72    pub fn user_turn(add_generation_prompt: bool) -> Self {
73        Self {
74            add_generation_prompt,
75            ..Self::default()
76        }
77    }
78
79    pub fn gemma4_thinking(add_generation_prompt: bool) -> Self {
80        Self {
81            add_generation_prompt,
82            enable_thinking: true,
83        }
84    }
85}
86
87/// HF chat templates call `.get(key)` / `.get(key, default)` on message
88/// dicts. minijinja maps from `serde_json` do not expose that method —
89/// wrap them so Gemma 4 / tool-use templates render.
90#[derive(Debug, Clone)]
91struct GettableValue(JsonValue);
92
93impl GettableValue {
94    fn from_json(v: JsonValue) -> Value {
95        match v {
96            JsonValue::Object(_) | JsonValue::Array(_) => Value::from_object(Self(v)),
97            JsonValue::String(s) => Value::from(s),
98            JsonValue::Number(n) => {
99                if let Some(i) = n.as_i64() {
100                    Value::from(i)
101                } else if let Some(u) = n.as_u64() {
102                    Value::from(u)
103                } else if let Some(f) = n.as_f64() {
104                    Value::from(f)
105                } else {
106                    Value::from(n.to_string())
107                }
108            }
109            JsonValue::Bool(b) => Value::from(b),
110            JsonValue::Null => Value::from(()),
111        }
112    }
113}
114
115impl Object for GettableValue {
116    fn get_value(self: &Arc<Self>, key: &Value) -> Option<Value> {
117        match &self.0 {
118            JsonValue::Object(map) => {
119                let k = key.as_str()?;
120                map.get(k).map(|v| GettableValue::from_json(v.clone()))
121            }
122            JsonValue::Array(items) => {
123                let idx = key.as_usize()?;
124                items.get(idx).map(|v| GettableValue::from_json(v.clone()))
125            }
126            _ => None,
127        }
128    }
129
130    fn call_method(
131        self: &Arc<Self>,
132        _state: &State<'_, '_>,
133        name: &str,
134        args: &[Value],
135    ) -> Result<Value, JinjaError> {
136        if name != "get" {
137            return Err(JinjaError::new(
138                ErrorKind::UnknownMethod,
139                format!("GettableValue has no method named {name}"),
140            ));
141        }
142        let key = args
143            .first()
144            .and_then(|v| v.as_str())
145            .ok_or_else(|| JinjaError::new(ErrorKind::InvalidOperation, "get() needs a key"))?;
146        let default = args.get(1).cloned().unwrap_or(Value::UNDEFINED);
147        Ok(match &self.0 {
148            JsonValue::Object(map) => map
149                .get(key)
150                .map(|v| GettableValue::from_json(v.clone()))
151                .unwrap_or(default),
152            _ => default,
153        })
154    }
155
156    fn render(self: &Arc<Self>, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157        match &self.0 {
158            JsonValue::String(s) => write!(f, "{s}"),
159            JsonValue::Number(n) => write!(f, "{n}"),
160            JsonValue::Bool(b) => write!(f, "{b}"),
161            JsonValue::Null => write!(f, "null"),
162            JsonValue::Array(_) => write!(f, "[...]"),
163            JsonValue::Object(_) => write!(f, "{{...}}"),
164        }
165    }
166}
167
168fn chat_messages_to_values(messages: &[ChatMessage]) -> Vec<Value> {
169    messages
170        .iter()
171        .map(|m| {
172            GettableValue::from_json(serde_json::json!({
173                "role": m.role,
174                "content": m.content,
175            }))
176        })
177        .collect()
178}
179
180impl ChatMessage {
181    pub fn user(content: impl Into<String>) -> Self {
182        Self {
183            role: "user".into(),
184            content: content.into(),
185        }
186    }
187    pub fn system(content: impl Into<String>) -> Self {
188        Self {
189            role: "system".into(),
190            content: content.into(),
191        }
192    }
193    pub fn assistant(content: impl Into<String>) -> Self {
194        Self {
195            role: "assistant".into(),
196            content: content.into(),
197        }
198    }
199}
200
201/// Where a [`ChatTemplate`] was loaded from. Useful for diagnostics and
202/// for letting a caller round-trip the source string into config.
203#[derive(Debug, Clone)]
204pub enum ChatTemplateSource {
205    Inline,
206    GgufMetadata(String),
207}
208
209/// Compiled Jinja chat template + BOS/EOS strings.
210pub struct ChatTemplate {
211    env: Environment<'static>,
212    source_text: String,
213    source_kind: ChatTemplateSource,
214    bos_token: Option<String>,
215    eos_token: Option<String>,
216}
217
218const TEMPLATE_NAME: &str = "chat";
219
220fn build_env(source: String) -> Result<Environment<'static>> {
221    let mut env = Environment::new();
222    // HF templates occasionally call `raise_exception(msg)` for invariant
223    // checks (e.g. "system must come first"). Wire it to a Jinja error.
224    env.add_function(
225        "raise_exception",
226        |msg: String| -> Result<Value, JinjaError> {
227            Err(JinjaError::new(ErrorKind::InvalidOperation, msg))
228        },
229    );
230    env.add_template_owned(TEMPLATE_NAME, source)
231        .context("compiling chat template")?;
232    Ok(env)
233}
234
235impl ChatTemplate {
236    /// Compile a chat template from a raw Jinja string.
237    pub fn from_source(src: impl Into<String>) -> Result<Self> {
238        let source_text: String = src.into();
239        let env = build_env(source_text.clone())?;
240        Ok(Self {
241            env,
242            source_text,
243            source_kind: ChatTemplateSource::Inline,
244            bos_token: None,
245            eos_token: None,
246        })
247    }
248
249    /// Override BOS/EOS strings (passed to the template as `bos_token` /
250    /// `eos_token` Jinja variables).
251    pub fn with_tokens(mut self, bos: Option<String>, eos: Option<String>) -> Self {
252        self.bos_token = bos;
253        self.eos_token = eos;
254        self
255    }
256
257    /// Load template + BOS/EOS from a GGUF file. Reads
258    /// `tokenizer.chat_template` first, then `tokenizer.ggml.chat_template`.
259    pub fn from_gguf(path: &Path) -> Result<Self> {
260        let raw = GgufFile::from_path(path).with_context(|| format!("opening GGUF {path:?}"))?;
261        Self::from_gguf_file(&raw)
262    }
263
264    /// Same as [`from_gguf`](Self::from_gguf), but reuses an already-parsed file.
265    pub fn from_gguf_file(raw: &GgufFile) -> Result<Self> {
266        let (key, src) = pick_chat_template_meta(raw).ok_or_else(|| {
267            anyhow!("no tokenizer.chat_template or tokenizer.ggml.chat_template in GGUF metadata")
268        })?;
269        let env = build_env(src.clone())?;
270        let bos = resolve_special_token(raw, "tokenizer.ggml.bos_token_id");
271        let eos = resolve_special_token(raw, "tokenizer.ggml.eos_token_id");
272        Ok(Self {
273            env,
274            source_text: src,
275            source_kind: ChatTemplateSource::GgufMetadata(key.to_owned()),
276            bos_token: bos,
277            eos_token: eos,
278        })
279    }
280
281    pub fn source_text(&self) -> &str {
282        &self.source_text
283    }
284
285    pub fn source_kind(&self) -> &ChatTemplateSource {
286        &self.source_kind
287    }
288
289    pub fn bos_token(&self) -> Option<&str> {
290        self.bos_token.as_deref()
291    }
292
293    pub fn eos_token(&self) -> Option<&str> {
294        self.eos_token.as_deref()
295    }
296
297    /// Render the template with the given messages.
298    ///
299    /// The template sees Jinja variables: `messages` (list of
300    /// `{role, content}` maps), `add_generation_prompt` (bool), and
301    /// `bos_token` / `eos_token` strings (empty if unknown).
302    pub fn render(&self, messages: &[ChatMessage], add_generation_prompt: bool) -> Result<String> {
303        self.render_with_options(
304            messages,
305            ChatRenderOptions {
306                add_generation_prompt,
307                ..ChatRenderOptions::default()
308            },
309        )
310    }
311
312    /// Render with extra template knobs (`enable_thinking`, …).
313    pub fn render_with_options(
314        &self,
315        messages: &[ChatMessage],
316        opts: ChatRenderOptions,
317    ) -> Result<String> {
318        let msgs = chat_messages_to_values(messages);
319        let ctx = minijinja::context! {
320            messages => Value::from(msgs),
321            add_generation_prompt => opts.add_generation_prompt,
322            enable_thinking => opts.enable_thinking,
323            bos_token => self.bos_token.clone().unwrap_or_default(),
324            eos_token => self.eos_token.clone().unwrap_or_default(),
325        };
326        let tmpl = self
327            .env
328            .get_template(TEMPLATE_NAME)
329            .expect("template registered in build_env");
330        tmpl.render(ctx).context("rendering chat template")
331    }
332}
333
334fn pick_chat_template_meta(raw: &GgufFile) -> Option<(&'static str, String)> {
335    for key in ["tokenizer.chat_template", "tokenizer.ggml.chat_template"] {
336        if let Some(MetaValue::String(s)) = raw.metadata.get(key) {
337            return Some((key, s.clone()));
338        }
339    }
340    None
341}
342
343fn resolve_special_token(raw: &GgufFile, id_key: &str) -> Option<String> {
344    let id = raw.metadata.get(id_key).and_then(MetaValue::as_u32)? as usize;
345    let toks = raw.metadata.get("tokenizer.ggml.tokens")?;
346    let MetaValue::Array(arr) = toks else {
347        return None;
348    };
349    match arr.get(id)? {
350        MetaValue::String(s) => Some(s.clone()),
351        _ => None,
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    // Minimal Qwen / ChatML-style template — same shape as Qwen3's, simplified
360    // enough that test failures point at our rendering plumbing not at
361    // upstream Jinja quirks. Whitespace-trim markers are intentionally
362    // avoided so the literal `\n` inside the template survives.
363    const QWEN_TEMPLATE: &str = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}";
364
365    // Minimal Llama-3-style template using bos_token + headers.
366    const LLAMA3_TEMPLATE: &str = "{% for m in messages %}{% if loop.first %}{{ bos_token }}{% endif %}<|start_header_id|>{{ m.role }}<|end_header_id|>\n\n{{ m.content }}<|eot_id|>{% endfor %}{% if add_generation_prompt %}<|start_header_id|>assistant<|end_header_id|>\n\n{% endif %}";
367
368    // Minimal Gemma-style template.
369    const GEMMA_TEMPLATE: &str = "{% for m in messages %}{% set role = 'user' if m.role == 'system' else m.role %}<start_of_turn>{{ role }}\n{{ m.content }}<end_of_turn>\n{% endfor %}{% if add_generation_prompt %}<start_of_turn>model\n{% endif %}";
370
371    fn sample_conv() -> Vec<ChatMessage> {
372        vec![ChatMessage::system("be concise"), ChatMessage::user("hi")]
373    }
374
375    #[test]
376    fn qwen_template_renders_with_generation_prompt() {
377        let t = ChatTemplate::from_source(QWEN_TEMPLATE).unwrap();
378        let out = t.render(&sample_conv(), true).unwrap();
379        let expected = "<|im_start|>system\nbe concise<|im_end|>\n\
380                        <|im_start|>user\nhi<|im_end|>\n\
381                        <|im_start|>assistant\n";
382        assert_eq!(out, expected);
383    }
384
385    #[test]
386    fn qwen_template_omits_generation_prompt_when_disabled() {
387        let t = ChatTemplate::from_source(QWEN_TEMPLATE).unwrap();
388        let out = t.render(&sample_conv(), false).unwrap();
389        assert!(out.ends_with("<|im_end|>\n"));
390        assert!(!out.contains("<|im_start|>assistant\n"));
391    }
392
393    #[test]
394    fn llama3_template_uses_bos_token() {
395        let t = ChatTemplate::from_source(LLAMA3_TEMPLATE)
396            .unwrap()
397            .with_tokens(Some("<|begin_of_text|>".into()), Some("<|eot_id|>".into()));
398        let out = t.render(&sample_conv(), true).unwrap();
399        let expected = "<|begin_of_text|>\
400                        <|start_header_id|>system<|end_header_id|>\n\nbe concise<|eot_id|>\
401                        <|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|>\
402                        <|start_header_id|>assistant<|end_header_id|>\n\n";
403        assert_eq!(out, expected);
404        assert_eq!(t.bos_token(), Some("<|begin_of_text|>"));
405        assert_eq!(t.eos_token(), Some("<|eot_id|>"));
406    }
407
408    #[test]
409    fn gemma_template_rewrites_system_to_user() {
410        let t = ChatTemplate::from_source(GEMMA_TEMPLATE).unwrap();
411        let out = t.render(&sample_conv(), true).unwrap();
412        let expected = "<start_of_turn>user\nbe concise<end_of_turn>\n\
413                        <start_of_turn>user\nhi<end_of_turn>\n\
414                        <start_of_turn>model\n";
415        assert_eq!(out, expected);
416    }
417
418    #[test]
419    fn dict_get_method_works_like_hf_templates() {
420        const TEMPLATE: &str =
421            "{% for m in messages %}{{ m.get('role') }}:{{ m.get('content') }};{% endfor %}";
422        let t = ChatTemplate::from_source(TEMPLATE).unwrap();
423        let out = t
424            .render(
425                &[ChatMessage::user("hi"), ChatMessage::assistant("yo")],
426                false,
427            )
428            .unwrap();
429        assert_eq!(out, "user:hi;assistant:yo;");
430    }
431
432    #[test]
433    fn enable_thinking_is_visible_to_template() {
434        const TEMPLATE: &str = "{% if enable_thinking %}think{% else %}plain{% endif %}";
435        let t = ChatTemplate::from_source(TEMPLATE).unwrap();
436        let on = t
437            .render_with_options(&[], ChatRenderOptions::gemma4_thinking(false))
438            .unwrap();
439        let off = t.render(&[], false).unwrap();
440        assert_eq!(on, "think");
441        assert_eq!(off, "plain");
442    }
443
444    #[test]
445    fn raise_exception_propagates_as_error() {
446        let t = ChatTemplate::from_source("{{ raise_exception('nope') }}").unwrap();
447        let err = t.render(&[], false).unwrap_err();
448        assert!(format!("{err:#}").contains("nope"));
449    }
450
451    /// Builds a minimal GGUF in a temp file with a chat_template + token
452    /// table, then verifies BOS/EOS resolve and rendering works.
453    #[test]
454    fn from_gguf_reads_template_and_special_tokens() {
455        // We build a v3 GGUF with three metadata keys:
456        //   tokenizer.chat_template      (String)
457        //   tokenizer.ggml.tokens        (Array of String)
458        //   tokenizer.ggml.bos_token_id  (U32)
459        //   tokenizer.ggml.eos_token_id  (U32)
460        // and one tiny f32 tensor so the file passes the loader.
461        let mut buf: Vec<u8> = Vec::new();
462        buf.extend_from_slice(&rlx_gguf::GGUF_MAGIC.to_le_bytes());
463        buf.extend_from_slice(&3u32.to_le_bytes());
464        buf.extend_from_slice(&1u64.to_le_bytes()); // tensor count
465        buf.extend_from_slice(&4u64.to_le_bytes()); // kv count
466
467        let write_string_kv = |buf: &mut Vec<u8>, k: &str, v: &str| {
468            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
469            buf.extend_from_slice(k.as_bytes());
470            buf.extend_from_slice(&8u32.to_le_bytes());
471            buf.extend_from_slice(&(v.len() as u64).to_le_bytes());
472            buf.extend_from_slice(v.as_bytes());
473        };
474        let write_u32_kv = |buf: &mut Vec<u8>, k: &str, v: u32| {
475            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
476            buf.extend_from_slice(k.as_bytes());
477            buf.extend_from_slice(&4u32.to_le_bytes());
478            buf.extend_from_slice(&v.to_le_bytes());
479        };
480        let write_string_array_kv = |buf: &mut Vec<u8>, k: &str, items: &[&str]| {
481            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
482            buf.extend_from_slice(k.as_bytes());
483            // type = Array(9)
484            buf.extend_from_slice(&9u32.to_le_bytes());
485            // element type = String(8)
486            buf.extend_from_slice(&8u32.to_le_bytes());
487            // length (u64)
488            buf.extend_from_slice(&(items.len() as u64).to_le_bytes());
489            for s in items {
490                buf.extend_from_slice(&(s.len() as u64).to_le_bytes());
491                buf.extend_from_slice(s.as_bytes());
492            }
493        };
494
495        write_string_kv(&mut buf, "tokenizer.chat_template", QWEN_TEMPLATE);
496        write_string_array_kv(
497            &mut buf,
498            "tokenizer.ggml.tokens",
499            &["<pad>", "<bos>", "<eos>", "hi"],
500        );
501        write_u32_kv(&mut buf, "tokenizer.ggml.bos_token_id", 1);
502        write_u32_kv(&mut buf, "tokenizer.ggml.eos_token_id", 2);
503
504        // tiny f32 tensor
505        let name = "w";
506        buf.extend_from_slice(&(name.len() as u64).to_le_bytes());
507        buf.extend_from_slice(name.as_bytes());
508        buf.extend_from_slice(&1u32.to_le_bytes());
509        buf.extend_from_slice(&4u64.to_le_bytes());
510        buf.extend_from_slice(&(rlx_gguf::GgmlType::F32 as u32).to_le_bytes());
511        buf.extend_from_slice(&0u64.to_le_bytes());
512        while !buf
513            .len()
514            .is_multiple_of(rlx_gguf::DEFAULT_ALIGNMENT as usize)
515        {
516            buf.push(0);
517        }
518        for _ in 0..4 {
519            buf.extend_from_slice(&1.0f32.to_le_bytes());
520        }
521        let path = std::env::temp_dir().join("rlx_chat_template_from_gguf.gguf");
522        std::fs::write(&path, &buf).unwrap();
523
524        let t = ChatTemplate::from_gguf(&path).expect("from_gguf");
525        assert_eq!(t.bos_token(), Some("<bos>"));
526        assert_eq!(t.eos_token(), Some("<eos>"));
527        let out = t.render(&sample_conv(), true).unwrap();
528        assert!(out.contains("<|im_start|>assistant\n"));
529        match t.source_kind() {
530            ChatTemplateSource::GgufMetadata(k) => assert_eq!(k, "tokenizer.chat_template"),
531            other => panic!("unexpected source: {other:?}"),
532        }
533        std::fs::remove_file(&path).ok();
534    }
535}
rlx_text/chat.rs

rlx_text/
chat.rs
