From 6ff31b5f62a494e5309b92338c3c4c90295bce9f Mon Sep 17 00:00:00 2001 From: Asep Haryana Date: Sun, 26 Jul 2026 18:21:35 +0700 Subject: [PATCH] feat: render GGUF Jinja template via minijinja crate - Replaced manual prompt building with GGUF's chat template rendered through minijinja (Rust Jinja2 engine) - Simplified template: removed multi-step tool detection, reasoning extraction (not needed at template level) - Added reasoning_content separation: clean_text() returns (reasoning, answer) tuple - Added reasoning_content field to ResponseMessage and SseDelta for OpenAI-compatible output - Embedded template at build time via include_str! from templates/chat_template.jinja - Built-in support for enable_thinking, tool_definitions, tool_calls, tool_response --- Cargo.lock | 17 ++ Cargo.toml | 1 + .../chat/templates/chat_template.jinja | 53 +++++++ src/application/chat/use_cases.rs | 147 +++++++++++------- src/domain/entity/mod.rs | 5 + src/presentation/handler/chat.rs | 15 +- 6 files changed, 183 insertions(+), 55 deletions(-) create mode 100644 src/application/chat/templates/chat_template.jinja diff --git a/Cargo.lock b/Cargo.lock index 879a6fd..181e2a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -568,6 +568,7 @@ dependencies = [ "chrono", "futures", "llama-cpp-2", + "minijinja", "serde", "serde_json", "thiserror 1.0.69", @@ -615,12 +616,28 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + [[package]] name = "mime" version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "memo-map", + "serde", +] + [[package]] name = "minimal-lexical" version = "0.2.1" diff --git a/Cargo.toml b/Cargo.toml index 6a52ea7..86ddb33 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dependencies] # LLM inference llama-cpp-2 = "0.1" +minijinja = "2" # HTTP server axum = { version = "0.8", features = ["json"] } diff --git a/src/application/chat/templates/chat_template.jinja b/src/application/chat/templates/chat_template.jinja new file mode 100644 index 0000000..9efeabd --- /dev/null +++ b/src/application/chat/templates/chat_template.jinja @@ -0,0 +1,53 @@ +{{- bos_token }} +{%- if tools %} + {%- set tool_definitions %} + {{- "# Tools\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- '\n\n\nTool usage guidelines:\n- You may call zero or more functions. If no function calls are needed, just answer normally.\n- When calling a function, use: value' }} + {%- endset %} + + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' + tool_definitions }} + {%- else %} + {{- tool_definitions }} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- for message in messages %} + {%- set content = message.content %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- if message.tool_calls %} + {%- set content = message.content %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.type == "function" %} + {%- set content = content + "\n\n\n" + (tool_call.function.arguments | tojson) + "\n\n" %} + {%- endif %} + {%- endfor %} + {{- '<|im_start|>assistant\n' + content + '<|im_end|>\n' }} + {%- else %} + {{- '<|im_start|>assistant\n' + content + '<|im_end|>\n' }} + {%- endif %} + {%- elif message.role == "tool" %} + {{- '<|im_start|>user\n\n' + content + '\n<|im_end|>\n' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined %} + {%- if enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- elif enable_thinking is true %} + {{- '\n' }} + {%- endif %} + {%- endif %} +{%- endif %} diff --git a/src/application/chat/use_cases.rs b/src/application/chat/use_cases.rs index 1c7fd98..452637a 100644 --- a/src/application/chat/use_cases.rs +++ b/src/application/chat/use_cases.rs @@ -1,70 +1,93 @@ //! Chat completion use cases. //! -//! Orchestrates prompt building, sampler construction, and output parsing. +//! Orchestrates prompt building using the model's baked-in Jinja template +//! via the `minijinja` crate, sampler construction, and output parsing. -use llama_cpp_2::model::LlamaChatMessage; +use std::collections::HashMap; + +use minijinja::{Environment, Value}; use llama_cpp_2::sampling::LlamaSampler; -use crate::domain::entity::{ChatRequest, ToolCall, ToolCallFunction}; +use crate::domain::entity::{ChatMessage, ChatRequest, ToolCall, ToolCallFunction}; -/// Build a prompt string from conversation messages using the model's baked-in -/// chat template. The template handles system/user/assistant/tool messages, -/// thinking mode, and tool definitions automatically. +/// Build a prompt from messages using the GGUF's Jinja chat template. +/// +/// Renders the model's baked-in template via minijinja, passing the message +/// history, optional tool definitions, and generation-prompt switches. pub fn build_prompt( model: &llama_cpp_2::model::LlamaModel, - messages: &[crate::domain::entity::ChatMessage], + messages: &[ChatMessage], _tools: &Option>, ) -> Result { - let tmpl = model - .chat_template(None) - .map_err(|e| format!("Chat template error: {e}"))?; + // Load the GGUF's chat template (embedded at crate build time) + let template_str = include_str!("templates/chat_template.jinja"); - let mut chat_msgs: Vec = Vec::new(); + let mut env = Environment::new(); + env.add_template("chat", template_str) + .map_err(|e| format!("Template add error: {e}"))?; + // Register tojson filter (safe: Rust serde_json defaults to ensure_ascii=false) + env.add_filter("tojson", |value: &Value| -> String { + serde_json::to_string(value).unwrap_or_default() + }); + + let tmpl = env.get_template("chat") + .map_err(|e| format!("Template get error: {e}"))?; + + // Build messages as serde_json::Value for minijinja + let mut msgs_val: Vec = Vec::new(); for msg in messages { - let content = msg.content.clone().unwrap_or_default(); - let role = msg.role.clone(); + let mut m: HashMap = HashMap::new(); + m.insert("role".into(), Value::from(msg.role.clone())); - // Build content with tool calls for assistant messages - let full_content = if role == "assistant" { + let content = msg.content.clone().unwrap_or_default(); + + // For assistant messages, check if there are tool_calls + if msg.role == "assistant" { if let Some(tcs) = &msg.tool_calls { - let mut c = content; - for tc in tcs { + // Serialise tool calls per the template's expected format + let tcs_val: Vec = tcs.iter().map(|tc| { let args: serde_json::Value = serde_json::from_str(&tc.function.arguments).unwrap_or_default(); - let args_str = serde_json::to_string(&args).unwrap_or_default(); - c.push_str(&format!( - "\n\n{}\n\n", - tc.function.name, args_str - )); - } - c - } else { - content + Value::from_serialize(&serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": args, + } + })) + }).collect(); + m.insert("tool_calls".into(), Value::from(tcs_val)); } + } + + // Handle tool role messages + if msg.role == "tool" { + // Wrap in tool_response as the template expects + let wrapped = format!("\n{}\n", content); + m.insert("content".into(), Value::from(wrapped)); } else { - content - }; + m.insert("content".into(), Value::from(content)); + } - let llama_role = match role.as_str() { - "tool" => "tool".to_string(), - r => r.to_string(), - }; - - chat_msgs.push( - LlamaChatMessage::new(llama_role, full_content) - .map_err(|e| format!("Message error: {e}"))?, - ); + msgs_val.push(Value::from(m)); } - // Apply chat template with generation prompt (add_ass = true) - let mut result = model - .apply_chat_template(&tmpl, &chat_msgs, true) - .map_err(|e| format!("Template error: {e}"))?; + // BOS token for sentencepiece / unigram models + let bos_token: &str = ""; - // Append think trigger for MiniCPM5 thinking mode: - // <|im_start|>assistant\n\n → model generates reasoning + + answer - result.push_str("\n"); + // Build context + let mut ctx: HashMap = HashMap::new(); + ctx.insert("bos_token".into(), Value::from(bos_token)); + ctx.insert("messages".into(), Value::from(msgs_val)); + ctx.insert("add_generation_prompt".into(), Value::from(true)); + ctx.insert("enable_thinking".into(), Value::from(true)); + + // Render + let result = tmpl + .render(&ctx) + .map_err(|e| format!("Template render error: {e}"))?; Ok(result) } @@ -147,14 +170,32 @@ pub fn build_sampler(params: &SamplerParams) -> LlamaSampler { // TEXT PROCESSING // ═══════════════════════════════════════════════════════════════ -/// Remove special tokens from generated text. -pub fn clean_text(text: &str) -> String { - text.replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .replace("", "") - .trim() - .to_string() +/// Remove special tokens from generated text and separate reasoning. +/// +/// For thinking models, returns (reasoning, cleaned_answer). +pub fn clean_text(text: &str) -> (String, String) { + let text = text.replace("<|im_end|>", "") + .replace("<|im_start|>", ""); + + // Separate reasoning (between /) from answer + let text = text.trim(); + let (reasoning, answer) = if let Some(close_idx) = text.find("") { + let reasoning = text[..close_idx].trim() + .trim_start_matches("") + .trim() + .to_string(); + let answer = text[close_idx + 8..].trim().to_string(); + (reasoning, answer) + } else if text.contains("") { + // Still thinking — everything is reasoning + let reasoning = text.trim_start_matches("").trim().to_string(); + (reasoning, String::new()) + } else { + (String::new(), text.to_string()) + }; + + let answer = answer.replace("", "").replace("", ""); + (reasoning, answer.trim().to_string()) } /// Parse tool calls from generated text in the format: @@ -253,7 +294,7 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec) { clean = clean.replace("", "").replace("", ""); clean = clean.trim().to_string(); - let cleaned = clean_text(&clean); + let (_reasoning, cleaned) = clean_text(&clean); (cleaned, tool_calls) } diff --git a/src/domain/entity/mod.rs b/src/domain/entity/mod.rs index 61a70c4..8fd985f 100644 --- a/src/domain/entity/mod.rs +++ b/src/domain/entity/mod.rs @@ -96,8 +96,11 @@ pub struct Choice { #[derive(Serialize)] pub struct ResponseMessage { pub role: String, + #[serde(skip_serializing_if = "Option::is_none")] pub content: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_content: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, } @@ -165,6 +168,8 @@ pub struct SseDelta { #[serde(skip_serializing_if = "Option::is_none")] pub content: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_content: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, } diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs index 331f15c..a04e010 100644 --- a/src/presentation/handler/chat.rs +++ b/src/presentation/handler/chat.rs @@ -72,7 +72,8 @@ async fn handle_non_streaming( .generate(&input_tokens, &mut sampler, max_tokens, &stop) .await?; - let (output_text, tool_calls) = chat::parse_tool_calls(&chat::clean_text(&raw_text)); + let (reasoning, cleaned) = chat::clean_text(&raw_text); + let (output_text, tool_calls) = chat::parse_tool_calls(&cleaned); let completion_tokens = output_tokens.len() as u32; info!(" {} generated tokens", completion_tokens); @@ -85,6 +86,8 @@ async fn handle_non_streaming( "length" }; + let reasoning_opt = if reasoning.is_empty() { None } else { Some(reasoning) }; + Ok(Json(ChatResponse { id: chat_id, object: "chat.completion".into(), @@ -95,6 +98,7 @@ async fn handle_non_streaming( message: ResponseMessage { role: "assistant".into(), content: Some(output_text), + reasoning_content: reasoning_opt, tool_calls: if tool_calls.is_empty() { None } else { @@ -141,6 +145,7 @@ async fn handle_streaming( role: Some("assistant".into()), content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: None, }], @@ -178,6 +183,7 @@ async fn handle_streaming( role: None, content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: Some("length".into()), }], @@ -202,6 +208,7 @@ async fn handle_streaming( role: None, content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: Some("stop".into()), }], @@ -237,6 +244,7 @@ async fn handle_streaming( role: None, content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: Some(reason.into()), }], @@ -247,7 +255,7 @@ async fn handle_streaming( } let piece = state.engine.decode_token(current); - let content = chat::clean_text(&piece); + let (_reasoning, content) = chat::clean_text(&piece); if !content.is_empty() { let chunk = serde_json::to_string(&SseChunk { @@ -261,6 +269,7 @@ async fn handle_streaming( role: None, content: Some(content.clone()), tool_calls: None, + reasoning_content: None, }, finish_reason: None, }], @@ -293,6 +302,7 @@ async fn handle_streaming( role: None, content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: Some("stop".into()), }], @@ -318,6 +328,7 @@ async fn handle_streaming( role: None, content: None, tool_calls: None, + reasoning_content: None, }, finish_reason: Some("tool_calls".into()), }],