feat: harden llama chat template rendering

crates/inference-server/src/compat.rs

@@ -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(
+        messages.push(PromptMessage {
-            &mut prompt,
+            role: Role::System,
-            template,
+            name: None,
-            Role::System,
+            content: config.default_system_prompt.clone(),
-            None,
+        });
-            &config.default_system_prompt,
-        );
     }
 
     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(']');
         }
 
+        messages.push(PromptMessage {
+            role: msg.role.clone(),
+            name: msg.name.clone(),
+            content,
+        });
+    }
+
+    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,
-            msg.role.clone(),
+                &message.role,
-            msg.name.as_deref(),
+                message.name.as_deref(),
-            &content,
+                &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(
+fn append_response_format_instruction(messages: &mut Vec<PromptMessage>, format: &ResponseFormat) {
-    prompt: &mut String,
-    template: ChatTemplate,
-    format: &ResponseFormat,
-) {
     match format.format_type.as_str() {
         "json_object" => {
-            append_chat_message(
+            messages.push(PromptMessage {
-                prompt,
+                role: Role::System,
-                template,
+                name: None,
-                Role::System,
+                content: "Return valid JSON only. Do not wrap it in markdown fences.".to_owned(),
-                None,
+            });
-                "Return valid JSON only. Do not wrap it in markdown fences.",
-            );
         }
         "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,
+    messages: &mut Vec<PromptMessage>,
-    template: ChatTemplate,
     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(
+            messages.push(PromptMessage {
-                prompt,
+                role: Role::System,
-                template,
+                name: None,
-                Role::System,
+                content: "You must call a tool before replying.".to_owned(),
-                None,
+            });
-                "You must call a tool before replying.",
-            );
         }
         Some(ToolChoice::Mode(mode)) if mode == "none" => {
-            append_chat_message(
+            messages.push(PromptMessage {
-                prompt,
+                role: Role::System,
-                template,
+                name: None,
-                Role::System,
+                content: "Do not call tools in this response.".to_owned(),
-                None,
+            });
-                "Do not call tools in this response.",
-            );
         }
         Some(ToolChoice::Specific(choice)) => {
-            let content = format!("You must call the tool named {}.", choice.function.name);
+            messages.push(PromptMessage {
-            append_chat_message(prompt, template, Role::System, None, &content);
+                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:"),