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
This commit is contained in:
Generated
+17
@@ -568,6 +568,7 @@ dependencies = [
|
|||||||
"chrono",
|
"chrono",
|
||||||
"futures",
|
"futures",
|
||||||
"llama-cpp-2",
|
"llama-cpp-2",
|
||||||
|
"minijinja",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"thiserror 1.0.69",
|
"thiserror 1.0.69",
|
||||||
@@ -615,12 +616,28 @@ version = "2.8.3"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
|
checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "memo-map"
|
||||||
|
version = "0.3.3"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "mime"
|
name = "mime"
|
||||||
version = "0.3.17"
|
version = "0.3.17"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
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]]
|
[[package]]
|
||||||
name = "minimal-lexical"
|
name = "minimal-lexical"
|
||||||
version = "0.2.1"
|
version = "0.2.1"
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ edition = "2021"
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
# LLM inference
|
# LLM inference
|
||||||
llama-cpp-2 = "0.1"
|
llama-cpp-2 = "0.1"
|
||||||
|
minijinja = "2"
|
||||||
|
|
||||||
# HTTP server
|
# HTTP server
|
||||||
axum = { version = "0.8", features = ["json"] }
|
axum = { version = "0.8", features = ["json"] }
|
||||||
|
|||||||
@@ -0,0 +1,53 @@
|
|||||||
|
{{- bos_token }}
|
||||||
|
{%- if tools %}
|
||||||
|
{%- set tool_definitions %}
|
||||||
|
{{- "# Tools\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||||
|
{%- for tool in tools %}
|
||||||
|
{{- "\n" }}
|
||||||
|
{{- tool | tojson }}
|
||||||
|
{%- endfor %}
|
||||||
|
{{- '\n</tools>\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: <function name="name"><param name="key">value</param></function>' }}
|
||||||
|
{%- 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<tool_call>\n<function=" + tool_call.function.name + ">\n" + (tool_call.function.arguments | tojson) + "\n</function>\n</tool_call>" %}
|
||||||
|
{%- 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<tool_response>\n' + content + '\n</tool_response><|im_end|>\n' }}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endfor %}
|
||||||
|
{%- if add_generation_prompt %}
|
||||||
|
{{- '<|im_start|>assistant\n' }}
|
||||||
|
{%- if enable_thinking is defined %}
|
||||||
|
{%- if enable_thinking is false %}
|
||||||
|
{{- '<think>\n\n</think>\n\n' }}
|
||||||
|
{%- elif enable_thinking is true %}
|
||||||
|
{{- '<think>\n' }}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endif %}
|
||||||
@@ -1,70 +1,93 @@
|
|||||||
//! Chat completion use cases.
|
//! 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 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
|
/// Build a prompt from messages using the GGUF's Jinja chat template.
|
||||||
/// chat template. The template handles system/user/assistant/tool messages,
|
///
|
||||||
/// thinking mode, and tool definitions automatically.
|
/// Renders the model's baked-in template via minijinja, passing the message
|
||||||
|
/// history, optional tool definitions, and generation-prompt switches.
|
||||||
pub fn build_prompt(
|
pub fn build_prompt(
|
||||||
model: &llama_cpp_2::model::LlamaModel,
|
model: &llama_cpp_2::model::LlamaModel,
|
||||||
messages: &[crate::domain::entity::ChatMessage],
|
messages: &[ChatMessage],
|
||||||
_tools: &Option<Vec<crate::domain::entity::ToolDef>>,
|
_tools: &Option<Vec<crate::domain::entity::ToolDef>>,
|
||||||
) -> Result<String, String> {
|
) -> Result<String, String> {
|
||||||
let tmpl = model
|
// Load the GGUF's chat template (embedded at crate build time)
|
||||||
.chat_template(None)
|
let template_str = include_str!("templates/chat_template.jinja");
|
||||||
.map_err(|e| format!("Chat template error: {e}"))?;
|
|
||||||
|
|
||||||
let mut chat_msgs: Vec<LlamaChatMessage> = 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<Value> = Vec::new();
|
||||||
for msg in messages {
|
for msg in messages {
|
||||||
let content = msg.content.clone().unwrap_or_default();
|
let mut m: HashMap<String, Value> = HashMap::new();
|
||||||
let role = msg.role.clone();
|
m.insert("role".into(), Value::from(msg.role.clone()));
|
||||||
|
|
||||||
// Build content with tool calls for assistant messages
|
let content = msg.content.clone().unwrap_or_default();
|
||||||
let full_content = if role == "assistant" {
|
|
||||||
|
// For assistant messages, check if there are tool_calls
|
||||||
|
if msg.role == "assistant" {
|
||||||
if let Some(tcs) = &msg.tool_calls {
|
if let Some(tcs) = &msg.tool_calls {
|
||||||
let mut c = content;
|
// Serialise tool calls per the template's expected format
|
||||||
for tc in tcs {
|
let tcs_val: Vec<Value> = tcs.iter().map(|tc| {
|
||||||
let args: serde_json::Value =
|
let args: serde_json::Value =
|
||||||
serde_json::from_str(&tc.function.arguments).unwrap_or_default();
|
serde_json::from_str(&tc.function.arguments).unwrap_or_default();
|
||||||
let args_str = serde_json::to_string(&args).unwrap_or_default();
|
Value::from_serialize(&serde_json::json!({
|
||||||
c.push_str(&format!(
|
"id": tc.id,
|
||||||
"<tool_call>\n<function={}>\n{}\n</function>\n</tool_call>",
|
"type": "function",
|
||||||
tc.function.name, args_str
|
"function": {
|
||||||
));
|
"name": tc.function.name,
|
||||||
}
|
"arguments": args,
|
||||||
c
|
}
|
||||||
} else {
|
}))
|
||||||
content
|
}).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!("<tool_response>\n{}\n</tool_response>", content);
|
||||||
|
m.insert("content".into(), Value::from(wrapped));
|
||||||
} else {
|
} else {
|
||||||
content
|
m.insert("content".into(), Value::from(content));
|
||||||
};
|
}
|
||||||
|
|
||||||
let llama_role = match role.as_str() {
|
msgs_val.push(Value::from(m));
|
||||||
"tool" => "tool".to_string(),
|
|
||||||
r => r.to_string(),
|
|
||||||
};
|
|
||||||
|
|
||||||
chat_msgs.push(
|
|
||||||
LlamaChatMessage::new(llama_role, full_content)
|
|
||||||
.map_err(|e| format!("Message error: {e}"))?,
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Apply chat template with generation prompt (add_ass = true)
|
// BOS token for sentencepiece / unigram models
|
||||||
let mut result = model
|
let bos_token: &str = "<s>";
|
||||||
.apply_chat_template(&tmpl, &chat_msgs, true)
|
|
||||||
.map_err(|e| format!("Template error: {e}"))?;
|
|
||||||
|
|
||||||
// Append think trigger for MiniCPM5 thinking mode:
|
// Build context
|
||||||
// <|im_start|>assistant\n<think>\n → model generates reasoning + </think> + answer
|
let mut ctx: HashMap<String, Value> = HashMap::new();
|
||||||
result.push_str("<think>\n");
|
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)
|
Ok(result)
|
||||||
}
|
}
|
||||||
@@ -147,14 +170,32 @@ pub fn build_sampler(params: &SamplerParams) -> LlamaSampler {
|
|||||||
// TEXT PROCESSING
|
// TEXT PROCESSING
|
||||||
// ═══════════════════════════════════════════════════════════════
|
// ═══════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
/// Remove special tokens from generated text.
|
/// Remove special tokens from generated text and separate reasoning.
|
||||||
pub fn clean_text(text: &str) -> String {
|
///
|
||||||
text.replace("<|im_end|>", "")
|
/// For thinking models, returns (reasoning, cleaned_answer).
|
||||||
.replace("<|im_start|>", "")
|
pub fn clean_text(text: &str) -> (String, String) {
|
||||||
.replace("<think>", "")
|
let text = text.replace("<|im_end|>", "")
|
||||||
.replace("</think>", "")
|
.replace("<|im_start|>", "");
|
||||||
.trim()
|
|
||||||
.to_string()
|
// Separate reasoning (between <think>/</think>) from answer
|
||||||
|
let text = text.trim();
|
||||||
|
let (reasoning, answer) = if let Some(close_idx) = text.find("</think>") {
|
||||||
|
let reasoning = text[..close_idx].trim()
|
||||||
|
.trim_start_matches("<think>")
|
||||||
|
.trim()
|
||||||
|
.to_string();
|
||||||
|
let answer = text[close_idx + 8..].trim().to_string();
|
||||||
|
(reasoning, answer)
|
||||||
|
} else if text.contains("<think>") {
|
||||||
|
// Still thinking — everything is reasoning
|
||||||
|
let reasoning = text.trim_start_matches("<think>").trim().to_string();
|
||||||
|
(reasoning, String::new())
|
||||||
|
} else {
|
||||||
|
(String::new(), text.to_string())
|
||||||
|
};
|
||||||
|
|
||||||
|
let answer = answer.replace("<think>", "").replace("</think>", "");
|
||||||
|
(reasoning, answer.trim().to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Parse tool calls from generated text in the format:
|
/// Parse tool calls from generated text in the format:
|
||||||
@@ -253,7 +294,7 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
|
|||||||
|
|
||||||
clean = clean.replace("<tool_call>", "").replace("</tool_call>", "");
|
clean = clean.replace("<tool_call>", "").replace("</tool_call>", "");
|
||||||
clean = clean.trim().to_string();
|
clean = clean.trim().to_string();
|
||||||
let cleaned = clean_text(&clean);
|
let (_reasoning, cleaned) = clean_text(&clean);
|
||||||
|
|
||||||
(cleaned, tool_calls)
|
(cleaned, tool_calls)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -96,8 +96,11 @@ pub struct Choice {
|
|||||||
#[derive(Serialize)]
|
#[derive(Serialize)]
|
||||||
pub struct ResponseMessage {
|
pub struct ResponseMessage {
|
||||||
pub role: String,
|
pub role: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub content: Option<String>,
|
pub content: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub reasoning_content: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub tool_calls: Option<Vec<ToolCall>>,
|
pub tool_calls: Option<Vec<ToolCall>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -165,6 +168,8 @@ pub struct SseDelta {
|
|||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub content: Option<String>,
|
pub content: Option<String>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub reasoning_content: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub tool_calls: Option<Vec<ToolCall>>,
|
pub tool_calls: Option<Vec<ToolCall>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -72,7 +72,8 @@ async fn handle_non_streaming(
|
|||||||
.generate(&input_tokens, &mut sampler, max_tokens, &stop)
|
.generate(&input_tokens, &mut sampler, max_tokens, &stop)
|
||||||
.await?;
|
.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;
|
let completion_tokens = output_tokens.len() as u32;
|
||||||
info!(" {} generated tokens", completion_tokens);
|
info!(" {} generated tokens", completion_tokens);
|
||||||
@@ -85,6 +86,8 @@ async fn handle_non_streaming(
|
|||||||
"length"
|
"length"
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let reasoning_opt = if reasoning.is_empty() { None } else { Some(reasoning) };
|
||||||
|
|
||||||
Ok(Json(ChatResponse {
|
Ok(Json(ChatResponse {
|
||||||
id: chat_id,
|
id: chat_id,
|
||||||
object: "chat.completion".into(),
|
object: "chat.completion".into(),
|
||||||
@@ -95,6 +98,7 @@ async fn handle_non_streaming(
|
|||||||
message: ResponseMessage {
|
message: ResponseMessage {
|
||||||
role: "assistant".into(),
|
role: "assistant".into(),
|
||||||
content: Some(output_text),
|
content: Some(output_text),
|
||||||
|
reasoning_content: reasoning_opt,
|
||||||
tool_calls: if tool_calls.is_empty() {
|
tool_calls: if tool_calls.is_empty() {
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
@@ -141,6 +145,7 @@ async fn handle_streaming(
|
|||||||
role: Some("assistant".into()),
|
role: Some("assistant".into()),
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: None,
|
finish_reason: None,
|
||||||
}],
|
}],
|
||||||
@@ -178,6 +183,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: Some("length".into()),
|
finish_reason: Some("length".into()),
|
||||||
}],
|
}],
|
||||||
@@ -202,6 +208,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: Some("stop".into()),
|
finish_reason: Some("stop".into()),
|
||||||
}],
|
}],
|
||||||
@@ -237,6 +244,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: Some(reason.into()),
|
finish_reason: Some(reason.into()),
|
||||||
}],
|
}],
|
||||||
@@ -247,7 +255,7 @@ async fn handle_streaming(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let piece = state.engine.decode_token(current);
|
let piece = state.engine.decode_token(current);
|
||||||
let content = chat::clean_text(&piece);
|
let (_reasoning, content) = chat::clean_text(&piece);
|
||||||
|
|
||||||
if !content.is_empty() {
|
if !content.is_empty() {
|
||||||
let chunk = serde_json::to_string(&SseChunk {
|
let chunk = serde_json::to_string(&SseChunk {
|
||||||
@@ -261,6 +269,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: Some(content.clone()),
|
content: Some(content.clone()),
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: None,
|
finish_reason: None,
|
||||||
}],
|
}],
|
||||||
@@ -293,6 +302,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: Some("stop".into()),
|
finish_reason: Some("stop".into()),
|
||||||
}],
|
}],
|
||||||
@@ -318,6 +328,7 @@ async fn handle_streaming(
|
|||||||
role: None,
|
role: None,
|
||||||
content: None,
|
content: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
reasoning_content: None,
|
||||||
},
|
},
|
||||||
finish_reason: Some("tool_calls".into()),
|
finish_reason: Some("tool_calls".into()),
|
||||||
}],
|
}],
|
||||||
|
|||||||
Reference in New Issue
Block a user