diff --git a/crates/inference-server/src/compat.rs b/crates/inference-server/src/compat.rs index 01ceef9..b5201ca 100644 --- a/crates/inference-server/src/compat.rs +++ b/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, + 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, 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, 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("<>\n"); + block.push_str(&pending_system.join("\n\n")); + block.push_str("\n<>\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] <>\n"); + prompt.push_str(&pending_system.join("\n\n")); + prompt.push_str("\n<> [/INST]\n"); + } +} + fn append_assistant_prefix(prompt: &mut String, template: ChatTemplate) { match template { ChatTemplate::Generic => prompt.push_str("Assistant:"),