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}