Skip to main content

llama_cpp_4/
chat.rs

1//! Chat templates, tool calling, and JSON-Schema-driven grammars.
2//!
3//! This wraps llama.cpp's `common/chat.h` — the layer its own server uses to
4//! turn an OpenAI-shaped request into a prompt, a GBNF grammar, and a parser
5//! that can read tool calls back out of whatever the model emits.
6//!
7//! The value over hand-rolling it is [`ChatParams::grammar`] plus
8//! [`ChatParams::grammar_triggers`]. Naive grammar forcing breaks models that
9//! emit reasoning before a tool call: constrain from token zero and the model
10//! can never open its `<think>` block. A *lazy* grammar stays dormant until a
11//! trigger fires — usually the `<tool_call>` marker — so reasoning flows
12//! unconstrained and only the call itself is forced to be well-formed.
13//!
14//! # Example
15//!
16//! ```no_run
17//! # use llama_cpp_4::chat::{ChatTemplates, ChatApplyParams, ToolChoice};
18//! # fn f(model: &llama_cpp_4::model::LlamaModel) -> Result<(), Box<dyn std::error::Error>> {
19//! let templates = ChatTemplates::from_model(model, None)?;
20//!
21//! let applied = templates.apply(
22//!     &ChatApplyParams::new(r#"[{"role":"user","content":"Weather in Tokyo?"}]"#)
23//!         .with_tools(r#"[{"type":"function","function":{"name":"get_weather",
24//!             "parameters":{"type":"object","properties":{"city":{"type":"string"}}}}}]"#)
25//!         .with_tool_choice(ToolChoice::Required),
26//! )?;
27//!
28//! // Feed `applied.prompt` to the model, constrain with `applied.grammar`,
29//! // then read the tool calls back out:
30//! let msg = applied.parse("<tool_call>{\"name\":\"get_weather\"}</tool_call>", false)?;
31//! # Ok(())
32//! # }
33//! ```
34
35use std::ffi::{c_char, CString};
36use std::ptr::NonNull;
37
38use llama_cpp_sys_4 as sys;
39
40use crate::model::LlamaModel;
41
42/// Errors from the chat layer.
43///
44/// An alias for [`ShimError`](crate::shim::ShimError) — every shim-backed
45/// module shares one error type, since they share one status enum and one error
46/// buffer.
47pub type ChatError = crate::shim::ShimError;
48
49use crate::shim::{check_status, last_error, read_string};
50
51/// What the model is allowed to do with the supplied tools.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
53pub enum ToolChoice {
54    /// The model decides whether to call a tool.
55    #[default]
56    Auto,
57    /// The model must call a tool. This is enforced by grammar, not by prompt.
58    Required,
59    /// Tools are visible to the template but must not be called.
60    None,
61}
62
63impl ToolChoice {
64    // The C enum is unsigned (no negative discriminants) but the struct field
65    // it feeds is `int32_t`, so the two need bridging in both directions.
66    #[allow(clippy::cast_possible_wrap)]
67    fn as_raw(self) -> i32 {
68        let raw = match self {
69            Self::Auto => sys::CHAT_SHIM_TOOL_CHOICE_AUTO,
70            Self::Required => sys::CHAT_SHIM_TOOL_CHOICE_REQUIRED,
71            Self::None => sys::CHAT_SHIM_TOOL_CHOICE_NONE,
72        };
73        raw as i32
74    }
75
76    #[allow(clippy::cast_possible_wrap)]
77    fn from_raw(raw: i32) -> Option<Self> {
78        if raw == sys::CHAT_SHIM_TOOL_CHOICE_AUTO as i32 {
79            Some(Self::Auto)
80        } else if raw == sys::CHAT_SHIM_TOOL_CHOICE_REQUIRED as i32 {
81            Some(Self::Required)
82        } else if raw == sys::CHAT_SHIM_TOOL_CHOICE_NONE as i32 {
83            Some(Self::None)
84        } else {
85            None
86        }
87    }
88
89    /// Parse an `OpenAI` `tool_choice` value: `"auto"`, `"required"` or
90    /// `"none"`.
91    ///
92    /// # Errors
93    ///
94    /// Returns [`ChatError::Failed`] if llama.cpp does not recognise `value`.
95    pub fn parse_oaicompat(value: &str) -> Result<Self, ChatError> {
96        let c_value = CString::new(value)?;
97        let mut raw: i32 = 0;
98        let status = unsafe {
99            sys::chat_shim_tool_choice_parse_oaicompat(c_value.as_ptr(), &raw mut raw)
100        };
101        check_status(status)?;
102        Self::from_raw(raw).ok_or_else(|| ChatError::Failed(format!("unknown tool_choice {raw}")))
103    }
104}
105
106/// How reasoning (`<think>` blocks) should be handled when parsing output.
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
108pub enum ReasoningFormat {
109    /// Do not treat reasoning specially.
110    #[default]
111    None,
112    /// Split reasoning into `reasoning_content`. The usual choice.
113    Auto,
114    /// Like [`Self::Auto`], but leaves reasoning inline when streaming.
115    DeepSeekLegacy,
116    /// Split reasoning out, including in streaming deltas.
117    DeepSeek,
118}
119
120impl ReasoningFormat {
121    #[allow(clippy::cast_possible_wrap)]
122    fn as_raw(self) -> i32 {
123        let raw = match self {
124            Self::None => sys::CHAT_SHIM_REASONING_NONE,
125            Self::Auto => sys::CHAT_SHIM_REASONING_AUTO,
126            Self::DeepSeekLegacy => sys::CHAT_SHIM_REASONING_DEEPSEEK_LEGACY,
127            Self::DeepSeek => sys::CHAT_SHIM_REASONING_DEEPSEEK,
128        };
129        raw as i32
130    }
131}
132
133/// Convert a JSON Schema into a GBNF grammar.
134///
135/// The output goes straight to
136/// [`LlamaSampler::grammar`](crate::sampling::LlamaSampler::grammar), which is
137/// how you implement `OpenAI`'s `response_format: json_schema` — the model then
138/// *cannot* emit anything the schema rejects, rather than being asked nicely.
139///
140/// `force_gbnf` mirrors upstream's flag: leave it `false` and the converter may
141/// pick a more compact representation for schemas it recognises.
142///
143/// ```
144/// # use llama_cpp_4::chat::json_schema_to_grammar;
145/// let gbnf = json_schema_to_grammar(r#"{"type":"integer"}"#, false).unwrap();
146/// assert!(gbnf.contains("root"));
147/// ```
148///
149/// # Errors
150///
151/// Returns [`ChatError::BadJson`] if `schema_json` is not valid JSON, or
152/// [`ChatError::Failed`] if it is not a schema llama.cpp can convert.
153pub fn json_schema_to_grammar(schema_json: &str, force_gbnf: bool) -> Result<String, ChatError> {
154    let c_schema = CString::new(schema_json)?;
155    read_string(|buf, len, expected| unsafe {
156        sys::common_json_schema_to_grammar_c(c_schema.as_ptr(), force_gbnf, buf, len, expected)
157    })
158}
159
160/// The chat templates a model ships, ready to apply.
161///
162/// Wraps `common_chat_templates`, which owns parsed Jinja programs — building
163/// it is not free, so hold one per model rather than one per request.
164#[derive(Debug)]
165pub struct ChatTemplates {
166    raw: NonNull<sys::chat_shim_templates>,
167}
168
169// SAFETY: the handle owns a `common_chat_templates_ptr` with no interior
170// mutability reachable through `&self` — `apply` and the getters only read it.
171// The shim's only mutable global is a `thread_local` error buffer.
172unsafe impl Send for ChatTemplates {}
173unsafe impl Sync for ChatTemplates {}
174
175impl Drop for ChatTemplates {
176    fn drop(&mut self) {
177        unsafe { sys::chat_shim_templates_free(self.raw.as_ptr()) }
178    }
179}
180
181impl ChatTemplates {
182    /// Build the template set for `model`.
183    ///
184    /// Pass `None` for `template_override` to use whatever the model ships, or
185    /// Jinja source to override it. Overriding is what lets you reach a
186    /// model's `tool_use` variant — fetch it with
187    /// [`LlamaModel::chat_template`] and hand it back here.
188    ///
189    /// # Errors
190    ///
191    /// Returns [`ChatError::Init`] if the model has no usable template
192    /// or the override does not parse.
193    pub fn from_model(
194        model: &LlamaModel,
195        template_override: Option<&str>,
196    ) -> Result<Self, ChatError> {
197        let c_override = template_override.map(CString::new).transpose()?;
198        let override_ptr = c_override.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
199        let raw = unsafe { sys::chat_shim_templates_init(model.model.as_ptr(), override_ptr) };
200        NonNull::new(raw)
201            .map(|raw| Self { raw })
202            .ok_or_else(|| ChatError::Init(last_error()))
203    }
204
205    /// Where the active template came from, e.g. `"model"` or a builtin name.
206    ///
207    /// # Errors
208    ///
209    /// Returns [`ChatError::Failed`] if llama.cpp could not report a source.
210    pub fn source(&self, variant: Option<&str>) -> Result<String, ChatError> {
211        let c_variant = variant.map(CString::new).transpose()?;
212        let variant_ptr = c_variant.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
213        read_string(|buf, len, expected| unsafe {
214            sys::chat_shim_templates_source(self.raw.as_ptr(), variant_ptr, buf, len, expected)
215        })
216    }
217
218    /// Whether the caller supplied the template rather than the model.
219    #[must_use]
220    pub fn was_explicit(&self) -> bool {
221        unsafe { sys::chat_shim_templates_was_explicit(self.raw.as_ptr()) }
222    }
223
224    /// Whether this template understands `enable_thinking`.
225    #[must_use]
226    pub fn supports_enable_thinking(&self) -> bool {
227        unsafe { sys::chat_shim_templates_support_enable_thinking(self.raw.as_ptr()) }
228    }
229
230    /// Template capabilities as a JSON object of `name -> bool`.
231    ///
232    /// This is what upstream's server reports on `/props`; it tells you whether
233    /// the template can handle tools, parallel calls, a system role, and so on
234    /// *before* you send a request it cannot render.
235    ///
236    /// # Errors
237    ///
238    /// Returns [`ChatError::Failed`] if llama.cpp could not report caps.
239    pub fn caps_json(&self) -> Result<String, ChatError> {
240        read_string(|buf, len, expected| unsafe {
241            sys::chat_shim_templates_get_caps(self.raw.as_ptr(), buf, len, expected)
242        })
243    }
244
245    /// Render messages and tools into a prompt plus its sampling constraints.
246    ///
247    /// # Errors
248    ///
249    /// Returns [`ChatError::BadJson`] for malformed `messages`/`tools` JSON, or
250    /// [`ChatError::Failed`] if the template cannot render the request.
251    pub fn apply(&self, params: &ChatApplyParams) -> Result<ChatParams, ChatError> {
252        let messages = CString::new(params.messages_json.as_str())?;
253        let tools = params.tools_json.as_deref().map(CString::new).transpose()?;
254        let grammar = params.grammar.as_deref().map(CString::new).transpose()?;
255        let schema = params.json_schema.as_deref().map(CString::new).transpose()?;
256        let kwargs = params
257            .template_kwargs_json
258            .as_deref()
259            .map(CString::new)
260            .transpose()?;
261
262        let raw_params = sys::chat_shim_apply_params {
263            messages_json: messages.as_ptr(),
264            tools_json: opt_ptr(tools.as_ref()),
265            grammar: opt_ptr(grammar.as_ref()),
266            json_schema: opt_ptr(schema.as_ref()),
267            template_kwargs_json: opt_ptr(kwargs.as_ref()),
268            tool_choice: params.tool_choice.as_raw(),
269            reasoning_format: params.reasoning_format.as_raw(),
270            add_generation_prompt: params.add_generation_prompt,
271            enable_thinking: params.enable_thinking,
272            parallel_tool_calls: params.parallel_tool_calls,
273            use_jinja: params.use_jinja,
274            add_bos: params.add_bos,
275            add_eos: params.add_eos,
276        };
277
278        // Two-call protocol: size, then fill. The shim fills `result` on both
279        // calls, so the template is only rendered once per call — but we take
280        // the second call's offsets, which describe the buffer we actually got.
281        // bindgen derives no `Default` for this struct; zeroing is correct
282        // here because every field is a `usize` offset, an `i32`, or a `bool`,
283        // and the shim overwrites all of them before reporting success.
284        let mut result: sys::chat_shim_apply_result = unsafe { std::mem::zeroed() };
285        let mut needed: usize = 0;
286        let status = unsafe {
287            sys::chat_shim_templates_apply(
288                self.raw.as_ptr(),
289                &raw const raw_params,
290                &raw mut result,
291                std::ptr::null_mut(),
292                0,
293                &raw mut needed,
294            )
295        };
296        if status != sys::LLAMA_SHIM_BUFFER_TOO_SMALL {
297            check_status(status)?;
298        }
299
300        let mut buf = vec![0u8; needed];
301        let status = unsafe {
302            sys::chat_shim_templates_apply(
303                self.raw.as_ptr(),
304                &raw const raw_params,
305                &raw mut result,
306                buf.as_mut_ptr().cast::<c_char>(),
307                buf.len(),
308                &raw mut needed,
309            )
310        };
311        check_status(status)?;
312
313        ChatParams::from_packed(&buf, &result)
314    }
315}
316
317/// A request to render: messages, optional tools, and how to constrain output.
318///
319/// `messages_json` and `tools_json` are `OpenAI`-shaped JSON. They are passed as
320/// text rather than typed structs so this crate need not take a JSON dependency
321/// or track upstream's message schema, which grows most releases.
322// Mirrors `common_chat_templates_inputs`, which is mostly independent flags;
323// grouping them into sub-structs would diverge from the C layout for no gain.
324#[allow(clippy::struct_excessive_bools)]
325#[derive(Debug, Clone)]
326pub struct ChatApplyParams {
327    messages_json: String,
328    tools_json: Option<String>,
329    grammar: Option<String>,
330    json_schema: Option<String>,
331    template_kwargs_json: Option<String>,
332    tool_choice: ToolChoice,
333    reasoning_format: ReasoningFormat,
334    add_generation_prompt: bool,
335    enable_thinking: bool,
336    parallel_tool_calls: bool,
337    use_jinja: bool,
338    add_bos: bool,
339    add_eos: bool,
340}
341
342impl ChatApplyParams {
343    /// Start from an OpenAI-shaped `messages` JSON array.
344    #[must_use]
345    pub fn new(messages_json: impl Into<String>) -> Self {
346        Self {
347            messages_json: messages_json.into(),
348            tools_json: None,
349            grammar: None,
350            json_schema: None,
351            template_kwargs_json: None,
352            tool_choice: ToolChoice::Auto,
353            reasoning_format: ReasoningFormat::Auto,
354            add_generation_prompt: true,
355            enable_thinking: true,
356            parallel_tool_calls: false,
357            use_jinja: true,
358            add_bos: false,
359            add_eos: false,
360        }
361    }
362
363    /// Supply an OpenAI-shaped `tools` JSON array.
364    #[must_use]
365    pub fn with_tools(mut self, tools_json: impl Into<String>) -> Self {
366        self.tools_json = Some(tools_json.into());
367        self
368    }
369
370    /// Constrain output with a GBNF grammar directly.
371    #[must_use]
372    pub fn with_grammar(mut self, grammar: impl Into<String>) -> Self {
373        self.grammar = Some(grammar.into());
374        self
375    }
376
377    /// Constrain output with a JSON Schema. llama.cpp converts it to GBNF.
378    #[must_use]
379    pub fn with_json_schema(mut self, schema_json: impl Into<String>) -> Self {
380        self.json_schema = Some(schema_json.into());
381        self
382    }
383
384    /// Extra Jinja variables, as a JSON object.
385    #[must_use]
386    pub fn with_template_kwargs(mut self, kwargs_json: impl Into<String>) -> Self {
387        self.template_kwargs_json = Some(kwargs_json.into());
388        self
389    }
390
391    /// Set what the model may do with the tools. Defaults to
392    /// [`ToolChoice::Auto`].
393    #[must_use]
394    pub fn with_tool_choice(mut self, choice: ToolChoice) -> Self {
395        self.tool_choice = choice;
396        self
397    }
398
399    /// Set how reasoning is handled. Defaults to [`ReasoningFormat::Auto`].
400    #[must_use]
401    pub fn with_reasoning_format(mut self, format: ReasoningFormat) -> Self {
402        self.reasoning_format = format;
403        self
404    }
405
406    /// Append the assistant generation prompt. Defaults to `true`; set `false`
407    /// to render a transcript rather than a request.
408    #[must_use]
409    pub fn with_add_generation_prompt(mut self, add: bool) -> Self {
410        self.add_generation_prompt = add;
411        self
412    }
413
414    /// Let the model think before answering, on templates that support it.
415    /// Defaults to `true`.
416    #[must_use]
417    pub fn with_enable_thinking(mut self, enable: bool) -> Self {
418        self.enable_thinking = enable;
419        self
420    }
421
422    /// Allow several tool calls in one turn. Defaults to `false`.
423    #[must_use]
424    pub fn with_parallel_tool_calls(mut self, parallel: bool) -> Self {
425        self.parallel_tool_calls = parallel;
426        self
427    }
428
429    /// Render with the Jinja engine. Defaults to `true`; `false` selects the
430    /// legacy `llama_chat_apply_template` path, which ignores tools.
431    #[must_use]
432    pub fn with_use_jinja(mut self, use_jinja: bool) -> Self {
433        self.use_jinja = use_jinja;
434        self
435    }
436
437    /// Prepend BOS / append EOS to the rendered prompt. Both default to
438    /// `false`, since tokenization usually adds BOS itself.
439    // The two parameter names mirror the upstream fields; renaming either to
440    // satisfy `similar_names` would obscure which C field it sets.
441    #[allow(clippy::similar_names)]
442    #[must_use]
443    pub fn with_add_bos_eos(mut self, add_bos: bool, add_eos: bool) -> Self {
444        self.add_bos = add_bos;
445        self.add_eos = add_eos;
446        self
447    }
448}
449
450/// A lazy-grammar trigger: what has to appear before the grammar engages.
451#[derive(Debug, Clone, PartialEq, Eq)]
452pub struct GrammarTrigger {
453    /// `"token"`, `"word"`, `"pattern"` or `"pattern_full"`.
454    pub kind: String,
455    /// The literal, word, or regex the trigger matches.
456    pub value: String,
457    /// Token id for `"token"` triggers; `-1` otherwise.
458    pub token: i32,
459}
460
461/// The rendered prompt and everything needed to constrain and parse generation.
462#[derive(Debug, Clone)]
463pub struct ChatParams {
464    /// The prompt to feed the model.
465    pub prompt: String,
466    /// GBNF grammar constraining output. Empty when unconstrained.
467    pub grammar: String,
468    /// When true, `grammar` must not be applied until a trigger in
469    /// [`Self::grammar_triggers`] fires. Applying it from token zero is what
470    /// stops thinking models emitting their reasoning prefix.
471    pub grammar_lazy: bool,
472    /// Triggers that activate a lazy grammar. Raw JSON, and parsed into
473    /// [`Self::grammar_triggers`].
474    pub grammar_triggers_json: String,
475    /// Parsed form of [`Self::grammar_triggers_json`].
476    pub grammar_triggers: Vec<GrammarTrigger>,
477    /// JSON array of strings to keep verbatim while sampling.
478    pub preserved_tokens_json: String,
479    /// JSON array of extra stop strings this format needs.
480    pub additional_stops_json: String,
481    /// Whether this template supports reasoning.
482    pub supports_thinking: bool,
483    /// Opening reasoning tag, e.g. `"<think>"`. Empty when unsupported.
484    pub thinking_start_tag: String,
485    /// JSON array of closing reasoning tags.
486    pub thinking_end_tags_json: String,
487    /// `common_chat_format` discriminant, needed to parse output back.
488    pub format: i32,
489    /// Serialized PEG parser produced alongside the prompt. Opaque.
490    parser: String,
491    /// Prefix the parser expects to see before generated text.
492    generation_prompt: String,
493    reasoning_format: ReasoningFormat,
494}
495
496impl ChatParams {
497    fn from_packed(
498        buf: &[u8],
499        result: &sys::chat_shim_apply_result,
500    ) -> Result<Self, ChatError> {
501        let at = |off: usize| -> Result<String, ChatError> {
502            let rest = buf.get(off..).ok_or(ChatError::CorruptResult)?;
503            let end = rest
504                .iter()
505                .position(|b| *b == 0)
506                .ok_or(ChatError::CorruptResult)?;
507            String::from_utf8(rest[..end].to_vec()).map_err(ChatError::from)
508        };
509
510        let grammar_triggers_json = at(result.grammar_triggers_off)?;
511        let grammar_triggers = parse_triggers(&grammar_triggers_json);
512
513        Ok(Self {
514            prompt: at(result.prompt_off)?,
515            grammar: at(result.grammar_off)?,
516            grammar_lazy: result.grammar_lazy,
517            grammar_triggers,
518            grammar_triggers_json,
519            preserved_tokens_json: at(result.preserved_tokens_off)?,
520            additional_stops_json: at(result.additional_stops_off)?,
521            supports_thinking: result.supports_thinking,
522            thinking_start_tag: at(result.thinking_start_tag_off)?,
523            thinking_end_tags_json: at(result.thinking_end_tags_off)?,
524            format: result.format,
525            parser: at(result.parser_off)?,
526            generation_prompt: at(result.generation_prompt_off)?,
527            reasoning_format: ReasoningFormat::Auto,
528        })
529    }
530
531    /// Human-readable name of this chat format, e.g. `"Hermes 2 Pro"`.
532    ///
533    /// # Errors
534    ///
535    /// Returns [`ChatError::Failed`] if llama.cpp cannot name the format.
536    pub fn format_name(&self) -> Result<String, ChatError> {
537        format_name(self.format)
538    }
539
540    /// Convert [`Self::grammar_triggers`] into the `(patterns, tokens)` pair
541    /// [`LlamaSampler::grammar_lazy_patterns`](crate::sampling::LlamaSampler::grammar_lazy_patterns)
542    /// expects.
543    ///
544    /// The four trigger kinds do not map onto the sampler one-for-one, and
545    /// getting the translation wrong silently produces a grammar that never
546    /// activates:
547    ///
548    /// - `word` is a **literal**, so it is regex-escaped before becoming a
549    ///   pattern. A raw `<tool_call>` would otherwise be a character class.
550    /// - `pattern` passes through unchanged.
551    /// - `pattern_full` is anchored with `^`/`$` unless it already is.
552    /// - `token` becomes a trigger token rather than a pattern.
553    ///
554    /// This mirrors `common/sampling.cpp`, so callers do not have to.
555    #[must_use]
556    pub fn sampler_triggers(&self) -> (Vec<String>, Vec<crate::token::LlamaToken>) {
557        let mut patterns = Vec::new();
558        let mut tokens = Vec::new();
559        for trigger in &self.grammar_triggers {
560            match trigger.kind.as_str() {
561                "word" => patterns.push(regex_escape(&trigger.value)),
562                "pattern" => patterns.push(trigger.value.clone()),
563                "pattern_full" => patterns.push(anchor_pattern(&trigger.value)),
564                "token" => tokens.push(crate::token::LlamaToken(trigger.token)),
565                // An unknown kind from a newer llama.cpp: dropping it is safer
566                // than guessing, and the grammar simply stays dormant for it.
567                _ => {}
568            }
569        }
570        (patterns, tokens)
571    }
572
573    /// The prefix the template already placed at the end of the prompt, e.g.
574    /// `"<|im_start|>assistant\n"`.
575    ///
576    /// This is not part of the model's output, but the grammar and the parser
577    /// are both written as if it were — see [`Self::grammar_sampler`].
578    #[must_use]
579    pub fn generation_prompt(&self) -> &str {
580        &self.generation_prompt
581    }
582
583    /// Build the grammar sampler this render needs, or `None` when
584    /// unconstrained.
585    ///
586    /// Prefer this over constructing the sampler yourself: it handles three
587    /// things that are each silently wrong if missed.
588    ///
589    /// - **Lazy vs eager.** A lazy grammar must be built with its triggers, or
590    ///   it never activates. One built eagerly from a lazy grammar constrains
591    ///   from token zero and blocks a thinking model's reasoning prefix.
592    /// - **Trigger translation.** Literal, regex and token triggers map onto
593    ///   the sampler differently — see [`Self::sampler_triggers`].
594    /// - **Generation-prompt prefill.** llama.cpp writes tool-call grammars to
595    ///   match `generation_prompt + output`, because that is what the parser
596    ///   later sees. The sampler only sees `output`, so without advancing the
597    ///   grammar past that prefix it forces the model to *re-emit*
598    ///   `<|im_start|>assistant` as generated text. Prefill applies only to
599    ///   non-lazy grammars — a lazy one has not started matching yet.
600    ///
601    /// # Panics
602    ///
603    /// Panics if llama.cpp cannot parse the grammar it just produced, or if
604    /// that grammar contains an interior NUL.
605    #[must_use]
606    pub fn grammar_sampler(&self, model: &LlamaModel) -> Option<crate::sampling::LlamaSampler> {
607        use crate::sampling::LlamaSampler;
608
609        if self.grammar.is_empty() {
610            return None;
611        }
612
613        let (patterns, tokens) = self.sampler_triggers();
614        let lazy = self.grammar_lazy && !(patterns.is_empty() && tokens.is_empty());
615
616        let mut sampler = if lazy {
617            let refs: Vec<&str> = patterns.iter().map(String::as_str).collect();
618            LlamaSampler::grammar_lazy_patterns(model, &self.grammar, "root", &refs, &tokens)
619        } else {
620            LlamaSampler::grammar(model, &self.grammar, "root")
621        };
622
623        if !lazy {
624            for token in self.generation_prompt_tokens(model) {
625                sampler.accept(token);
626            }
627        }
628        Some(sampler)
629    }
630
631    /// Tokenize [`Self::generation_prompt`] the way llama.cpp does when
632    /// prefilling a grammar.
633    ///
634    /// Some tokenizers prepend a space to the first token; upstream drops it
635    /// when the prompt itself does not start with whitespace, since that space
636    /// is an artefact rather than something the template emitted.
637    fn generation_prompt_tokens(&self, model: &LlamaModel) -> Vec<crate::token::LlamaToken> {
638        if self.generation_prompt.is_empty() {
639            return Vec::new();
640        }
641        let Ok(tokens) = model.str_to_token(&self.generation_prompt, crate::model::AddBos::Never)
642        else {
643            return Vec::new();
644        };
645        let starts_with_space = self
646            .generation_prompt
647            .starts_with(char::is_whitespace);
648        let mut out = Vec::with_capacity(tokens.len());
649        for (i, token) in tokens.into_iter().enumerate() {
650            if i == 0 && !starts_with_space {
651                if let Ok(piece) = model.token_to_str(token, crate::model::Special::Tokenize) {
652                    if piece.starts_with(char::is_whitespace) {
653                        continue;
654                    }
655                }
656            }
657            out.push(token);
658        }
659        out
660    }
661
662    /// Parse model output back into an OpenAI-shaped message JSON object with
663    /// `role`, `content`, `reasoning_content` and `tool_calls`.
664    ///
665    /// This uses the parser the template produced, so it understands that
666    /// model family's tool-call syntax rather than scraping for a fixed marker.
667    ///
668    /// Set `is_partial` while streaming: the parser then tolerates a truncated
669    /// tail instead of rejecting the buffer.
670    ///
671    /// # Errors
672    ///
673    /// Returns [`ChatError::Failed`] if the output cannot be parsed.
674    pub fn parse(&self, text: &str, is_partial: bool) -> Result<String, ChatError> {
675        self.parse_with(text, is_partial, true, false)
676    }
677
678    /// [`Self::parse`] with control over tool-call parsing and whether
679    /// reasoning is left inline in `content`.
680    ///
681    /// # Errors
682    ///
683    /// Returns [`ChatError::Failed`] if the output cannot be parsed.
684    pub fn parse_with(
685        &self,
686        text: &str,
687        is_partial: bool,
688        parse_tool_calls: bool,
689        reasoning_in_content: bool,
690    ) -> Result<String, ChatError> {
691        let c_text = CString::new(text)?;
692        let c_parser = CString::new(self.parser.as_str())?;
693        let c_gen_prompt = CString::new(self.generation_prompt.as_str())?;
694        let params = sys::chat_shim_parse_params {
695            text: c_text.as_ptr(),
696            parser: c_parser.as_ptr(),
697            generation_prompt: c_gen_prompt.as_ptr(),
698            format: self.format,
699            reasoning_format: self.reasoning_format.as_raw(),
700            is_partial,
701            parse_tool_calls,
702            reasoning_in_content,
703        };
704        read_string(|buf, len, expected| unsafe {
705            sys::chat_shim_parse(&raw const params, buf, len, expected)
706        })
707    }
708}
709
710/// Human-readable name of a `common_chat_format` discriminant.
711///
712/// # Errors
713///
714/// Returns [`ChatError::Failed`] if llama.cpp cannot name the format.
715pub fn format_name(format: i32) -> Result<String, ChatError> {
716    read_string(|buf, len, expected| unsafe {
717        sys::chat_shim_format_name(format, buf, len, expected)
718    })
719}
720
721/// Validate and normalise an OpenAI-shaped `messages` array.
722///
723/// Expands the shorthands llama.cpp accepts — typed content parts, legacy
724/// `function_call` — into the canonical form, and rejects malformed input
725/// before it reaches a template.
726///
727/// # Errors
728///
729/// Returns [`ChatError::BadJson`] if the input is not valid JSON, or
730/// [`ChatError::Failed`] if it is not a valid message array.
731pub fn parse_messages_oaicompat(messages_json: &str) -> Result<String, ChatError> {
732    let c_messages = CString::new(messages_json)?;
733    read_string(|buf, len, expected| unsafe {
734        sys::chat_shim_msgs_parse_oaicompat(c_messages.as_ptr(), buf, len, expected)
735    })
736}
737
738/// Validate and normalise an OpenAI-shaped `tools` array into a JSON array of
739/// `{"name", "description", "parameters"}`.
740///
741/// # Errors
742///
743/// Returns [`ChatError::BadJson`] if the input is not valid JSON, or
744/// [`ChatError::Failed`] if it is not a valid tool array.
745pub fn parse_tools_oaicompat(tools_json: &str) -> Result<String, ChatError> {
746    let c_tools = CString::new(tools_json)?;
747    read_string(|buf, len, expected| unsafe {
748        sys::chat_shim_tools_parse_oaicompat(c_tools.as_ptr(), buf, len, expected)
749    })
750}
751
752// ── plumbing ────────────────────────────────────────────────────────────────
753
754/// Escape the characters `std::regex` treats as special, matching
755/// `regex_escape` in `common/common.cpp`.
756fn regex_escape(s: &str) -> String {
757    const SPECIAL: &[char] = &[
758        '.', '^', '$', '|', '(', ')', '*', '+', '?', '[', ']', '{', '}', '\\',
759    ];
760    let mut out = String::with_capacity(s.len());
761    for c in s.chars() {
762        if SPECIAL.contains(&c) {
763            out.push('\\');
764        }
765        out.push(c);
766    }
767    out
768}
769
770/// Anchor a `pattern_full` trigger, matching `common/sampling.cpp`. An empty
771/// pattern becomes `^$` — matching only the empty string — rather than `^^$$`.
772fn anchor_pattern(pattern: &str) -> String {
773    if pattern.is_empty() {
774        return "^$".to_owned();
775    }
776    let mut out = String::with_capacity(pattern.len() + 2);
777    if !pattern.starts_with('^') {
778        out.push('^');
779    }
780    out.push_str(pattern);
781    if !pattern.ends_with('$') {
782        out.push('$');
783    }
784    out
785}
786
787fn opt_ptr(s: Option<&CString>) -> *const c_char {
788    s.map_or(std::ptr::null(), |c| c.as_ptr())
789}
790
791/// Pull the fields out of the shim's trigger JSON without taking a JSON
792/// dependency. The shape is fixed by `chat_shim.cpp`, so a hand parser is
793/// enough — anything unexpected yields no triggers, and
794/// [`ChatParams::grammar_triggers_json`] still carries the raw text.
795fn parse_triggers(json: &str) -> Vec<GrammarTrigger> {
796    let mut out = Vec::new();
797    for chunk in json.split('{').skip(1) {
798        let kind = json_str_field(chunk, "type");
799        let value = json_str_field(chunk, "value");
800        let token = json_int_field(chunk, "token");
801        if let (Some(kind), Some(value)) = (kind, value) {
802            out.push(GrammarTrigger {
803                kind,
804                value,
805                token: token.unwrap_or(-1),
806            });
807        }
808    }
809    out
810}
811
812fn json_str_field(chunk: &str, key: &str) -> Option<String> {
813    let needle = format!("\"{key}\":");
814    let rest = &chunk[chunk.find(&needle)? + needle.len()..];
815    let rest = rest.trim_start();
816    let mut chars = rest.strip_prefix('"')?.chars();
817    let mut value = String::new();
818    while let Some(c) = chars.next() {
819        match c {
820            '"' => return Some(value),
821            '\\' => match chars.next()? {
822                'n' => value.push('\n'),
823                'r' => value.push('\r'),
824                't' => value.push('\t'),
825                'u' => {
826                    let hex: String = chars.by_ref().take(4).collect();
827                    let code = u32::from_str_radix(&hex, 16).ok()?;
828                    value.push(char::from_u32(code)?);
829                }
830                other => value.push(other),
831            },
832            other => value.push(other),
833        }
834    }
835    None
836}
837
838fn json_int_field(chunk: &str, key: &str) -> Option<i32> {
839    let needle = format!("\"{key}\":");
840    let rest = &chunk[chunk.find(&needle)? + needle.len()..];
841    let rest = rest.trim_start();
842    let end = rest
843        .find(|c: char| !c.is_ascii_digit() && c != '-')
844        .unwrap_or(rest.len());
845    rest[..end].parse().ok()
846}
847
848#[cfg(test)]
849mod tests {
850    use super::*;
851
852    /// The schema the `structured` example ships in its docs must convert and
853    /// be loadable by the sampler — an object schema with a required key is the
854    /// shape every `response_format: json_schema` caller sends.
855    #[test]
856    fn realistic_object_schema_converts() {
857        let gbnf = json_schema_to_grammar(
858            r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#,
859            false,
860        )
861        .unwrap();
862        eprintln!("GBNF:\n{gbnf}");
863        assert!(gbnf.contains("root"));
864    }
865
866    #[test]
867    fn json_schema_to_grammar_produces_a_root_rule() {
868        let gbnf = json_schema_to_grammar(r#"{"type":"integer"}"#, false).unwrap();
869        assert!(gbnf.contains("root"), "no root rule in: {gbnf}");
870    }
871
872    /// An object schema must constrain its keys, otherwise the grammar is not
873    /// actually enforcing the schema.
874    #[test]
875    fn json_schema_to_grammar_constrains_object_keys() {
876        let gbnf = json_schema_to_grammar(
877            r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#,
878            false,
879        )
880        .unwrap();
881        assert!(gbnf.contains("city"), "key not in grammar: {gbnf}");
882    }
883
884    /// Invalid JSON must be reported as such rather than crossing the FFI
885    /// boundary as an exception.
886    #[test]
887    fn json_schema_to_grammar_rejects_invalid_json() {
888        let err = json_schema_to_grammar("{not json", false).unwrap_err();
889        assert!(
890            matches!(err, ChatError::BadJson(_)),
891            "expected BadJson, got {err:?}"
892        );
893    }
894
895    #[test]
896    fn json_schema_to_grammar_rejects_empty_input() {
897        assert!(json_schema_to_grammar("", false).is_err());
898    }
899
900    /// Interior NULs must be caught in Rust — a `CString` cannot hold one, and
901    /// truncating silently would change the schema.
902    #[test]
903    fn json_schema_to_grammar_rejects_interior_nul() {
904        let err = json_schema_to_grammar("{\"type\":\"in\0teger\"}", false).unwrap_err();
905        assert!(matches!(err, ChatError::Nul(_)), "got {err:?}");
906    }
907
908    #[test]
909    fn tool_choice_parses_openai_values() {
910        assert_eq!(ToolChoice::parse_oaicompat("auto").unwrap(), ToolChoice::Auto);
911        assert_eq!(
912            ToolChoice::parse_oaicompat("required").unwrap(),
913            ToolChoice::Required
914        );
915        assert_eq!(ToolChoice::parse_oaicompat("none").unwrap(), ToolChoice::None);
916    }
917
918    #[test]
919    fn tool_choice_rejects_unknown_value() {
920        assert!(ToolChoice::parse_oaicompat("sometimes").is_err());
921    }
922
923    #[test]
924    fn tools_parse_oaicompat_extracts_name_and_parameters() {
925        let normalised = parse_tools_oaicompat(
926            r#"[{"type":"function","function":{"name":"get_weather",
927                "description":"Get weather","parameters":{"type":"object"}}}]"#,
928        )
929        .unwrap();
930        assert!(normalised.contains("get_weather"), "got {normalised}");
931    }
932
933    #[test]
934    fn tools_parse_oaicompat_rejects_garbage() {
935        assert!(parse_tools_oaicompat("[[[").is_err());
936    }
937
938    #[test]
939    fn messages_parse_oaicompat_roundtrips_a_simple_turn() {
940        let normalised =
941            parse_messages_oaicompat(r#"[{"role":"user","content":"hi"}]"#).unwrap();
942        assert!(normalised.contains("user"), "got {normalised}");
943        assert!(normalised.contains("hi"), "got {normalised}");
944    }
945
946    #[test]
947    fn messages_parse_oaicompat_rejects_non_array() {
948        assert!(parse_messages_oaicompat(r#"{"role":"user"}"#).is_err());
949    }
950
951    /// The trigger parser is hand-rolled, so pin it against the exact shape
952    /// `chat_shim.cpp` emits, including escapes.
953    #[test]
954    fn trigger_parser_reads_shim_output() {
955        let json = r#"[{"type":"word","value":"<tool_call>","token":-1},
956                       {"type":"token","value":"a\nb","token":42}]"#;
957        let triggers = parse_triggers(json);
958        assert_eq!(triggers.len(), 2);
959        assert_eq!(triggers[0].kind, "word");
960        assert_eq!(triggers[0].value, "<tool_call>");
961        assert_eq!(triggers[0].token, -1);
962        assert_eq!(triggers[1].kind, "token");
963        assert_eq!(triggers[1].value, "a\nb");
964        assert_eq!(triggers[1].token, 42);
965    }
966
967    /// A literal trigger like `<tool_call>` must be escaped before it becomes
968    /// a regex — unescaped, `[` and `]` would make it a character class and the
969    /// grammar would never fire.
970    #[test]
971    fn word_triggers_are_regex_escaped() {
972        let params = ChatParams {
973            grammar_triggers: vec![GrammarTrigger {
974                kind: "word".to_owned(),
975                value: "a.b[c]".to_owned(),
976                token: -1,
977            }],
978            ..stub_params()
979        };
980        let (patterns, tokens) = params.sampler_triggers();
981        assert_eq!(patterns, vec![r"a\.b\[c\]".to_owned()]);
982        assert!(tokens.is_empty());
983    }
984
985    #[test]
986    fn pattern_triggers_pass_through_unescaped() {
987        let params = ChatParams {
988            grammar_triggers: vec![GrammarTrigger {
989                kind: "pattern".to_owned(),
990                value: "a.b".to_owned(),
991                token: -1,
992            }],
993            ..stub_params()
994        };
995        assert_eq!(params.sampler_triggers().0, vec!["a.b".to_owned()]);
996    }
997
998    #[test]
999    fn pattern_full_triggers_are_anchored_once() {
1000        let cases = [
1001            ("abc", "^abc$"),
1002            ("^abc", "^abc$"),
1003            ("abc$", "^abc$"),
1004            ("^abc$", "^abc$"),
1005            ("", "^$"),
1006        ];
1007        for (input, want) in cases {
1008            let params = ChatParams {
1009                grammar_triggers: vec![GrammarTrigger {
1010                    kind: "pattern_full".to_owned(),
1011                    value: input.to_owned(),
1012                    token: -1,
1013                }],
1014                ..stub_params()
1015            };
1016            assert_eq!(
1017                params.sampler_triggers().0,
1018                vec![want.to_owned()],
1019                "anchoring {input:?}"
1020            );
1021        }
1022    }
1023
1024    #[test]
1025    fn token_triggers_become_tokens_not_patterns() {
1026        let params = ChatParams {
1027            grammar_triggers: vec![GrammarTrigger {
1028                kind: "token".to_owned(),
1029                value: String::new(),
1030                token: 42,
1031            }],
1032            ..stub_params()
1033        };
1034        let (patterns, tokens) = params.sampler_triggers();
1035        assert!(patterns.is_empty());
1036        assert_eq!(tokens, vec![crate::token::LlamaToken(42)]);
1037    }
1038
1039    /// A kind from a newer llama.cpp must be dropped rather than guessed at.
1040    #[test]
1041    fn unknown_trigger_kinds_are_dropped() {
1042        let params = ChatParams {
1043            grammar_triggers: vec![GrammarTrigger {
1044                kind: "something_new".to_owned(),
1045                value: "x".to_owned(),
1046                token: -1,
1047            }],
1048            ..stub_params()
1049        };
1050        let (patterns, tokens) = params.sampler_triggers();
1051        assert!(patterns.is_empty());
1052        assert!(tokens.is_empty());
1053    }
1054
1055    fn stub_params() -> ChatParams {
1056        ChatParams {
1057            prompt: String::new(),
1058            grammar: String::new(),
1059            grammar_lazy: false,
1060            grammar_triggers_json: String::new(),
1061            grammar_triggers: Vec::new(),
1062            preserved_tokens_json: String::new(),
1063            additional_stops_json: String::new(),
1064            supports_thinking: false,
1065            thinking_start_tag: String::new(),
1066            thinking_end_tags_json: String::new(),
1067            format: 0,
1068            parser: String::new(),
1069            generation_prompt: String::new(),
1070            reasoning_format: ReasoningFormat::Auto,
1071        }
1072    }
1073
1074    #[test]
1075    fn trigger_parser_handles_empty_array() {
1076        assert!(parse_triggers("[]").is_empty());
1077    }
1078}