Asep Haryana
495b9ed126
- Removed /think trigger from prompt (model generates better without it) - Added <think> and </think> back to clean_text (model uses plain text thinking tags, not special tokens)
302 lines
11 KiB
Rust
302 lines
11 KiB
Rust
//! Chat completion use cases.
|
|
//!
|
|
//! Orchestrates prompt building, sampler construction, and output parsing.
|
|
//! These are pure functions with no framework dependencies.
|
|
|
|
use llama_cpp_2::sampling::LlamaSampler;
|
|
|
|
use crate::domain::entity::{ChatMessage, ChatRequest, ToolCall, ToolCallFunction, ToolDef};
|
|
|
|
/// Build a prompt string from conversation messages and optional tool definitions.
|
|
///
|
|
/// Uses ChatML format with `<|im_start|>` / `<|im_end|>` delimiters. Tool
|
|
/// definitions are injected into the first system message.
|
|
pub fn build_prompt(messages: &[ChatMessage], tools: &Option<Vec<ToolDef>>) -> String {
|
|
let mut prompt = String::new();
|
|
|
|
for (i, msg) in messages.iter().enumerate() {
|
|
match msg.role.as_str() {
|
|
"system" => {
|
|
let mut content = msg.content.clone().unwrap_or_default();
|
|
// Inject tools into the system message (first occurrence)
|
|
if i == 0 {
|
|
if let Some(tools_list) = tools {
|
|
if !tools_list.is_empty() {
|
|
let mut tools_text = String::from(
|
|
"\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>",
|
|
);
|
|
for tool in tools_list {
|
|
tools_text.push('\n');
|
|
tools_text.push_str(
|
|
&serde_json::to_string(tool).unwrap_or_default(),
|
|
);
|
|
}
|
|
tools_text.push_str(
|
|
"\n</tools>\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>",
|
|
);
|
|
content.push_str(&tools_text);
|
|
}
|
|
}
|
|
}
|
|
prompt.push_str(&format!("<|im_start|>system\n{}<|im_end|>\n", content));
|
|
}
|
|
"user" => {
|
|
let content = msg.content.as_deref().unwrap_or("");
|
|
if msg.tool_call_id.is_some() || msg.name.is_some() {
|
|
prompt.push_str(&format!(
|
|
"<|im_start|>user\n<tool_response>\n{}\n</tool_response><|im_end|>\n",
|
|
content
|
|
));
|
|
} else {
|
|
prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content));
|
|
}
|
|
}
|
|
"assistant" => {
|
|
let content = msg.content.as_deref().unwrap_or("");
|
|
if let Some(tcs) = &msg.tool_calls {
|
|
let mut asst = format!("<|im_start|>assistant\n{}", content);
|
|
for tc in tcs {
|
|
let args: serde_json::Value =
|
|
serde_json::from_str(&tc.function.arguments).unwrap_or_default();
|
|
asst.push_str(&format!(
|
|
"<tool_call>\n<function={}>\n",
|
|
tc.function.name
|
|
));
|
|
if let Some(obj) = args.as_object() {
|
|
for (k, v) in obj {
|
|
let val = match v {
|
|
serde_json::Value::String(s) => s.clone(),
|
|
other => serde_json::to_string(other).unwrap_or_default(),
|
|
};
|
|
asst.push_str(&format!("<parameter={}>\n{}\n</parameter>\n", k, val));
|
|
}
|
|
}
|
|
asst.push_str("</function>\n</tool_call>");
|
|
}
|
|
asst.push_str("<|im_end|>\n");
|
|
prompt.push_str(&asst);
|
|
} else {
|
|
prompt.push_str(&format!(
|
|
"<|im_start|>assistant\n{}<|im_end|>\n",
|
|
content
|
|
));
|
|
}
|
|
}
|
|
_ => {
|
|
let content = msg.content.as_deref().unwrap_or("");
|
|
prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", content));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Generation prompt
|
|
prompt.push_str("<|im_start|>assistant\n");
|
|
prompt
|
|
}
|
|
|
|
/// Parameters for building a [`LlamaSampler`] chain.
|
|
pub struct SamplerParams {
|
|
pub temperature: Option<f32>,
|
|
pub top_p: Option<f32>,
|
|
pub top_k: Option<u32>,
|
|
pub min_p: Option<f32>,
|
|
pub repeat_penalty: Option<f32>,
|
|
pub frequency_penalty: Option<f32>,
|
|
pub presence_penalty: Option<f32>,
|
|
pub seed: Option<u32>,
|
|
}
|
|
|
|
impl SamplerParams {
|
|
pub fn from_request(req: &ChatRequest) -> Self {
|
|
Self {
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
repeat_penalty: req.repeat_penalty,
|
|
frequency_penalty: req.frequency_penalty,
|
|
presence_penalty: req.presence_penalty,
|
|
seed: req.seed,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Build a [`LlamaSampler`] chain from [`SamplerParams`].
|
|
pub fn build_sampler(params: &SamplerParams) -> LlamaSampler {
|
|
use llama_cpp_2::sampling::LlamaSampler as LS;
|
|
|
|
let temperature = params.temperature;
|
|
let top_p = params.top_p;
|
|
let top_k = params.top_k;
|
|
let min_p = params.min_p;
|
|
let repeat_penalty = params.repeat_penalty;
|
|
let frequency_penalty = params.frequency_penalty;
|
|
let presence_penalty = params.presence_penalty;
|
|
let seed = params.seed;
|
|
let mut samplers: Vec<LlamaSampler> = Vec::new();
|
|
|
|
// Repetition/frequency/presence penalties
|
|
let repeat = repeat_penalty.unwrap_or(1.0);
|
|
let freq = frequency_penalty.unwrap_or(0.0);
|
|
let present = presence_penalty.unwrap_or(0.0);
|
|
if (repeat - 1.0).abs() > 1e-6 || freq > 0.0 || present > 0.0 {
|
|
samplers.push(LS::penalties(64, repeat, freq, present));
|
|
}
|
|
|
|
// top_k
|
|
if let Some(k) = top_k {
|
|
samplers.push(LS::top_k(k as i32));
|
|
}
|
|
|
|
// top_p
|
|
if let Some(p) = top_p {
|
|
samplers.push(LS::top_p(p, 1));
|
|
}
|
|
|
|
// min_p
|
|
if let Some(p) = min_p {
|
|
samplers.push(LS::min_p(p, 1));
|
|
}
|
|
|
|
// Temperature + final selector
|
|
let temp = temperature.unwrap_or(0.0);
|
|
if temp <= 0.0 {
|
|
samplers.push(LS::greedy());
|
|
} else {
|
|
if (temp - 1.0).abs() > 1e-6 {
|
|
samplers.push(LS::temp(temp));
|
|
}
|
|
let s = seed.unwrap_or(0);
|
|
samplers.push(LS::dist(s));
|
|
}
|
|
|
|
LlamaSampler::chain_simple(samplers)
|
|
}
|
|
|
|
|
|
// ═══════════════════════════════════════════════════════════════
|
|
// TEXT PROCESSING
|
|
// ═══════════════════════════════════════════════════════════════
|
|
|
|
/// Remove special tokens from generated text.
|
|
pub fn clean_text(text: &str) -> String {
|
|
text.replace("<|im_end|>", "")
|
|
.replace("<|im_start|>", "")
|
|
.replace("<|thought_begin|>", "")
|
|
.replace("<|thought_end|>", "")
|
|
.replace("<|tool_call|>", "")
|
|
.replace("<|execute_start|>", "")
|
|
.replace("<|execute_end|>", "")
|
|
.replace("<think>", "")
|
|
.replace("</think>", "")
|
|
.replace("/think", "")
|
|
.replace("/no_think", "")
|
|
.trim()
|
|
.to_string()
|
|
}
|
|
|
|
/// Parse tool calls from generated text in the format:
|
|
///
|
|
/// ```xml
|
|
/// <tool_call>
|
|
/// <function=name>
|
|
/// <parameter=key>value</parameter>
|
|
/// </function>
|
|
/// </tool_call>
|
|
/// ```
|
|
pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
|
|
let mut clean = text.to_string();
|
|
let mut tool_calls: Vec<ToolCall> = Vec::new();
|
|
|
|
let mut idx = 0;
|
|
loop {
|
|
let start_tag = "<tool_call>";
|
|
let end_tag = "</tool_call>";
|
|
|
|
let start = match clean[idx..].find(start_tag) {
|
|
Some(s) => idx + s,
|
|
None => break,
|
|
};
|
|
|
|
let end = match clean[start..].find(end_tag) {
|
|
Some(e) => start + e + end_tag.len(),
|
|
None => break,
|
|
};
|
|
|
|
let block = &clean[start + start_tag.len()..end - end_tag.len()];
|
|
let trimmed = block.trim();
|
|
|
|
// Parse function name
|
|
let func_name = trimmed
|
|
.lines()
|
|
.next()
|
|
.and_then(|l| {
|
|
let l = l.trim();
|
|
l.strip_prefix("<function=")
|
|
.and_then(|s| s.strip_suffix('>'))
|
|
.map(|s| s.trim().to_string())
|
|
})
|
|
.unwrap_or_default();
|
|
|
|
// Parse parameters
|
|
let mut args_map = serde_json::Map::new();
|
|
let lines = trimmed.lines();
|
|
let mut current_param: Option<String> = None;
|
|
let mut current_value = String::new();
|
|
let mut in_param = false;
|
|
|
|
for line in lines {
|
|
let line = line.trim();
|
|
if let Some(param) =
|
|
line.strip_prefix("<parameter=")
|
|
.and_then(|s| s.strip_suffix('>'))
|
|
{
|
|
if let Some(p) = current_param.take() {
|
|
args_map.insert(
|
|
p,
|
|
serde_json::Value::String(current_value.trim().to_string()),
|
|
);
|
|
current_value = String::new();
|
|
}
|
|
current_param = Some(param.to_string());
|
|
in_param = true;
|
|
} else if line == "</parameter>" {
|
|
in_param = false;
|
|
} else if in_param {
|
|
if !current_value.is_empty() {
|
|
current_value.push('\n');
|
|
}
|
|
current_value.push_str(line);
|
|
} else if line.starts_with("<function=") || line.starts_with("</function>") {
|
|
continue;
|
|
}
|
|
}
|
|
// Save last param
|
|
if let Some(p) = current_param.take() {
|
|
args_map.insert(p, serde_json::Value::String(current_value.trim().to_string()));
|
|
}
|
|
|
|
let args_json = serde_json::Value::Object(args_map).to_string();
|
|
|
|
tool_calls.push(ToolCall {
|
|
id: format!("call_{}", uuid::Uuid::new_v4().to_string().replace('-', "")),
|
|
call_type: "function".into(),
|
|
function: ToolCallFunction {
|
|
name: func_name,
|
|
arguments: args_json,
|
|
},
|
|
});
|
|
|
|
idx = end;
|
|
}
|
|
|
|
// Remove tool_call blocks from the text
|
|
clean = clean.replace("<tool_call>", "").replace("</tool_call>", "");
|
|
// XML-like tags are already fully parsed; remaining text is the content
|
|
clean = clean.trim().to_string();
|
|
// Strip remaining XML tags that aren't part of clean
|
|
let cleaned = clean_text(&clean);
|
|
|
|
(cleaned, tool_calls)
|
|
}
|