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()),
}],