feat: harden llama chat template rendering
This commit is contained in:
Evandro Fonseca Junior committed
1 parent
f96b42c356
commit
28d3fa7b51
1 file changed
+113
-52
@@ -191,6 +191,13 @@ pub struct FinalizedAssistant {
|
||||
pub finish_reason: FinishReason,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct PromptMessage {
|
||||
role: Role,
|
||||
name: Option<String>,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ChatTemplate {
|
||||
Generic,
|
||||
@@ -304,28 +311,26 @@ pub fn validate_request(
|
||||
|
||||
pub fn prepare_prompt(req: &ChatCompletionRequest, config: &ServerConfig) -> PreparedPrompt {
|
||||
let template = select_chat_template(req, config);
|
||||
let mut prompt = String::with_capacity(1024);
|
||||
let mut messages = Vec::with_capacity(req.messages.len() + 4);
|
||||
let has_system = req
|
||||
.messages
|
||||
.iter()
|
||||
.any(|message| matches!(message.role, Role::System | Role::Developer));
|
||||
|
||||
if !has_system {
|
||||
append_chat_message(
|
||||
&mut prompt,
|
||||
template,
|
||||
Role::System,
|
||||
None,
|
||||
&config.default_system_prompt,
|
||||
);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: config.default_system_prompt.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(format) = &req.response_format {
|
||||
append_response_format_instruction(&mut prompt, template, format);
|
||||
append_response_format_instruction(&mut messages, format);
|
||||
}
|
||||
|
||||
if !req.tools.is_empty() {
|
||||
append_tool_instructions(&mut prompt, template, &req.tools, req.tool_choice.as_ref());
|
||||
append_tool_instructions(&mut messages, &req.tools, req.tool_choice.as_ref());
|
||||
}
|
||||
|
||||
for msg in &req.messages {
|
||||
@@ -341,16 +346,29 @@ pub fn prepare_prompt(req: &ChatCompletionRequest, config: &ServerConfig) -> Pre
|
||||
content.push(']');
|
||||
}
|
||||
|
||||
append_chat_message(
|
||||
&mut prompt,
|
||||
template,
|
||||
msg.role.clone(),
|
||||
msg.name.as_deref(),
|
||||
&content,
|
||||
);
|
||||
messages.push(PromptMessage {
|
||||
role: msg.role.clone(),
|
||||
name: msg.name.clone(),
|
||||
content,
|
||||
});
|
||||
}
|
||||
|
||||
append_assistant_prefix(&mut prompt, template);
|
||||
let mut prompt = String::with_capacity(1024);
|
||||
if template == ChatTemplate::Llama {
|
||||
render_llama_prompt(&mut prompt, &messages);
|
||||
} else {
|
||||
for message in &messages {
|
||||
append_chat_message(
|
||||
&mut prompt,
|
||||
template,
|
||||
&message.role,
|
||||
message.name.as_deref(),
|
||||
&message.content,
|
||||
);
|
||||
}
|
||||
|
||||
append_assistant_prefix(&mut prompt, template);
|
||||
}
|
||||
|
||||
PreparedPrompt {
|
||||
prompt,
|
||||
@@ -419,35 +437,32 @@ fn select_chat_template(req: &ChatCompletionRequest, config: &ServerConfig) -> C
|
||||
}
|
||||
}
|
||||
|
||||
fn append_response_format_instruction(
|
||||
prompt: &mut String,
|
||||
template: ChatTemplate,
|
||||
format: &ResponseFormat,
|
||||
) {
|
||||
fn append_response_format_instruction(messages: &mut Vec<PromptMessage>, format: &ResponseFormat) {
|
||||
match format.format_type.as_str() {
|
||||
"json_object" => {
|
||||
append_chat_message(
|
||||
prompt,
|
||||
template,
|
||||
Role::System,
|
||||
None,
|
||||
"Return valid JSON only. Do not wrap it in markdown fences.",
|
||||
);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: "Return valid JSON only. Do not wrap it in markdown fences.".to_owned(),
|
||||
});
|
||||
}
|
||||
"json_schema" => {
|
||||
let mut content = "Return valid JSON only and follow this schema exactly: ".to_owned();
|
||||
if let Some(schema) = &format.json_schema {
|
||||
content.push_str(&schema.to_string());
|
||||
}
|
||||
append_chat_message(prompt, template, Role::System, None, &content);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content,
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn append_tool_instructions(
|
||||
prompt: &mut String,
|
||||
template: ChatTemplate,
|
||||
messages: &mut Vec<PromptMessage>,
|
||||
tools: &[ToolDefinition],
|
||||
tool_choice: Option<&ToolChoice>,
|
||||
) {
|
||||
@@ -465,37 +480,44 @@ fn append_tool_instructions(
|
||||
let mut available_tools = "Available tools: ".to_owned();
|
||||
available_tools
|
||||
.push_str(&serde_json::to_string(&tool_summaries).unwrap_or_else(|_| "[]".to_owned()));
|
||||
append_chat_message(prompt, template, Role::System, None, &available_tools);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: available_tools,
|
||||
});
|
||||
|
||||
let tool_usage = format!(
|
||||
"If you need a tool, respond with JSON only using one of these shapes: {} or {}. If no tool is needed, answer normally.",
|
||||
r#"{"tool_calls":[{"name":"tool_name","arguments":{}}]}"#,
|
||||
r#"{"name":"tool_name","arguments":{}}"#
|
||||
);
|
||||
append_chat_message(prompt, template, Role::System, None, &tool_usage);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: tool_usage,
|
||||
});
|
||||
|
||||
match tool_choice {
|
||||
Some(ToolChoice::Mode(mode)) if mode == "required" => {
|
||||
append_chat_message(
|
||||
prompt,
|
||||
template,
|
||||
Role::System,
|
||||
None,
|
||||
"You must call a tool before replying.",
|
||||
);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: "You must call a tool before replying.".to_owned(),
|
||||
});
|
||||
}
|
||||
Some(ToolChoice::Mode(mode)) if mode == "none" => {
|
||||
append_chat_message(
|
||||
prompt,
|
||||
template,
|
||||
Role::System,
|
||||
None,
|
||||
"Do not call tools in this response.",
|
||||
);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: "Do not call tools in this response.".to_owned(),
|
||||
});
|
||||
}
|
||||
Some(ToolChoice::Specific(choice)) => {
|
||||
let content = format!("You must call the tool named {}.", choice.function.name);
|
||||
append_chat_message(prompt, template, Role::System, None, &content);
|
||||
messages.push(PromptMessage {
|
||||
role: Role::System,
|
||||
name: None,
|
||||
content: format!("You must call the tool named {}.", choice.function.name),
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
@@ -504,7 +526,7 @@ fn append_tool_instructions(
|
||||
fn append_chat_message(
|
||||
prompt: &mut String,
|
||||
template: ChatTemplate,
|
||||
role: Role,
|
||||
role: &Role,
|
||||
name: Option<&str>,
|
||||
content: &str,
|
||||
) {
|
||||
@@ -561,6 +583,45 @@ fn append_chat_message(
|
||||
}
|
||||
}
|
||||
|
||||
fn render_llama_prompt(prompt: &mut String, messages: &[PromptMessage]) {
|
||||
let mut pending_system = Vec::new();
|
||||
|
||||
for message in messages {
|
||||
let content = if let Some(name) = message.name.as_deref() {
|
||||
format!("{} ({})", message.content, name)
|
||||
} else {
|
||||
message.content.clone()
|
||||
};
|
||||
|
||||
match message.role {
|
||||
Role::System | Role::Developer => pending_system.push(content),
|
||||
Role::User | Role::Tool => {
|
||||
let mut block = String::new();
|
||||
if !pending_system.is_empty() {
|
||||
block.push_str("<<SYS>>\n");
|
||||
block.push_str(&pending_system.join("\n\n"));
|
||||
block.push_str("\n<</SYS>>\n\n");
|
||||
pending_system.clear();
|
||||
}
|
||||
block.push_str(&content);
|
||||
prompt.push_str("[INST] ");
|
||||
prompt.push_str(&block);
|
||||
prompt.push_str(" [/INST]\n");
|
||||
}
|
||||
Role::Assistant => {
|
||||
prompt.push_str(&content);
|
||||
prompt.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !pending_system.is_empty() {
|
||||
prompt.push_str("[INST] <<SYS>>\n");
|
||||
prompt.push_str(&pending_system.join("\n\n"));
|
||||
prompt.push_str("\n<</SYS>> [/INST]\n");
|
||||
}
|
||||
}
|
||||
|
||||
fn append_assistant_prefix(prompt: &mut String, template: ChatTemplate) {
|
||||
match template {
|
||||
ChatTemplate::Generic => prompt.push_str("Assistant:"),
|
||||
|
||||
Reference in new issue
Block a user