1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
impl ChatTemplateEngine for RawTemplate {
fn format_message(&self, _role: &str, content: &str) -> Result<String, RealizarError> {
// Sanitize content to prevent prompt injection (F-SEC-220)
// Even raw templates should sanitize to prevent special token attacks
Ok(sanitize_special_tokens(content))
}
fn format_conversation(&self, messages: &[ChatMessage]) -> Result<String, RealizarError> {
// Sanitize content to prevent prompt injection (F-SEC-220)
let result: String = messages
.iter()
.map(|m| sanitize_special_tokens(&m.content))
.collect();
Ok(result)
}
fn special_tokens(&self) -> &SpecialTokens {
&self.special_tokens
}
fn format(&self) -> TemplateFormat {
TemplateFormat::Raw
}
fn supports_system_prompt(&self) -> bool {
true
}
}
// ============================================================================
// Auto-Detection
// ============================================================================
/// Auto-detect template format from model name or path
///
/// # Arguments
/// * `model_name` - Model name or path (e.g., "TinyLlama/TinyLlama-1.1B-Chat")
///
/// # Returns
/// Detected `TemplateFormat`
///
/// # Example
///
/// ```
/// use realizar::chat_template::{detect_format_from_name, TemplateFormat};
///
/// assert_eq!(detect_format_from_name("TinyLlama-1.1B-Chat"), TemplateFormat::Zephyr);
/// assert_eq!(detect_format_from_name("Qwen2-0.5B-Instruct"), TemplateFormat::ChatML);
/// ```
#[must_use]
pub fn detect_format_from_name(model_name: &str) -> TemplateFormat {
let name_lower = model_name.to_lowercase();
// Pattern rules ordered by specificity (more specific patterns first)
// Format: (patterns, format) - check patterns before formats that share prefixes
// PMAT-181: Qwen3 gets special no-think template (before generic "qwen" match)
if name_lower.contains("qwen3") {
return TemplateFormat::Qwen3NoThink;
}
let rules: &[(&[&str], TemplateFormat)] = &[
// ChatML: Qwen (2.x), OpenHermes, Yi
(&["qwen", "openhermes", "yi-"], TemplateFormat::ChatML),
// Zephyr: TinyLlama, Zephyr, StableLM (check BEFORE llama!)
(&["tinyllama", "zephyr", "stablelm"], TemplateFormat::Zephyr),
// Mistral/Mixtral (check before LLaMA since both use [INST])
(&["mistral", "mixtral"], TemplateFormat::Mistral),
// LLaMA 2 / Vicuna
(&["llama", "vicuna"], TemplateFormat::Llama2),
// Phi variants
(&["phi-", "phi2", "phi3"], TemplateFormat::Phi),
// Alpaca
(&["alpaca"], TemplateFormat::Alpaca),
];
for (patterns, format) in rules {
if patterns.iter().any(|p| name_lower.contains(p)) {
return *format;
}
}
TemplateFormat::Raw
}
/// Auto-detect template format from special tokens
#[must_use]
pub fn detect_format_from_tokens(special_tokens: &SpecialTokens) -> TemplateFormat {
if special_tokens.im_start_token.is_some() || special_tokens.im_end_token.is_some() {
return TemplateFormat::ChatML;
}
if special_tokens.inst_start.is_some() || special_tokens.inst_end.is_some() {
return TemplateFormat::Llama2;
}
TemplateFormat::Raw
}
/// Create a template engine for a given format
#[must_use]
pub fn create_template(format: TemplateFormat) -> Box<dyn ChatTemplateEngine> {
match format {
TemplateFormat::ChatML => Box::new(ChatMLTemplate::new()),
TemplateFormat::Qwen3NoThink => Box::new(Qwen3NoThinkTemplate::new()),
TemplateFormat::Llama2 => Box::new(Llama2Template::new()),
TemplateFormat::Zephyr => Box::new(ZephyrTemplate::new()),
TemplateFormat::Mistral => Box::new(MistralTemplate::new()),
TemplateFormat::Phi => Box::new(PhiTemplate::new()),
TemplateFormat::Alpaca => Box::new(AlpacaTemplate::new()),
TemplateFormat::Custom | TemplateFormat::Raw => Box::new(RawTemplate::new()),
}
}
/// Auto-detect and create template from model name
#[must_use]
pub fn auto_detect_template(model_name: &str) -> Box<dyn ChatTemplateEngine> {
let format = detect_format_from_name(model_name);
create_template(format)
}
/// Format chat messages using auto-detected template
///
/// This is the main entry point for the API. It replaces the naive
/// "System: ...\nUser: ...\nAssistant: " format with proper model-specific
/// templates.
///
/// # Arguments
/// * `messages` - The chat messages to format
/// * `model_name` - Optional model name for auto-detection (defaults to Raw)
///
/// # Returns
/// Formatted prompt string ready for tokenization
///
/// # Example
///
/// ```
/// use realizar::chat_template::{ChatMessage, format_messages};
///
/// let messages = vec![
/// ChatMessage::system("You are helpful."),
/// ChatMessage::user("Hello!"),
/// ];
///
/// // With model name - uses ChatML format
/// let prompt = format_messages(&messages, Some("Qwen2-0.5B")).expect("prompt");
/// assert!(prompt.contains("<|im_start|>"));
///
/// // Without model name - uses Raw format
/// let prompt = format_messages(&messages, None).expect("prompt");
/// assert!(prompt.contains("You are helpful."));
/// ```
pub fn format_messages(
messages: &[ChatMessage],
model_name: Option<&str>,
) -> Result<String, RealizarError> {
let template = model_name.map_or_else(
|| Box::new(RawTemplate::new()) as Box<dyn ChatTemplateEngine>,
auto_detect_template,
);
template.format_conversation(messages)
}