use std::sync::Arc;
use ferrox_edge::{
derive_think_gears, probe_effort_profile, probe_thinking_profile, EffortProfile, ThinkGears,
};
use ferrox_models::chat_template::{
BuiltinTemplate, ChatTemplate as JinjaTemplate, RenderOptions, TemplateError,
};
use serde_json::{json, Map, Value};
use crate::{ChatMessage, ToolDef};
#[derive(Clone)]
pub(crate) struct PromptTemplate {
inner: Arc<Inner>,
}
struct Inner {
template: JinjaTemplate,
end_of_turn: Option<&'static str>,
bos_token: Option<String>,
eos_token: Option<String>,
thinking: ferrox_edge::ThinkingProfile,
handles_tools: bool,
}
impl std::fmt::Debug for PromptTemplate {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "PromptTemplate({})", self.inner.template.describe())
}
}
impl PromptTemplate {
pub(crate) fn from_gguf_metadata(
source: Option<&str>,
arch: Option<&str>,
byte_tokenizer: bool,
bos_token: Option<String>,
eos_token: Option<String>,
) -> Self {
Self::build(
JinjaTemplate::from_gguf_metadata(source, arch, byte_tokenizer),
source,
bos_token,
eos_token,
)
}
pub(crate) fn from_source(
source: Option<&str>,
bos_token: Option<String>,
eos_token: Option<String>,
) -> Self {
Self::build(
JinjaTemplate::from_gguf_metadata(source, None, false),
source,
bos_token,
eos_token,
)
}
pub(crate) fn plain() -> Self {
Self::build(
JinjaTemplate::builtin(BuiltinTemplate::Plain),
None,
None,
None,
)
}
fn build(
template: JinjaTemplate,
source: Option<&str>,
bos_token: Option<String>,
eos_token: Option<String>,
) -> Self {
let (bos, eos) = (bos_token.as_deref(), eos_token.as_deref());
let thinking = probe_thinking_profile(
probe_render(&template, bos, eos),
probe_efforts(&template, bos, eos),
);
let handles_tools = probe_tools_consumed(&template, bos, eos);
Self {
inner: Arc::new(Inner {
end_of_turn: end_of_turn_marker(source),
template,
bos_token,
eos_token,
thinking,
handles_tools,
}),
}
}
pub(crate) fn describe(&self) -> String {
self.inner.template.describe()
}
pub(crate) fn end_of_turn(&self) -> Option<&'static str> {
self.inner.end_of_turn
}
pub(crate) fn handles_tools(&self) -> bool {
self.inner.handles_tools
}
pub(crate) fn efforts(&self) -> &EffortProfile {
&self.inner.thinking.efforts
}
pub(crate) fn think_gears(&self, parser_configured: bool) -> ThinkGears {
derive_think_gears(&self.inner.thinking, parser_configured)
}
pub(crate) fn render(
&self,
messages: &[ChatMessage],
tools: &[ToolDef],
extra: Map<String, Value>,
) -> Result<String, TemplateError> {
let structured = self.handles_tools() && !tools.is_empty();
let json: Vec<Value> = messages
.iter()
.map(|m| message_json(m, structured))
.collect();
let opts = RenderOptions {
add_generation_prompt: true,
bos_token: self.inner.bos_token.clone(),
eos_token: self.inner.eos_token.clone(),
tools: if structured {
tools.iter().map(tool_json).collect()
} else {
Vec::new()
},
extra,
};
self.inner.template.render(&json, &opts)
}
}
fn message_json(m: &ChatMessage, structured: bool) -> Value {
let mut obj = Map::new();
obj.insert("role".into(), json!(m.role));
let text = m
.content
.as_ref()
.map(crate::MessageContent::as_text)
.unwrap_or_default();
match (&m.tool_calls, structured) {
(Some(calls), true) => {
obj.insert("content".into(), json!(text));
obj.insert(
"tool_calls".into(),
Value::Array(calls.iter().map(tool_call_json).collect()),
);
}
(Some(_), false) => {
obj.insert("content".into(), json!(m.rendered_content()));
}
(None, _) => {
obj.insert("content".into(), json!(text));
}
}
if let Some(id) = &m.tool_call_id {
obj.insert("tool_call_id".into(), json!(id));
}
if let Some(reasoning) = m.reasoning_content.as_deref().filter(|r| !r.is_empty()) {
obj.insert("reasoning_content".into(), json!(reasoning));
obj.insert("reasoning".into(), json!(reasoning));
let harmony_would_raise = m.tool_calls.is_some() && !text.is_empty();
if !harmony_would_raise {
obj.insert("thinking".into(), json!(reasoning));
}
}
Value::Object(obj)
}
fn tool_call_json(call: &crate::ToolCallIn) -> Value {
let raw = call.function.arguments.as_str();
let args = match serde_json::from_str::<Value>(raw) {
Ok(v @ Value::Object(_)) => v,
_ => json!(raw),
};
json!({
"type": "function",
"id": call.id,
"function": {"name": call.function.name, "arguments": args},
})
}
pub(crate) fn tool_json(t: &ToolDef) -> Value {
json!({
"type": "function",
"function": {
"name": t.function.name,
"description": t.function.description.as_deref().unwrap_or(""),
"parameters": t.function.parameters.clone().unwrap_or_else(|| json!({
"type": "object",
"properties": {},
})),
},
})
}
fn probe_tools_consumed(
template: &JinjaTemplate,
bos_token: Option<&str>,
eos_token: Option<&str>,
) -> bool {
let messages = vec![json!({"role": "user", "content": "probe"})];
let render = |tools: Vec<Value>| {
template.render(
&messages,
&RenderOptions {
add_generation_prompt: true,
bos_token: bos_token.map(str::to_string),
eos_token: eos_token.map(str::to_string),
tools,
extra: Map::new(),
},
)
};
let probe = json!({
"type": "function",
"function": {
"name": "ferrox_probe_tool",
"description": "No-op probe tool.",
"parameters": {"type": "object", "properties": {}},
},
});
match (render(Vec::new()), render(vec![probe])) {
(Ok(without), Ok(with)) => with != without,
(_, Err(_)) => false,
(Err(_), Ok(_)) => true,
}
}
fn probe_render<'a>(
template: &'a JinjaTemplate,
bos_token: Option<&'a str>,
eos_token: Option<&'a str>,
) -> impl FnMut(&Map<String, Value>, Option<&[Value]>) -> Result<String, TemplateError> + 'a {
let messages = vec![json!({"role": "user", "content": "probe"})];
move |kwargs, tools| {
let opts = RenderOptions {
add_generation_prompt: true,
bos_token: bos_token.map(str::to_string),
eos_token: eos_token.map(str::to_string),
tools: tools.map(<[Value]>::to_vec).unwrap_or_default(),
extra: kwargs.clone(),
};
template.render(&messages, &opts)
}
}
fn probe_efforts(
template: &JinjaTemplate,
bos_token: Option<&str>,
eos_token: Option<&str>,
) -> EffortProfile {
probe_effort_profile(probe_render(template, bos_token, eos_token))
}
fn end_of_turn_marker(source: Option<&str>) -> Option<&'static str> {
let src = source?;
if src.contains("<|turn>") || src.contains("<turn|>") {
Some("<turn|>")
} else if src.contains("<start_of_turn>") {
Some("<end_of_turn>")
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::MessageContent;
fn tool_def(name: &str) -> ToolDef {
serde_json::from_value(json!({
"type": "function",
"function": {
"name": name,
"description": "a tool",
"parameters": {"type": "object", "properties": {}},
},
}))
.expect("tool def")
}
fn tool_call_in(name: &str, arguments: &str) -> crate::ToolCallIn {
serde_json::from_value(json!({
"id": "call_0",
"type": "function",
"function": {"name": name, "arguments": arguments},
}))
.expect("tool call")
}
fn msg(role: &str, content: &str) -> ChatMessage {
ChatMessage {
role: role.to_string(),
content: Some(MessageContent::Text(content.to_string())),
tool_calls: None,
tool_call_id: None,
reasoning_content: None,
}
}
const CHATML: &str = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}";
#[test]
fn a_template_the_old_sniffer_could_not_recognise_now_renders_correctly() {
let mistral = "{% for m in messages %}{% if m.role == 'user' %}[INST] {{ m.content }} [/INST]{% else %}{{ m.content }}</s>{% endif %}{% endfor %}";
let tmpl =
PromptTemplate::from_gguf_metadata(Some(mistral), Some("llama"), false, None, None);
let rendered = tmpl
.render(&[msg("user", "hi")], &[], Map::new())
.expect("renders");
assert_eq!(rendered, "[INST] hi [/INST]");
}
#[test]
fn a_checkpoint_with_no_template_falls_back_to_chatml_like_llama_cpp() {
let tmpl = PromptTemplate::from_gguf_metadata(None, Some("olmoe"), false, None, None);
let rendered = tmpl
.render(&[msg("user", "hi")], &[], Map::new())
.expect("renders");
assert_eq!(
rendered,
"<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n"
);
}
#[test]
fn a_byte_tokenizer_checkpoint_falls_back_to_role_labeled_lines() {
let tmpl = PromptTemplate::from_gguf_metadata(None, Some("olmoe"), true, None, None);
let rendered = tmpl
.render(
&[msg("user", "hello"), msg("assistant", "hi back")],
&[],
Map::new(),
)
.expect("renders");
assert_eq!(rendered, "user: hello\nassistant: hi back");
}
#[test]
fn chat_template_kwargs_reach_the_template() {
let src = "{% if enable_thinking %}THINK{% endif %}{{ messages[0].content }}";
let tmpl = PromptTemplate::from_gguf_metadata(Some(src), Some("qwen3"), false, None, None);
let mut extra = Map::new();
extra.insert("enable_thinking".into(), json!(true));
assert_eq!(
tmpl.render(&[msg("user", "hi")], &[], extra)
.expect("renders"),
"THINKhi"
);
assert_eq!(
tmpl.render(&[msg("user", "hi")], &[], Map::new())
.expect("renders"),
"hi"
);
}
#[test]
fn a_template_that_reads_tools_is_given_them_structurally() {
let src =
"{% for t in tools %}TOOL:{{ t.function.name }}{% endfor %}{{ messages[0].content }}";
let tmpl = PromptTemplate::from_gguf_metadata(Some(src), Some("qwen3"), false, None, None);
assert!(tmpl.handles_tools());
let tools = vec![tool_def("get_weather")];
assert_eq!(
tmpl.render(&[msg("user", "hi")], &tools, Map::new())
.expect("renders"),
"TOOL:get_weatherhi"
);
}
#[test]
fn a_template_that_only_mentions_tools_does_not_count_as_handling_them() {
let mentions = "{# tools are described by the system prompt #}\
{% for m in messages %}{{ m.role }}: {{ m.content }}\n{% endfor %}";
let tmpl =
PromptTemplate::from_gguf_metadata(Some(mentions), Some("qwen2"), false, None, None);
assert!(
!tmpl.handles_tools(),
"the word alone must not skip the preamble"
);
}
#[test]
fn a_template_that_ignores_tools_sees_replayed_calls_as_marker_text() {
let tmpl =
PromptTemplate::from_gguf_metadata(Some(CHATML), Some("qwen2"), false, None, None);
assert!(!tmpl.handles_tools());
let replayed = ChatMessage {
role: "assistant".to_string(),
content: None,
tool_calls: Some(vec![tool_call_in("get_weather", r#"{"city":"Paris"}"#)]),
tool_call_id: None,
reasoning_content: None,
};
let rendered = tmpl.render(&[replayed], &[], Map::new()).expect("renders");
assert!(
rendered.contains(
r#"<tool_call>{"name": "get_weather", "arguments": {"city":"Paris"}}</tool_call>"#
),
"{rendered}"
);
}
#[test]
fn replayed_call_arguments_reach_a_structural_template_as_an_object() {
let src = "{% for m in messages %}{% for c in m.tool_calls %}{{ c.function.arguments.city }}{% endfor %}{% endfor %}tools:{{ tools | length }}";
let tmpl = PromptTemplate::from_gguf_metadata(Some(src), Some("qwen3"), false, None, None);
let replayed = ChatMessage {
role: "assistant".to_string(),
content: None,
tool_calls: Some(vec![tool_call_in("get_weather", r#"{"city":"Paris"}"#)]),
tool_call_id: None,
reasoning_content: None,
};
let tools = vec![tool_def("get_weather")];
assert_eq!(
tmpl.render(&[replayed], &tools, Map::new())
.expect("renders"),
"Paristools:1"
);
}
#[test]
fn a_broken_template_fails_the_request_instead_of_guessing() {
let tmpl = PromptTemplate::from_gguf_metadata(
Some("{% for m in messages %}"),
None,
false,
None,
None,
);
assert!(tmpl.render(&[msg("user", "hi")], &[], Map::new()).is_err());
}
#[test]
fn gemma_families_contribute_their_end_of_turn_marker_to_the_stop_set() {
assert_eq!(
end_of_turn_marker(Some("{{ bos_token }}<start_of_turn>user\n")),
Some("<end_of_turn>")
);
assert_eq!(
end_of_turn_marker(Some("{{- '<|turn>' + m.role -}}")),
Some("<turn|>")
);
assert_eq!(end_of_turn_marker(Some(CHATML)), None);
assert_eq!(end_of_turn_marker(None), None);
}
#[test]
fn the_effort_vocabulary_is_probed_at_load() {
let graded = "{% set allowed = ['low','medium','high'] %}\
{% if reasoning_effort %}\
{% if reasoning_effort not in allowed %}{{ raise_exception('bad effort') }}{% endif %}\
E:{{ reasoning_effort }}\
{% endif %}{{ messages[0].content }}";
let tmpl =
PromptTemplate::from_gguf_metadata(Some(graded), Some("qwen3"), false, None, None);
let profile = tmpl.efforts();
assert!(profile.consumes_effort);
assert!(profile.validates);
assert_eq!(
profile
.supported
.iter()
.map(|e| e.as_str())
.collect::<Vec<_>>(),
vec!["low", "medium", "high"]
);
let inert =
PromptTemplate::from_gguf_metadata(Some(CHATML), Some("qwen2"), false, None, None);
assert!(!inert.efforts().consumes_effort);
}
#[test]
fn a_replayed_chain_of_thought_reaches_a_template_under_both_spellings() {
let mut m = msg("assistant", "the answer is 4");
m.reasoning_content = Some("2 + 2".to_string());
let json = message_json(&m, false);
assert_eq!(json["reasoning_content"], "2 + 2");
assert_eq!(json["reasoning"], "2 + 2");
assert_eq!(json["thinking"], "2 + 2");
assert_eq!(
json["content"], "the answer is 4",
"reasoning must never be folded into what the model said"
);
}
#[test]
fn a_tool_call_turn_with_visible_text_withholds_only_the_harmony_spelling() {
let mut m = msg("assistant", "checking the weather");
m.reasoning_content = Some("I should call the tool".to_string());
m.tool_calls = Some(vec![tool_call_in("get_weather", "{}")]);
let json = message_json(&m, true);
assert!(
json.get("thinking").is_none(),
"harmony would raise on this turn"
);
assert_eq!(json["reasoning_content"], "I should call the tool");
assert_eq!(json["reasoning"], "I should call the tool");
m.content = None;
assert_eq!(message_json(&m, true)["thinking"], "I should call the tool");
}
#[test]
fn an_empty_or_absent_chain_of_thought_puts_no_key_in_front_of_a_template() {
for empty in [None, Some(String::new())] {
let mut m = msg("assistant", "hi");
m.reasoning_content = empty;
let json = message_json(&m, false);
assert!(json.get("reasoning_content").is_none());
assert!(json.get("reasoning").is_none());
assert!(json.get("thinking").is_none());
}
}
#[test]
fn a_replayed_turn_is_accepted_under_either_spelling_of_the_key() {
for key in ["reasoning_content", "reasoning"] {
let m: ChatMessage = serde_json::from_value(json!({
"role": "assistant",
"content": "4",
key: "2 + 2",
}))
.expect("deserializes");
assert_eq!(m.reasoning_content.as_deref(), Some("2 + 2"), "{key}");
}
}
}