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;
37use std::sync::Arc;
38
39use llama_cpp_sys_4 as sys;
40
41use crate::model::LlamaModel;
42
43/// Errors from the chat layer.
44///
45/// An alias for [`crate::ShimError`] — every shim-backed
46/// module shares one error type, since they share one status enum and one error
47/// buffer.
48pub type ChatError = crate::shim::ShimError;
49
50use crate::shim::{check_status, last_error, read_string};
51
52/// What the model is allowed to do with the supplied tools.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
54pub enum ToolChoice {
55    /// The model decides whether to call a tool.
56    #[default]
57    Auto,
58    /// The model must call a tool. This is enforced by grammar, not by prompt.
59    Required,
60    /// Tools are visible to the template but must not be called.
61    None,
62}
63
64impl ToolChoice {
65    // The C enum is unsigned (no negative discriminants) but the struct field
66    // it feeds is `int32_t`, so the two need bridging in both directions.
67    #[allow(clippy::cast_possible_wrap)]
68    fn as_raw(self) -> i32 {
69        let raw = match self {
70            Self::Auto => sys::CHAT_SHIM_TOOL_CHOICE_AUTO,
71            Self::Required => sys::CHAT_SHIM_TOOL_CHOICE_REQUIRED,
72            Self::None => sys::CHAT_SHIM_TOOL_CHOICE_NONE,
73        };
74        raw as i32
75    }
76
77    #[allow(clippy::cast_possible_wrap)]
78    fn from_raw(raw: i32) -> Option<Self> {
79        if raw == sys::CHAT_SHIM_TOOL_CHOICE_AUTO as i32 {
80            Some(Self::Auto)
81        } else if raw == sys::CHAT_SHIM_TOOL_CHOICE_REQUIRED as i32 {
82            Some(Self::Required)
83        } else if raw == sys::CHAT_SHIM_TOOL_CHOICE_NONE as i32 {
84            Some(Self::None)
85        } else {
86            None
87        }
88    }
89
90    /// Parse an `OpenAI` `tool_choice` value: `"auto"`, `"required"` or
91    /// `"none"`.
92    ///
93    /// # Errors
94    ///
95    /// Returns [`ChatError::Failed`] if llama.cpp does not recognise `value`.
96    pub fn parse_oaicompat(value: &str) -> Result<Self, ChatError> {
97        let c_value = CString::new(value)?;
98        let mut raw: i32 = 0;
99        let status =
100            // SAFETY: the `CString` arguments are alive for the duration of the call, and any
101            // out-param is a stack local the shim writes through.
102            unsafe { sys::chat_shim_tool_choice_parse_oaicompat(c_value.as_ptr(), &raw mut raw) };
103        check_status(status)?;
104        Self::from_raw(raw).ok_or_else(|| ChatError::Failed(format!("unknown tool_choice {raw}")))
105    }
106}
107
108/// How reasoning (`<think>` blocks) should be handled when parsing output.
109#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
110pub enum ReasoningFormat {
111    /// Do not treat reasoning specially.
112    #[default]
113    None,
114    /// Split reasoning into `reasoning_content`. The usual choice.
115    Auto,
116    /// Like [`Self::Auto`], but leaves reasoning inline when streaming.
117    DeepSeekLegacy,
118    /// Split reasoning out, including in streaming deltas.
119    DeepSeek,
120}
121
122impl ReasoningFormat {
123    #[allow(clippy::cast_possible_wrap)]
124    fn as_raw(self) -> i32 {
125        let raw = match self {
126            Self::None => sys::CHAT_SHIM_REASONING_NONE,
127            Self::Auto => sys::CHAT_SHIM_REASONING_AUTO,
128            Self::DeepSeekLegacy => sys::CHAT_SHIM_REASONING_DEEPSEEK_LEGACY,
129            Self::DeepSeek => sys::CHAT_SHIM_REASONING_DEEPSEEK,
130        };
131        raw as i32
132    }
133}
134
135/// Convert a JSON Schema into a GBNF grammar.
136///
137/// The output goes straight to
138/// [`LlamaSampler::grammar`](crate::sampling::LlamaSampler::grammar), which is
139/// how you implement `OpenAI`'s `response_format: json_schema` — the model then
140/// *cannot* emit anything the schema rejects, rather than being asked nicely.
141///
142/// `force_gbnf` mirrors upstream's flag: leave it `false` and the converter may
143/// pick a more compact representation for schemas it recognises.
144///
145/// ```
146/// # use llama_cpp_4::chat::json_schema_to_grammar;
147/// let gbnf = json_schema_to_grammar(r#"{"type":"integer"}"#, false).unwrap();
148/// assert!(gbnf.contains("root"));
149/// ```
150///
151/// # Errors
152///
153/// Returns [`ChatError::BadJson`] if `schema_json` is not valid JSON, or
154/// [`ChatError::Failed`] if it is not a schema llama.cpp can convert.
155pub fn json_schema_to_grammar(schema_json: &str, force_gbnf: bool) -> Result<String, ChatError> {
156    let c_schema = CString::new(schema_json)?;
157    // SAFETY: the `CString` arguments are alive for the duration of the call, and any
158    // out-param is a stack local the shim writes through.
159    read_string(|buf, len, expected| unsafe {
160        sys::common_json_schema_to_grammar_c(c_schema.as_ptr(), force_gbnf, buf, len, expected)
161    })
162}
163
164/// The chat templates a model ships, ready to apply.
165///
166/// Wraps `common_chat_templates`, which owns parsed Jinja programs — building
167/// it is not free, so hold one per model rather than one per request.
168#[derive(Debug)]
169pub struct ChatTemplates {
170    raw: NonNull<sys::chat_shim_templates>,
171}
172
173// SAFETY: the handle owns a `common_chat_templates_ptr` with no interior
174// mutability reachable through `&self` — `apply` and the getters only read it.
175// The shim's only mutable global is a `thread_local` error buffer.
176unsafe impl Send for ChatTemplates {}
177// SAFETY: every method that mutates this handle takes `&mut self`, so a shared
178// reference reaches only reads.
179unsafe impl Sync for ChatTemplates {}
180
181impl Drop for ChatTemplates {
182    fn drop(&mut self) {
183        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
184        // afterwards.
185        unsafe { sys::chat_shim_templates_free(self.raw.as_ptr()) }
186    }
187}
188
189impl ChatTemplates {
190    /// Build the template set for `model`.
191    ///
192    /// Pass `None` for `template_override` to use whatever the model ships, or
193    /// Jinja source to override it. Overriding is what lets you reach a
194    /// model's `tool_use` variant — fetch it with
195    /// [`LlamaModel::chat_template`] and hand it back here.
196    ///
197    /// # Errors
198    ///
199    /// Returns [`ChatError::Init`] if the model has no usable template
200    /// or the override does not parse.
201    pub fn from_model(
202        model: &LlamaModel,
203        template_override: Option<&str>,
204    ) -> Result<Self, ChatError> {
205        let c_override = template_override.map(CString::new).transpose()?;
206        let override_ptr = c_override.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
207        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
208        // the library owns whatever it returns.
209        let raw = unsafe { sys::chat_shim_templates_init(model.model.as_ptr(), override_ptr) };
210        NonNull::new(raw)
211            .map(|raw| Self { raw })
212            .ok_or_else(|| ChatError::Init(last_error()))
213    }
214
215    /// Where the active template came from, e.g. `"model"` or a builtin name.
216    ///
217    /// # Errors
218    ///
219    /// Returns [`ChatError::Failed`] if llama.cpp could not report a source.
220    pub fn source(&self, variant: Option<&str>) -> Result<String, ChatError> {
221        let c_variant = variant.map(CString::new).transpose()?;
222        let variant_ptr = c_variant.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
223        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
224        // and the shim reports how much it needed; nothing is read back unless it
225        // reported success.
226        read_string(|buf, len, expected| unsafe {
227            sys::chat_shim_templates_source(self.raw.as_ptr(), variant_ptr, buf, len, expected)
228        })
229    }
230
231    /// Whether the caller supplied the template rather than the model.
232    #[must_use]
233    pub fn was_explicit(&self) -> bool {
234        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
235        // arguments are scalars the library validates itself.
236        unsafe { sys::chat_shim_templates_was_explicit(self.raw.as_ptr()) }
237    }
238
239    /// Whether this template understands `enable_thinking`.
240    #[must_use]
241    pub fn supports_enable_thinking(&self) -> bool {
242        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
243        // arguments are scalars the library validates itself.
244        unsafe { sys::chat_shim_templates_support_enable_thinking(self.raw.as_ptr()) }
245    }
246
247    /// Template capabilities as a JSON object of `name -> bool`.
248    ///
249    /// This is what upstream's server reports on `/props`; it tells you whether
250    /// the template can handle tools, parallel calls, a system role, and so on
251    /// *before* you send a request it cannot render.
252    ///
253    /// # Errors
254    ///
255    /// Returns [`ChatError::Failed`] if llama.cpp could not report caps.
256    pub fn caps_json(&self) -> Result<String, ChatError> {
257        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
258        // and the shim reports how much it needed; nothing is read back unless it
259        // reported success.
260        read_string(|buf, len, expected| unsafe {
261            sys::chat_shim_templates_get_caps(self.raw.as_ptr(), buf, len, expected)
262        })
263    }
264
265    /// Render messages and tools into a prompt plus its sampling constraints.
266    ///
267    /// # Errors
268    ///
269    /// Returns [`ChatError::BadJson`] for malformed `messages`/`tools` JSON, or
270    /// [`ChatError::Failed`] if the template cannot render the request.
271    pub fn apply(&self, params: &ChatApplyParams) -> Result<ChatParams, ChatError> {
272        let messages = CString::new(params.messages_json.as_str())?;
273        let tools = params.tools_json.as_deref().map(CString::new).transpose()?;
274        let grammar = params.grammar.as_deref().map(CString::new).transpose()?;
275        let schema = params
276            .json_schema
277            .as_deref()
278            .map(CString::new)
279            .transpose()?;
280        let kwargs = params
281            .template_kwargs_json
282            .as_deref()
283            .map(CString::new)
284            .transpose()?;
285
286        let raw_params = sys::chat_shim_apply_params {
287            messages_json: messages.as_ptr(),
288            tools_json: opt_ptr(tools.as_ref()),
289            grammar: opt_ptr(grammar.as_ref()),
290            json_schema: opt_ptr(schema.as_ref()),
291            template_kwargs_json: opt_ptr(kwargs.as_ref()),
292            tool_choice: params.tool_choice.as_raw(),
293            reasoning_format: params.reasoning_format.as_raw(),
294            add_generation_prompt: params.add_generation_prompt,
295            enable_thinking: params.enable_thinking,
296            parallel_tool_calls: params.parallel_tool_calls,
297            use_jinja: params.use_jinja,
298            add_bos: params.add_bos,
299            add_eos: params.add_eos,
300            force_pure_content: params.force_pure_content,
301        };
302
303        // Two-call protocol: size, then fill. The shim fills `result` on both
304        // calls, so the template is only rendered once per call — but we take
305        // the second call's offsets, which describe the buffer we actually got.
306        // bindgen derives no `Default` for this struct; zeroing is correct
307        // here because every field is a `usize` offset, an `i32`, a `bool`, or
308        // the parser pointer — for which zero is the NULL the shim also writes
309        // on a size query — and the shim overwrites all of them before
310        // reporting success.
311        // SAFETY: the struct is a plain-old-data `repr(C)` type with no pointer or
312        // reference fields, so an all-zero value is a valid instance; the shim
313        // fills it in immediately below.
314        let mut result: sys::chat_shim_apply_result = unsafe { std::mem::zeroed() };
315        let mut needed: usize = 0;
316        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
317        // arguments are scalars the library validates itself.
318        let status = unsafe {
319            sys::chat_shim_templates_apply(
320                self.raw.as_ptr(),
321                &raw const raw_params,
322                &raw mut result,
323                std::ptr::null_mut(),
324                0,
325                &raw mut needed,
326            )
327        };
328        if status != sys::LLAMA_SHIM_BUFFER_TOO_SMALL {
329            check_status(status)?;
330        }
331
332        let mut buf = vec![0u8; needed];
333        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
334        // arguments are scalars the library validates itself.
335        let status = unsafe {
336            sys::chat_shim_templates_apply(
337                self.raw.as_ptr(),
338                &raw const raw_params,
339                &raw mut result,
340                buf.as_mut_ptr().cast::<c_char>(),
341                buf.len(),
342                &raw mut needed,
343            )
344        };
345        check_status(status)?;
346
347        // Take ownership of the parser before anything else fallible runs, or
348        // a later error would leak it. The shim only sets it on success.
349        let parser = NonNull::new(result.parser).map(|raw| Arc::new(ChatParser { raw }));
350
351        ChatParams::from_packed(&buf, &result, parser)
352    }
353}
354
355/// A request to render: messages, optional tools, and how to constrain output.
356///
357/// `messages_json` and `tools_json` are `OpenAI`-shaped JSON. They are passed as
358/// text rather than typed structs so this crate need not take a JSON dependency
359/// or track upstream's message schema, which grows most releases.
360// Mirrors `common_chat_templates_inputs`, which is mostly independent flags;
361// grouping them into sub-structs would diverge from the C layout for no gain.
362#[allow(clippy::struct_excessive_bools)]
363#[derive(Debug, Clone)]
364pub struct ChatApplyParams {
365    messages_json: String,
366    tools_json: Option<String>,
367    grammar: Option<String>,
368    json_schema: Option<String>,
369    template_kwargs_json: Option<String>,
370    tool_choice: ToolChoice,
371    reasoning_format: ReasoningFormat,
372    add_generation_prompt: bool,
373    enable_thinking: bool,
374    parallel_tool_calls: bool,
375    use_jinja: bool,
376    add_bos: bool,
377    add_eos: bool,
378    force_pure_content: bool,
379}
380
381impl ChatApplyParams {
382    /// Start from an OpenAI-shaped `messages` JSON array.
383    #[must_use]
384    pub fn new(messages_json: impl Into<String>) -> Self {
385        Self {
386            messages_json: messages_json.into(),
387            tools_json: None,
388            grammar: None,
389            json_schema: None,
390            template_kwargs_json: None,
391            tool_choice: ToolChoice::Auto,
392            reasoning_format: ReasoningFormat::Auto,
393            add_generation_prompt: true,
394            enable_thinking: true,
395            parallel_tool_calls: false,
396            use_jinja: true,
397            add_bos: false,
398            add_eos: false,
399            force_pure_content: false,
400        }
401    }
402
403    /// Supply an OpenAI-shaped `tools` JSON array.
404    #[must_use]
405    pub fn with_tools(mut self, tools_json: impl Into<String>) -> Self {
406        self.tools_json = Some(tools_json.into());
407        self
408    }
409
410    /// Constrain output with a GBNF grammar directly.
411    #[must_use]
412    pub fn with_grammar(mut self, grammar: impl Into<String>) -> Self {
413        self.grammar = Some(grammar.into());
414        self
415    }
416
417    /// Constrain output with a JSON Schema. llama.cpp converts it to GBNF.
418    #[must_use]
419    pub fn with_json_schema(mut self, schema_json: impl Into<String>) -> Self {
420        self.json_schema = Some(schema_json.into());
421        self
422    }
423
424    /// Extra Jinja variables, as a JSON object.
425    #[must_use]
426    pub fn with_template_kwargs(mut self, kwargs_json: impl Into<String>) -> Self {
427        self.template_kwargs_json = Some(kwargs_json.into());
428        self
429    }
430
431    /// Render a content-only template: reasoning and tool calls are left
432    /// inline in `content` instead of being parsed out.
433    ///
434    /// This is the replacement for the per-parse `parse_tool_calls` /
435    /// `reasoning_in_content` flags llama.cpp dropped in #30210. Because the
436    /// parser is now built once by `apply`, the decision has to be made here
437    /// rather than at [`ChatParams::parse`] time.
438    #[must_use]
439    pub fn with_force_pure_content(mut self, force: bool) -> Self {
440        self.force_pure_content = force;
441        self
442    }
443
444    /// Set what the model may do with the tools. Defaults to
445    /// [`ToolChoice::Auto`].
446    #[must_use]
447    pub fn with_tool_choice(mut self, choice: ToolChoice) -> Self {
448        self.tool_choice = choice;
449        self
450    }
451
452    /// Set how reasoning is handled. Defaults to [`ReasoningFormat::Auto`].
453    #[must_use]
454    pub fn with_reasoning_format(mut self, format: ReasoningFormat) -> Self {
455        self.reasoning_format = format;
456        self
457    }
458
459    /// Append the assistant generation prompt. Defaults to `true`; set `false`
460    /// to render a transcript rather than a request.
461    #[must_use]
462    pub fn with_add_generation_prompt(mut self, add: bool) -> Self {
463        self.add_generation_prompt = add;
464        self
465    }
466
467    /// Let the model think before answering, on templates that support it.
468    /// Defaults to `true`.
469    #[must_use]
470    pub fn with_enable_thinking(mut self, enable: bool) -> Self {
471        self.enable_thinking = enable;
472        self
473    }
474
475    /// Allow several tool calls in one turn. Defaults to `false`.
476    #[must_use]
477    pub fn with_parallel_tool_calls(mut self, parallel: bool) -> Self {
478        self.parallel_tool_calls = parallel;
479        self
480    }
481
482    /// Render with the Jinja engine. Defaults to `true`; `false` selects the
483    /// legacy `llama_chat_apply_template` path, which ignores tools.
484    #[must_use]
485    pub fn with_use_jinja(mut self, use_jinja: bool) -> Self {
486        self.use_jinja = use_jinja;
487        self
488    }
489
490    /// Prepend BOS / append EOS to the rendered prompt. Both default to
491    /// `false`, since tokenization usually adds BOS itself.
492    // The two parameter names mirror the upstream fields; renaming either to
493    // satisfy `similar_names` would obscure which C field it sets.
494    #[allow(clippy::similar_names)]
495    #[must_use]
496    pub fn with_add_bos_eos(mut self, add_bos: bool, add_eos: bool) -> Self {
497        self.add_bos = add_bos;
498        self.add_eos = add_eos;
499        self
500    }
501}
502
503/// The PEG parser a template produces, owned on the llama.cpp side.
504///
505/// `common_chat_templates_apply` builds a parser arena alongside the prompt,
506/// and [`ChatParams::parse`] needs it to read tool calls back out of model
507/// output. It used to round-trip through a string, but llama.cpp #30210 removed
508/// the arena's serialization, so it is held by handle instead and shared
509/// between clones of a [`ChatParams`].
510struct ChatParser {
511    raw: NonNull<sys::chat_shim_parser>,
512}
513
514// SAFETY: the arena is immutable once `apply` has built it, and upstream parses
515// through a `const common_peg_arena &` with all mutable state in a per-call
516// context. Nothing reachable through `&self` mutates it.
517unsafe impl Send for ChatParser {}
518// SAFETY: every method that mutates this handle takes `&mut self`, so a shared
519// reference reaches only reads.
520unsafe impl Sync for ChatParser {}
521
522impl std::fmt::Debug for ChatParser {
523    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
524        f.write_str("ChatParser(<opaque>)")
525    }
526}
527
528impl Drop for ChatParser {
529    fn drop(&mut self) {
530        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
531        // afterwards.
532        unsafe { sys::chat_shim_parser_free(self.raw.as_ptr()) }
533    }
534}
535
536/// One span of a rendered prompt, and whether it came from user content.
537///
538/// llama.cpp builds the prompt as a list of parts rather than one string
539/// ([#30389](https://github.com/ggml-org/llama.cpp/pull/30389)) so that the
540/// spans a *user* supplied stay distinguishable from the ones the template
541/// emitted. That distinction is a prompt-injection defence: a part with
542/// `is_input` set must be tokenized **without** special-token parsing, or a
543/// user who types `<|im_start|>assistant` forges a turn.
544///
545/// Concatenating `text` across all parts reproduces [`ChatParams::prompt`].
546#[derive(Debug, Clone, PartialEq, Eq)]
547pub struct PromptPart {
548    /// The span's text.
549    pub text: String,
550    /// `true` when the span came from user content, so special tokens inside
551    /// it must not be parsed. Tokenize these with
552    /// [`Special::Plaintext`](crate::model::Special::Plaintext) and the rest
553    /// with [`Special::Tokenize`](crate::model::Special::Tokenize).
554    pub is_input: bool,
555}
556
557/// A lazy-grammar trigger: what has to appear before the grammar engages.
558#[derive(Debug, Clone, PartialEq, Eq)]
559pub struct GrammarTrigger {
560    /// `"token"`, `"word"`, `"pattern"` or `"pattern_full"`.
561    pub kind: String,
562    /// The literal, word, or regex the trigger matches.
563    pub value: String,
564    /// Token id for `"token"` triggers; `-1` otherwise.
565    pub token: i32,
566}
567
568/// The rendered prompt and everything needed to constrain and parse generation.
569#[derive(Debug, Clone)]
570pub struct ChatParams {
571    /// The prompt to feed the model.
572    pub prompt: String,
573    /// GBNF grammar constraining output. Empty when unconstrained.
574    pub grammar: String,
575    /// When true, `grammar` must not be applied until a trigger in
576    /// [`Self::grammar_triggers`] fires. Applying it from token zero is what
577    /// stops thinking models emitting their reasoning prefix.
578    pub grammar_lazy: bool,
579    /// Triggers that activate a lazy grammar. Raw JSON, and parsed into
580    /// [`Self::grammar_triggers`].
581    pub grammar_triggers_json: String,
582    /// Parsed form of [`Self::grammar_triggers_json`].
583    pub grammar_triggers: Vec<GrammarTrigger>,
584    /// JSON array of strings to keep verbatim while sampling.
585    pub preserved_tokens_json: String,
586    /// JSON array of extra stop strings this format needs.
587    pub additional_stops_json: String,
588    /// Whether this template supports reasoning.
589    pub supports_thinking: bool,
590    /// Opening reasoning tag, e.g. `"<think>"`. Empty when unsupported.
591    pub thinking_start_tag: String,
592    /// JSON array of closing reasoning tags.
593    pub thinking_end_tags_json: String,
594    /// `common_chat_format` discriminant, needed to parse output back.
595    pub format: i32,
596    /// The prompt split into spans, flagging which came from user content.
597    /// See [`PromptPart`] — this is what lets you tokenize user spans without
598    /// special-token parsing. Raw JSON form in
599    /// [`Self::prompt_parts_json`].
600    pub prompt_parts: Vec<PromptPart>,
601    /// Raw JSON behind [`Self::prompt_parts`], a `[{"text", "is_input"}]`
602    /// array.
603    pub prompt_parts_json: String,
604    /// PEG parser produced alongside the prompt, or `None` when the template
605    /// produced none — [`Self::parse`] then falls back to plain content.
606    parser: Option<Arc<ChatParser>>,
607    /// Prefix the parser expects to see before generated text.
608    generation_prompt: String,
609}
610
611impl ChatParams {
612    fn from_packed(
613        buf: &[u8],
614        result: &sys::chat_shim_apply_result,
615        parser: Option<Arc<ChatParser>>,
616    ) -> Result<Self, ChatError> {
617        let at = |off: usize| -> Result<String, ChatError> {
618            let rest = buf.get(off..).ok_or(ChatError::CorruptResult)?;
619            let end = rest
620                .iter()
621                .position(|b| *b == 0)
622                .ok_or(ChatError::CorruptResult)?;
623            String::from_utf8(rest[..end].to_vec()).map_err(ChatError::from)
624        };
625
626        let grammar_triggers_json = at(result.grammar_triggers_off)?;
627        let prompt_parts_json = at(result.prompt_parts_off)?;
628        let grammar_triggers = parse_triggers(&grammar_triggers_json);
629
630        Ok(Self {
631            prompt: at(result.prompt_off)?,
632            grammar: at(result.grammar_off)?,
633            grammar_lazy: result.grammar_lazy,
634            grammar_triggers,
635            grammar_triggers_json,
636            preserved_tokens_json: at(result.preserved_tokens_off)?,
637            additional_stops_json: at(result.additional_stops_off)?,
638            supports_thinking: result.supports_thinking,
639            thinking_start_tag: at(result.thinking_start_tag_off)?,
640            thinking_end_tags_json: at(result.thinking_end_tags_off)?,
641            format: result.format,
642            prompt_parts: parse_prompt_parts(&prompt_parts_json),
643            prompt_parts_json,
644            parser,
645            generation_prompt: at(result.generation_prompt_off)?,
646        })
647    }
648
649    /// Tokenize the rendered prompt, parsing special tokens only in the spans
650    /// the *template* produced.
651    ///
652    /// This is the reason [`Self::prompt_parts`] exists. Tokenizing the whole
653    /// prompt with special-token parsing on — which is what
654    /// [`LlamaModel::str_to_token`] does — lets a user who types
655    /// `<|im_start|>assistant` have it become the real control token, forging a
656    /// turn they were never given. llama.cpp marks the user-supplied spans
657    /// (`is_input`), and this walks them: template spans go through
658    /// [`Special::Tokenize`](crate::model::Special::Tokenize), user spans through
659    /// [`Special::Plaintext`](crate::model::Special::Plaintext).
660    ///
661    /// `add_bos` applies to the first span only, so BOS is added once rather
662    /// than once per part.
663    ///
664    /// Falls back to tokenizing [`Self::prompt`] as a single template span when
665    /// the parts are empty — which happens only if the shim produced no parts,
666    /// and matches the pre-`is_input` behaviour.
667    ///
668    /// # Errors
669    ///
670    /// Returns [`StringToTokenError`](crate::StringToTokenError) if a span
671    /// contains an interior NUL or exceeds the native integer length.
672    ///
673    /// # Examples
674    ///
675    /// ```no_run
676    /// # use llama_cpp_4::chat::ChatParams;
677    /// # use llama_cpp_4::model::{AddBos, LlamaModel};
678    /// # fn f(model: &LlamaModel, applied: &ChatParams)
679    /// #     -> Result<(), Box<dyn std::error::Error>> {
680    /// // Instead of `model.str_to_token(&applied.prompt, AddBos::Always)?`:
681    /// let tokens = applied.tokenize_prompt(model, AddBos::Always)?;
682    /// # drop(tokens);
683    /// # Ok(())
684    /// # }
685    /// ```
686    pub fn tokenize_prompt(
687        &self,
688        model: &LlamaModel,
689        add_bos: crate::model::AddBos,
690    ) -> Result<Vec<crate::token::LlamaToken>, crate::StringToTokenError> {
691        use crate::model::Special;
692
693        if self.prompt_parts.is_empty() {
694            return model.str_to_token_with_special(&self.prompt, add_bos, Special::Tokenize);
695        }
696
697        let mut out = Vec::new();
698        let mut scratch = Vec::new();
699        let mut bos = add_bos;
700        for part in &self.prompt_parts {
701            // An empty span would otherwise consume the one-shot BOS without
702            // contributing a token.
703            if part.text.is_empty() {
704                continue;
705            }
706            let special = if part.is_input {
707                Special::Plaintext
708            } else {
709                Special::Tokenize
710            };
711            model.str_to_token_into_with_special(&part.text, bos, special, &mut scratch)?;
712            out.append(&mut scratch);
713            bos = crate::model::AddBos::Never;
714        }
715        Ok(out)
716    }
717
718    /// Human-readable name of this chat format, e.g. `"Hermes 2 Pro"`.
719    ///
720    /// # Errors
721    ///
722    /// Returns [`ChatError::Failed`] if llama.cpp cannot name the format.
723    pub fn format_name(&self) -> Result<String, ChatError> {
724        format_name(self.format)
725    }
726
727    /// Convert [`Self::grammar_triggers`] into the `(patterns, tokens)` pair
728    /// [`LlamaSampler::grammar_lazy_patterns`](crate::sampling::LlamaSampler::grammar_lazy_patterns)
729    /// expects.
730    ///
731    /// The four trigger kinds do not map onto the sampler one-for-one, and
732    /// getting the translation wrong silently produces a grammar that never
733    /// activates:
734    ///
735    /// - `word` is a **literal**, so it is regex-escaped before becoming a
736    ///   pattern. A raw `<tool_call>` would otherwise be a character class.
737    /// - `pattern` passes through unchanged.
738    /// - `pattern_full` is anchored with `^`/`$` unless it already is.
739    /// - `token` becomes a trigger token rather than a pattern.
740    ///
741    /// This mirrors `common/sampling.cpp`, so callers do not have to.
742    #[must_use]
743    pub fn sampler_triggers(&self) -> (Vec<String>, Vec<crate::token::LlamaToken>) {
744        let mut patterns = Vec::new();
745        let mut tokens = Vec::new();
746        for trigger in &self.grammar_triggers {
747            match trigger.kind.as_str() {
748                "word" => patterns.push(regex_escape(&trigger.value)),
749                "pattern" => patterns.push(trigger.value.clone()),
750                "pattern_full" => patterns.push(anchor_pattern(&trigger.value)),
751                "token" => tokens.push(crate::token::LlamaToken(trigger.token)),
752                // An unknown kind from a newer llama.cpp: dropping it is safer
753                // than guessing, and the grammar simply stays dormant for it.
754                _ => {}
755            }
756        }
757        (patterns, tokens)
758    }
759
760    /// The prefix the template already placed at the end of the prompt, e.g.
761    /// `"<|im_start|>assistant\n"`.
762    ///
763    /// This is not part of the model's output, but the grammar and the parser
764    /// are both written as if it were — see [`Self::grammar_sampler`].
765    #[must_use]
766    pub fn generation_prompt(&self) -> &str {
767        &self.generation_prompt
768    }
769
770    /// Build the grammar sampler this render needs, or `None` when
771    /// unconstrained.
772    ///
773    /// Prefer this over constructing the sampler yourself: it handles three
774    /// things that are each silently wrong if missed.
775    ///
776    /// - **Lazy vs eager.** A lazy grammar must be built with its triggers, or
777    ///   it never activates. One built eagerly from a lazy grammar constrains
778    ///   from token zero and blocks a thinking model's reasoning prefix.
779    /// - **Trigger translation.** Literal, regex and token triggers map onto
780    ///   the sampler differently — see [`Self::sampler_triggers`].
781    /// - **Generation-prompt prefill.** llama.cpp writes tool-call grammars to
782    ///   match `generation_prompt + output`, because that is what the parser
783    ///   later sees. The sampler only sees `output`, so without advancing the
784    ///   grammar past that prefix it forces the model to *re-emit*
785    ///   `<|im_start|>assistant` as generated text. Prefill applies only to
786    ///   non-lazy grammars — a lazy one has not started matching yet.
787    ///
788    /// # Panics
789    ///
790    /// Panics if llama.cpp cannot parse the grammar it just produced, or if
791    /// that grammar contains an interior NUL.
792    #[must_use]
793    pub fn grammar_sampler(&self, model: &LlamaModel) -> Option<crate::sampling::LlamaSampler> {
794        use crate::sampling::LlamaSampler;
795
796        if self.grammar.is_empty() {
797            return None;
798        }
799
800        let (patterns, tokens) = self.sampler_triggers();
801        let lazy = self.grammar_lazy && !(patterns.is_empty() && tokens.is_empty());
802
803        let mut sampler = if lazy {
804            let refs: Vec<&str> = patterns.iter().map(String::as_str).collect();
805            LlamaSampler::grammar_lazy_patterns(model, &self.grammar, "root", &refs, &tokens)
806        } else {
807            LlamaSampler::grammar(model, &self.grammar, "root")
808        };
809
810        if !lazy {
811            for token in self.generation_prompt_tokens(model) {
812                sampler.accept(token);
813            }
814        }
815        Some(sampler)
816    }
817
818    /// Tokenize [`Self::generation_prompt`] the way llama.cpp does when
819    /// prefilling a grammar.
820    ///
821    /// Some tokenizers prepend a space to the first token; upstream drops it
822    /// when the prompt itself does not start with whitespace, since that space
823    /// is an artefact rather than something the template emitted.
824    fn generation_prompt_tokens(&self, model: &LlamaModel) -> Vec<crate::token::LlamaToken> {
825        if self.generation_prompt.is_empty() {
826            return Vec::new();
827        }
828        let Ok(tokens) = model.str_to_token(&self.generation_prompt, crate::model::AddBos::Never)
829        else {
830            return Vec::new();
831        };
832        let starts_with_space = self.generation_prompt.starts_with(char::is_whitespace);
833        let mut out = Vec::with_capacity(tokens.len());
834        for (i, token) in tokens.into_iter().enumerate() {
835            if i == 0 && !starts_with_space {
836                if let Ok(piece) = model.token_to_str(token, crate::model::Special::Tokenize) {
837                    if piece.starts_with(char::is_whitespace) {
838                        continue;
839                    }
840                }
841            }
842            out.push(token);
843        }
844        out
845    }
846
847    /// Parse model output back into an OpenAI-shaped message JSON object with
848    /// `role`, `content`, `reasoning_content` and `tool_calls`.
849    ///
850    /// This uses the parser the template produced, so it understands that
851    /// model family's tool-call syntax rather than scraping for a fixed marker.
852    ///
853    /// Set `is_partial` while streaming: the parser then tolerates a truncated
854    /// tail instead of rejecting the buffer.
855    ///
856    /// # Errors
857    ///
858    /// Returns [`ChatError::Failed`] if the output cannot be parsed.
859    pub fn parse(&self, text: &str, is_partial: bool) -> Result<String, ChatError> {
860        let c_text = CString::new(text)?;
861        let c_gen_prompt = CString::new(self.generation_prompt.as_str())?;
862        let params = sys::chat_shim_parse_params {
863            text: c_text.as_ptr(),
864            parser: self
865                .parser
866                .as_ref()
867                .map_or(std::ptr::null(), |p| p.raw.as_ptr().cast_const()),
868            generation_prompt: c_gen_prompt.as_ptr(),
869            format: self.format,
870            is_partial,
871        };
872        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
873        // and the shim reports how much it needed; nothing is read back unless it
874        // reported success.
875        read_string(|buf, len, expected| unsafe {
876            sys::chat_shim_parse(&raw const params, buf, len, expected)
877        })
878    }
879}
880
881/// Human-readable name of a `common_chat_format` discriminant.
882///
883/// # Errors
884///
885/// Returns [`ChatError::Failed`] if llama.cpp cannot name the format.
886pub fn format_name(format: i32) -> Result<String, ChatError> {
887    // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
888    // and the shim reports how much it needed; nothing is read back unless it
889    // reported success.
890    read_string(|buf, len, expected| unsafe {
891        sys::chat_shim_format_name(format, buf, len, expected)
892    })
893}
894
895/// Validate and normalise an OpenAI-shaped `messages` array.
896///
897/// Expands the shorthands llama.cpp accepts — typed content parts, legacy
898/// `function_call` — into the canonical form, and rejects malformed input
899/// before it reaches a template.
900///
901/// # Errors
902///
903/// Returns [`ChatError::BadJson`] if the input is not valid JSON, or
904/// [`ChatError::Failed`] if it is not a valid message array.
905pub fn parse_messages_oaicompat(messages_json: &str) -> Result<String, ChatError> {
906    let c_messages = CString::new(messages_json)?;
907    // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
908    // and the shim reports how much it needed; nothing is read back unless it
909    // reported success.
910    read_string(|buf, len, expected| unsafe {
911        sys::chat_shim_msgs_parse_oaicompat(c_messages.as_ptr(), buf, len, expected)
912    })
913}
914
915/// Validate and normalise an OpenAI-shaped `tools` array into a JSON array of
916/// `{"name", "description", "parameters"}`.
917///
918/// # Errors
919///
920/// Returns [`ChatError::BadJson`] if the input is not valid JSON, or
921/// [`ChatError::Failed`] if it is not a valid tool array.
922pub fn parse_tools_oaicompat(tools_json: &str) -> Result<String, ChatError> {
923    let c_tools = CString::new(tools_json)?;
924    // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
925    // and the shim reports how much it needed; nothing is read back unless it
926    // reported success.
927    read_string(|buf, len, expected| unsafe {
928        sys::chat_shim_tools_parse_oaicompat(c_tools.as_ptr(), buf, len, expected)
929    })
930}
931
932// ── plumbing ────────────────────────────────────────────────────────────────
933
934/// Escape the characters `std::regex` treats as special, matching
935/// `regex_escape` in `common/common.cpp`.
936fn regex_escape(s: &str) -> String {
937    const SPECIAL: &[char] = &[
938        '.', '^', '$', '|', '(', ')', '*', '+', '?', '[', ']', '{', '}', '\\',
939    ];
940    let mut out = String::with_capacity(s.len());
941    for c in s.chars() {
942        if SPECIAL.contains(&c) {
943            out.push('\\');
944        }
945        out.push(c);
946    }
947    out
948}
949
950/// Anchor a `pattern_full` trigger, matching `common/sampling.cpp`. An empty
951/// pattern becomes `^$` — matching only the empty string — rather than `^^$$`.
952fn anchor_pattern(pattern: &str) -> String {
953    if pattern.is_empty() {
954        return "^$".to_owned();
955    }
956    let mut out = String::with_capacity(pattern.len() + 2);
957    if !pattern.starts_with('^') {
958        out.push('^');
959    }
960    out.push_str(pattern);
961    if !pattern.ends_with('$') {
962        out.push('$');
963    }
964    out
965}
966
967fn opt_ptr(s: Option<&CString>) -> *const c_char {
968    s.map_or(std::ptr::null(), |c| c.as_ptr())
969}
970
971/// Pull the fields out of the shim's trigger JSON without taking a JSON
972/// dependency. The shape is fixed by `chat_shim.cpp`, so a hand parser is
973/// enough — anything unexpected yields no triggers, and
974/// [`ChatParams::grammar_triggers_json`] still carries the raw text.
975/// Parse the shim's `[{"text", "is_input"}]` array.
976///
977/// Hand-rolled for the same reason [`parse_triggers`] is: this crate takes no
978/// JSON dependency. The shape is fixed by `chat_shim.cpp`, and
979/// `prompt_parts_reads_shim_output` pins it.
980/// Scan the `[{"text":...,"is_input":...}, ...]` array the shim emits.
981///
982/// Unlike [`parse_triggers`], this cannot chunk on a bare `{`: the text of a
983/// part is whatever the user typed, so a message containing
984/// `{"text":"x","is_input":false}` would split its own span into a chunk that
985/// reads as template text — re-enabling exactly the special-token forging that
986/// `is_input` exists to prevent. So this walks the string properly, tracking
987/// quoting and escapes, and never looks inside a string value.
988fn parse_prompt_parts(json: &str) -> Vec<PromptPart> {
989    let b = json.as_bytes();
990    let mut out = Vec::new();
991
992    let Some(mut i) = b.iter().position(|&c| c == b'[') else {
993        return out;
994    };
995    i += 1;
996
997    loop {
998        while i < b.len() && (b[i].is_ascii_whitespace() || b[i] == b',') {
999            i += 1;
1000        }
1001        if i >= b.len() || b[i] != b'{' {
1002            return out;
1003        }
1004        i += 1;
1005
1006        let mut text = None;
1007        let mut is_input = true;
1008        loop {
1009            while i < b.len() && (b[i].is_ascii_whitespace() || b[i] == b',') {
1010                i += 1;
1011            }
1012            if i >= b.len() {
1013                return out;
1014            }
1015            if b[i] == b'}' {
1016                i += 1;
1017                break;
1018            }
1019
1020            let Some((key, next)) = scan_json_string(json, i) else {
1021                return out;
1022            };
1023            i = next;
1024            while i < b.len() && b[i].is_ascii_whitespace() {
1025                i += 1;
1026            }
1027            if i >= b.len() || b[i] != b':' {
1028                return out;
1029            }
1030            i += 1;
1031            while i < b.len() && b[i].is_ascii_whitespace() {
1032                i += 1;
1033            }
1034            if i >= b.len() {
1035                return out;
1036            }
1037
1038            if b[i] == b'"' {
1039                let Some((value, next)) = scan_json_string(json, i) else {
1040                    return out;
1041                };
1042                i = next;
1043                if key == "text" {
1044                    text = Some(value);
1045                }
1046            } else {
1047                // A bare literal: `true`, `false`, or a number. Only `is_input`
1048                // is expected here, but skip any other the same way.
1049                let start = i;
1050                while i < b.len() && b[i] != b',' && b[i] != b'}' {
1051                    i += 1;
1052                }
1053                if key == "is_input" {
1054                    // Absent or unrecognised reads as `true`: treating a span as
1055                    // user input only ever *disables* special-token parsing,
1056                    // which is the safe direction.
1057                    is_input = json[start..i].trim() != "false";
1058                }
1059            }
1060        }
1061
1062        if let Some(text) = text {
1063            out.push(PromptPart { text, is_input });
1064        }
1065    }
1066}
1067
1068/// Read the JSON string starting at `start` (which must be its opening quote),
1069/// returning it unescaped along with the index just past its closing quote.
1070fn scan_json_string(json: &str, start: usize) -> Option<(String, usize)> {
1071    let b = json.as_bytes();
1072    if b.get(start) != Some(&b'"') {
1073        return None;
1074    }
1075    let mut value = String::new();
1076    let mut chars = json[start + 1..].char_indices();
1077    while let Some((offset, c)) = chars.next() {
1078        match c {
1079            '"' => return Some((value, start + 1 + offset + 1)),
1080            '\\' => match chars.next()?.1 {
1081                'n' => value.push('\n'),
1082                'r' => value.push('\r'),
1083                't' => value.push('\t'),
1084                'b' => value.push('\u{8}'),
1085                'f' => value.push('\u{c}'),
1086                'u' => {
1087                    let hex: String = (0..4).map_while(|_| chars.next().map(|(_, c)| c)).collect();
1088                    let code = u32::from_str_radix(&hex, 16).ok()?;
1089                    value.push(char::from_u32(code)?);
1090                }
1091                other => value.push(other),
1092            },
1093            other => value.push(other),
1094        }
1095    }
1096    None
1097}
1098
1099fn parse_triggers(json: &str) -> Vec<GrammarTrigger> {
1100    let mut out = Vec::new();
1101    for chunk in json.split('{').skip(1) {
1102        let kind = json_str_field(chunk, "type");
1103        let value = json_str_field(chunk, "value");
1104        let token = json_int_field(chunk, "token");
1105        if let (Some(kind), Some(value)) = (kind, value) {
1106            out.push(GrammarTrigger {
1107                kind,
1108                value,
1109                token: token.unwrap_or(-1),
1110            });
1111        }
1112    }
1113    out
1114}
1115
1116fn json_str_field(chunk: &str, key: &str) -> Option<String> {
1117    let needle = format!("\"{key}\":");
1118    let rest = &chunk[chunk.find(&needle)? + needle.len()..];
1119    let rest = rest.trim_start();
1120    let mut chars = rest.strip_prefix('"')?.chars();
1121    let mut value = String::new();
1122    while let Some(c) = chars.next() {
1123        match c {
1124            '"' => return Some(value),
1125            '\\' => match chars.next()? {
1126                'n' => value.push('\n'),
1127                'r' => value.push('\r'),
1128                't' => value.push('\t'),
1129                'u' => {
1130                    let hex: String = chars.by_ref().take(4).collect();
1131                    let code = u32::from_str_radix(&hex, 16).ok()?;
1132                    value.push(char::from_u32(code)?);
1133                }
1134                other => value.push(other),
1135            },
1136            other => value.push(other),
1137        }
1138    }
1139    None
1140}
1141
1142fn json_int_field(chunk: &str, key: &str) -> Option<i32> {
1143    let needle = format!("\"{key}\":");
1144    let rest = &chunk[chunk.find(&needle)? + needle.len()..];
1145    let rest = rest.trim_start();
1146    let end = rest
1147        .find(|c: char| !c.is_ascii_digit() && c != '-')
1148        .unwrap_or(rest.len());
1149    rest[..end].parse().ok()
1150}
1151
1152#[cfg(test)]
1153mod tests {
1154    use super::*;
1155
1156    /// The schema the `structured` example ships in its docs must convert and
1157    /// be loadable by the sampler — an object schema with a required key is the
1158    /// shape every `response_format: json_schema` caller sends.
1159    #[test]
1160    fn realistic_object_schema_converts() {
1161        let gbnf = json_schema_to_grammar(
1162            r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#,
1163            false,
1164        )
1165        .unwrap();
1166        eprintln!("GBNF:\n{gbnf}");
1167        assert!(gbnf.contains("root"));
1168    }
1169
1170    #[test]
1171    fn json_schema_to_grammar_produces_a_root_rule() {
1172        let gbnf = json_schema_to_grammar(r#"{"type":"integer"}"#, false).unwrap();
1173        assert!(gbnf.contains("root"), "no root rule in: {gbnf}");
1174    }
1175
1176    /// An object schema must constrain its keys, otherwise the grammar is not
1177    /// actually enforcing the schema.
1178    #[test]
1179    fn json_schema_to_grammar_constrains_object_keys() {
1180        let gbnf = json_schema_to_grammar(
1181            r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#,
1182            false,
1183        )
1184        .unwrap();
1185        assert!(gbnf.contains("city"), "key not in grammar: {gbnf}");
1186    }
1187
1188    /// Invalid JSON must be reported as such rather than crossing the FFI
1189    /// boundary as an exception.
1190    #[test]
1191    fn json_schema_to_grammar_rejects_invalid_json() {
1192        let err = json_schema_to_grammar("{not json", false).unwrap_err();
1193        assert!(
1194            matches!(err, ChatError::BadJson(_)),
1195            "expected BadJson, got {err:?}"
1196        );
1197    }
1198
1199    #[test]
1200    fn json_schema_to_grammar_rejects_empty_input() {
1201        assert!(json_schema_to_grammar("", false).is_err());
1202    }
1203
1204    /// Interior NULs must be caught in Rust — a `CString` cannot hold one, and
1205    /// truncating silently would change the schema.
1206    #[test]
1207    fn json_schema_to_grammar_rejects_interior_nul() {
1208        let err = json_schema_to_grammar("{\"type\":\"in\0teger\"}", false).unwrap_err();
1209        assert!(matches!(err, ChatError::Nul(_)), "got {err:?}");
1210    }
1211
1212    #[test]
1213    fn tool_choice_parses_openai_values() {
1214        assert_eq!(
1215            ToolChoice::parse_oaicompat("auto").unwrap(),
1216            ToolChoice::Auto
1217        );
1218        assert_eq!(
1219            ToolChoice::parse_oaicompat("required").unwrap(),
1220            ToolChoice::Required
1221        );
1222        assert_eq!(
1223            ToolChoice::parse_oaicompat("none").unwrap(),
1224            ToolChoice::None
1225        );
1226    }
1227
1228    #[test]
1229    fn tool_choice_rejects_unknown_value() {
1230        assert!(ToolChoice::parse_oaicompat("sometimes").is_err());
1231    }
1232
1233    #[test]
1234    fn tools_parse_oaicompat_extracts_name_and_parameters() {
1235        let normalised = parse_tools_oaicompat(
1236            r#"[{"type":"function","function":{"name":"get_weather",
1237                "description":"Get weather","parameters":{"type":"object"}}}]"#,
1238        )
1239        .unwrap();
1240        assert!(normalised.contains("get_weather"), "got {normalised}");
1241    }
1242
1243    #[test]
1244    fn tools_parse_oaicompat_rejects_garbage() {
1245        assert!(parse_tools_oaicompat("[[[").is_err());
1246    }
1247
1248    #[test]
1249    fn messages_parse_oaicompat_roundtrips_a_simple_turn() {
1250        let normalised = parse_messages_oaicompat(r#"[{"role":"user","content":"hi"}]"#).unwrap();
1251        assert!(normalised.contains("user"), "got {normalised}");
1252        assert!(normalised.contains("hi"), "got {normalised}");
1253    }
1254
1255    #[test]
1256    fn messages_parse_oaicompat_rejects_non_array() {
1257        assert!(parse_messages_oaicompat(r#"{"role":"user"}"#).is_err());
1258    }
1259
1260    /// The trigger parser is hand-rolled, so pin it against the exact shape
1261    /// `chat_shim.cpp` emits, including escapes.
1262    #[test]
1263    fn trigger_parser_reads_shim_output() {
1264        let json = r#"[{"type":"word","value":"<tool_call>","token":-1},
1265                       {"type":"token","value":"a\nb","token":42}]"#;
1266        let triggers = parse_triggers(json);
1267        assert_eq!(triggers.len(), 2);
1268        assert_eq!(triggers[0].kind, "word");
1269        assert_eq!(triggers[0].value, "<tool_call>");
1270        assert_eq!(triggers[0].token, -1);
1271        assert_eq!(triggers[1].kind, "token");
1272        assert_eq!(triggers[1].value, "a\nb");
1273        assert_eq!(triggers[1].token, 42);
1274    }
1275
1276    /// A literal trigger like `<tool_call>` must be escaped before it becomes
1277    /// a regex — unescaped, `[` and `]` would make it a character class and the
1278    /// grammar would never fire.
1279    #[test]
1280    fn word_triggers_are_regex_escaped() {
1281        let params = ChatParams {
1282            grammar_triggers: vec![GrammarTrigger {
1283                kind: "word".to_owned(),
1284                value: "a.b[c]".to_owned(),
1285                token: -1,
1286            }],
1287            ..stub_params()
1288        };
1289        let (patterns, tokens) = params.sampler_triggers();
1290        assert_eq!(patterns, vec![r"a\.b\[c\]".to_owned()]);
1291        assert!(tokens.is_empty());
1292    }
1293
1294    #[test]
1295    fn pattern_triggers_pass_through_unescaped() {
1296        let params = ChatParams {
1297            grammar_triggers: vec![GrammarTrigger {
1298                kind: "pattern".to_owned(),
1299                value: "a.b".to_owned(),
1300                token: -1,
1301            }],
1302            ..stub_params()
1303        };
1304        assert_eq!(params.sampler_triggers().0, vec!["a.b".to_owned()]);
1305    }
1306
1307    #[test]
1308    fn pattern_full_triggers_are_anchored_once() {
1309        let cases = [
1310            ("abc", "^abc$"),
1311            ("^abc", "^abc$"),
1312            ("abc$", "^abc$"),
1313            ("^abc$", "^abc$"),
1314            ("", "^$"),
1315        ];
1316        for (input, want) in cases {
1317            let params = ChatParams {
1318                grammar_triggers: vec![GrammarTrigger {
1319                    kind: "pattern_full".to_owned(),
1320                    value: input.to_owned(),
1321                    token: -1,
1322                }],
1323                ..stub_params()
1324            };
1325            assert_eq!(
1326                params.sampler_triggers().0,
1327                vec![want.to_owned()],
1328                "anchoring {input:?}"
1329            );
1330        }
1331    }
1332
1333    #[test]
1334    fn token_triggers_become_tokens_not_patterns() {
1335        let params = ChatParams {
1336            grammar_triggers: vec![GrammarTrigger {
1337                kind: "token".to_owned(),
1338                value: String::new(),
1339                token: 42,
1340            }],
1341            ..stub_params()
1342        };
1343        let (patterns, tokens) = params.sampler_triggers();
1344        assert!(patterns.is_empty());
1345        assert_eq!(tokens, vec![crate::token::LlamaToken(42)]);
1346    }
1347
1348    /// A kind from a newer llama.cpp must be dropped rather than guessed at.
1349    #[test]
1350    fn unknown_trigger_kinds_are_dropped() {
1351        let params = ChatParams {
1352            grammar_triggers: vec![GrammarTrigger {
1353                kind: "something_new".to_owned(),
1354                value: "x".to_owned(),
1355                token: -1,
1356            }],
1357            ..stub_params()
1358        };
1359        let (patterns, tokens) = params.sampler_triggers();
1360        assert!(patterns.is_empty());
1361        assert!(tokens.is_empty());
1362    }
1363
1364    fn stub_params() -> ChatParams {
1365        ChatParams {
1366            prompt: String::new(),
1367            grammar: String::new(),
1368            grammar_lazy: false,
1369            grammar_triggers_json: String::new(),
1370            grammar_triggers: Vec::new(),
1371            preserved_tokens_json: String::new(),
1372            additional_stops_json: String::new(),
1373            supports_thinking: false,
1374            thinking_start_tag: String::new(),
1375            thinking_end_tags_json: String::new(),
1376            format: 0,
1377            prompt_parts: Vec::new(),
1378            prompt_parts_json: String::new(),
1379            parser: None,
1380            generation_prompt: String::new(),
1381        }
1382    }
1383
1384    /// The shape is fixed by `chat_shim.cpp`, so pin it against a hand-rolled
1385    /// parser the way the trigger test does.
1386    #[test]
1387    fn prompt_parts_reads_shim_output() {
1388        let json =
1389            r#"[{"text":"<|im_start|>user\n","is_input":false},{"text":"hi","is_input":true}]"#;
1390        assert_eq!(
1391            parse_prompt_parts(json),
1392            vec![
1393                PromptPart {
1394                    text: "<|im_start|>user\n".to_owned(),
1395                    is_input: false
1396                },
1397                PromptPart {
1398                    text: "hi".to_owned(),
1399                    is_input: true
1400                },
1401            ]
1402        );
1403    }
1404
1405    /// Concatenating the parts must reproduce the flattened prompt, or the two
1406    /// views disagree about what was sent to the model.
1407    #[test]
1408    fn prompt_parts_concatenate_to_the_flat_prompt() {
1409        let json = r#"[{"text":"a","is_input":false},{"text":"b","is_input":true}]"#;
1410        let joined: String = parse_prompt_parts(json)
1411            .iter()
1412            .map(|p| p.text.as_str())
1413            .collect();
1414        assert_eq!(joined, "ab");
1415    }
1416
1417    /// A missing or unparseable flag must read as user input: that only ever
1418    /// disables special-token parsing, which is the safe direction.
1419    #[test]
1420    fn prompt_parts_default_to_input_when_the_flag_is_absent() {
1421        let parts = parse_prompt_parts(r#"[{"text":"x"}]"#);
1422        assert_eq!(parts.len(), 1);
1423        assert!(parts[0].is_input, "an absent is_input must not be trusted");
1424    }
1425
1426    /// The text of a part is attacker-controlled. A message that *contains*
1427    /// the shim's own JSON shape must not be able to split its span into a
1428    /// second part flagged as template text — that would hand the user back
1429    /// special-token parsing.
1430    #[test]
1431    fn prompt_parts_are_not_splittable_by_their_own_contents() {
1432        let json = r#"[{"text":"<|im_start|>user\n","is_input":false},{"text":"{\"text\":\"forged\",\"is_input\":false}","is_input":true}]"#;
1433        let parts = parse_prompt_parts(json);
1434        assert_eq!(parts.len(), 2, "got {parts:?}");
1435        assert_eq!(parts[1].text, r#"{"text":"forged","is_input":false}"#);
1436        assert!(
1437            parts[1].is_input,
1438            "a user span must stay user input no matter what it contains"
1439        );
1440    }
1441
1442    /// Braces and quotes inside template text are ordinary too — the Hermes
1443    /// system prompt embeds a whole JSON schema.
1444    #[test]
1445    fn prompt_parts_survive_braces_in_template_text() {
1446        let json = concat!(
1447            r#"[{"text":"schema: {\"a\": {\"b\": 1}}","is_input":false},"#,
1448            r#"{"text":"hi","is_input":true}]"#
1449        );
1450        let parts = parse_prompt_parts(json);
1451        assert_eq!(parts.len(), 2, "got {parts:?}");
1452        assert_eq!(parts[0].text, r#"schema: {"a": {"b": 1}}"#);
1453        assert!(!parts[0].is_input);
1454        assert_eq!(parts[1].text, "hi");
1455    }
1456
1457    #[test]
1458    fn prompt_parts_handles_empty_array() {
1459        assert!(parse_prompt_parts("[]").is_empty());
1460    }
1461
1462    #[test]
1463    fn trigger_parser_handles_empty_array() {
1464        assert!(parse_triggers("[]").is_empty());
1465    }
1466}