diff --git a/README.md b/README.md index eb8c87c..d8a38ff 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ OpenAI-compatible LLM inference server using `llama-cpp-2` (Rust). -**Model:** MiniCPM-V-4.6 Q4_K_M (505 MB) +**Model:** MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking (Q8_0) **Engine:** llama.cpp via `llama-cpp-2` crate **Domain:** [ai.asepharyana.my.id](https://ai.asepharyana.my.id) @@ -10,7 +10,7 @@ OpenAI-compatible LLM inference server using `llama-cpp-2` (Rust). ### `GET /health` ```json -{"status": "ok", "model": "minicpm-v-4.6-q4_k_m"} +{"status": "ok", "model": "minicpm5-1b-fable5-v2-thinking"} ``` ### `GET /v1/models` @@ -19,11 +19,13 @@ OpenAI-compatible model listing. ### `POST /v1/chat/completions` OpenAI-compatible chat completions. +The server serves a single model and rejects unknown model ids with `400`: + ```bash curl https://ai.asepharyana.my.id/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ - "model": "minicpm-v-4.6", + "model": "minicpm5-1b-fable5-v2-thinking", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100 }' @@ -42,6 +44,23 @@ MODEL_PATH=/path/to/model.gguf ./target/release/llm-api ./target/release/llm-api ``` +### Environment variables + +| Var | Default | Description | +|-----|---------|-------------| +| `MODEL_PATH` | `/models/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking-Q8_0.gguf` | GGUF model file | +| `API_KEY` | *(empty = auth off)* | Bearer token required on `/v1/chat/completions` | +| `SERVER_PORT` | `4010` | Listen port | +| `RUST_LOG` | `info` | Log level | +| `N_CTX` / `N_BATCH` / `N_THREADS` | `8192` / `512` / `4` | llama.cpp context/batch/threads | + +### Smoke test (setelah deploy) + +```bash +./scripts/smoke-test.sh http://127.0.0.1:4010 # tanpa auth +./scripts/smoke-test.sh https://ai.asepharyana.my.id "$API_KEY" +``` + ## Deploy (Nix + systemd) ```bash @@ -53,6 +72,9 @@ nix build .#default --impure --option sandbox false ## Benchmark +> *Historic* (MiniCPM-V-4.6). Kept for reference; numbers predate the current +> MiniCPM5-1B Thinking model. + | Framework | Model Size | tok/s | vs PyTorch | |-----------|-----------|-------|------------| | PyTorch BF16 | 2.48 GB | 0.97 | 1.0x | diff --git a/scripts/smoke-test.sh b/scripts/smoke-test.sh new file mode 100755 index 0000000..aa6d28e --- /dev/null +++ b/scripts/smoke-test.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +# Smoke test untuk llm-api — jalankan di VPS setelah deploy. +# +# Memverifikasi alur inti: +# /health, /v1/models, chat non-streaming, chat streaming (SSE + [DONE]), +# penolakan model tidak dikenal (400), dan request dengan tools. +# +# Usage: +# ./scripts/smoke-test.sh [BASE_URL] [API_KEY] +# BASE_URL default: http://127.0.0.1:4010 +set -euo pipefail + +BASE_URL="${1:-http://127.0.0.1:4010}" +API_KEY="${2:-}" +MODEL_ID="minicpm5-1b-fable5-v2-thinking" + +AUTH=() +if [[ -n "$API_KEY" ]]; then + AUTH=(-H "Authorization: Bearer $API_KEY") +fi + +pass=0 +fail=0 +check() { + local name="$1" ok="$2" + if [[ "$ok" == "0" ]]; then + echo " PASS: $name" + pass=$((pass + 1)) + else + echo " FAIL: $name" + fail=$((fail + 1)) + fi +} + +echo "== 1. Health ==" +health=$(curl -sf "$BASE_URL/health") || { echo "FAIL: /health unreachable"; exit 1; } +echo "$health" | jq -e ".status == \"ok\" and .model == \"$MODEL_ID\"" >/dev/null +check "health melaporkan $MODEL_ID" $? + +echo "== 2. Models ==" +models=$(curl -sf "$BASE_URL/v1/models") +echo "$models" | jq -e ".data[0].id == \"$MODEL_ID\"" >/dev/null +check "models mencantumkan $MODEL_ID" $? + +echo "== 3. Chat non-streaming ==" +resp=$(curl -sf "${AUTH[@]}" -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"Say hi\"}],\"max_tokens\":64}" \ + "$BASE_URL/v1/chat/completions") +echo "$resp" | jq -e '.choices[0].message.content | type == "string"' >/dev/null +check "non-streaming mengembalikan content" $? +echo "$resp" | jq -e '.usage.total_tokens > 0' >/dev/null +check "non-streaming berisi usage" $? + +echo "== 4. Chat streaming ==" +stream=$(curl -sfN "${AUTH[@]}" -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"Count from 1 to 5\"}],\"stream\":true,\"max_tokens\":64}" \ + "$BASE_URL/v1/chat/completions") + +echo "$stream" | grep -q '\[DONE\]' +check "stream diakhiri [DONE]" $? + +echo "$stream" | grep -q '"finish_reason":"stop"\|"finish_reason":"length"' +check "stream punya finish_reason" $? + +# Gabungkan semua delta content untuk memastikan output tidak kosong +# dan tidak bocor markup (mis. / <|im_end|>). +joined=$(echo "$stream" | grep '^data: ' | sed 's/^data: //' \ + | grep -v '\[DONE\]' \ + | jq -r 'select((.choices? // []) | length > 0) | .choices[0].delta.content // empty' 2>/dev/null \ + | tr -d '\n') +if [[ -z "$joined" ]]; then + # Sebagian output mungkin semua ber-label reasoning_content; cek keduanya. + joined=$(echo "$stream" | grep '^data: ' | sed 's/^data: //' \ + | grep -v '\[DONE\]' \ + | jq -r 'select((.choices? // []) | length > 0) | (.choices[0].delta.content // .choices[0].delta.reasoning_content) // empty' 2>/dev/null \ + | tr -d '\n') +fi +[[ -n "$joined" ]] +check "stream menghasilkan teks" $? +if [[ -n "$joined" ]]; then + if [[ "$joined" == *"<"* ]]; then + check "tidak ada markup bocor di stream" 1 + else + check "tidak ada markup bocor di stream" 0 + fi +fi + +echo "== 5. Model tidak dikenal ditolak ==" +status=$(curl -s -o /dev/null -w "%{http_code}" "${AUTH[@]}" -H "Content-Type: application/json" \ + -d '{"model":"minicpm-v-4.6","messages":[{"role":"user","content":"hi"}]}' \ + "$BASE_URL/v1/chat/completions") +[[ "$status" == "400" ]] +check "unknown model -> 400" $? + +echo "== 6. Request dengan tools diterima ==" +status2=$(curl -s -o /dev/null -w "%{http_code}" "${AUTH[@]}" -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"What's the weather in Jakarta?\"}],\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"description\":\"Get weather\",\"parameters\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\"}}}}}],\"max_tokens\":128}" \ + "$BASE_URL/v1/chat/completions") +[[ "$status2" == "200" ]] +check "tools request sukses (200)" $? + +echo +echo "Result: $pass passed, $fail failed" +[[ $fail -eq 0 ]] diff --git a/src/application/chat/mod.rs b/src/application/chat/mod.rs index c5e6615..184c6a7 100644 --- a/src/application/chat/mod.rs +++ b/src/application/chat/mod.rs @@ -1,3 +1,6 @@ pub mod use_cases; -pub use use_cases::{build_prompt, build_sampler, clean_text, parse_tool_calls, SamplerParams}; +pub use use_cases::{ + build_prompt, build_sampler, clean_text, parse_tool_calls, split_stream_chunk, validate_model, + SamplerParams, +}; diff --git a/src/application/chat/templates/chat_template.jinja b/src/application/chat/templates/chat_template.jinja index 9efeabd..65856f0 100644 --- a/src/application/chat/templates/chat_template.jinja +++ b/src/application/chat/templates/chat_template.jinja @@ -6,7 +6,7 @@ {{- "\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' }} + {{- '\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, wrap each call in tags:\n\n\nvalue\n\n' }} {%- endset %} {{- '<|im_start|>system\n' }} @@ -26,17 +26,8 @@ {%- 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 %} + {#- Tool-call XML from history is pre-embedded in content by build_prompt. #} + {{- '<|im_start|>assistant\n' + content + '<|im_end|>\n' }} {%- elif message.role == "tool" %} {{- '<|im_start|>user\n\n' + content + '\n<|im_end|>\n' }} {%- endif %} diff --git a/src/application/chat/use_cases.rs b/src/application/chat/use_cases.rs index 452637a..d0aa200 100644 --- a/src/application/chat/use_cases.rs +++ b/src/application/chat/use_cases.rs @@ -5,21 +5,21 @@ use std::collections::HashMap; -use minijinja::{Environment, Value}; use llama_cpp_2::sampling::LlamaSampler; +use minijinja::{Environment, Value}; use crate::domain::entity::{ChatMessage, ChatRequest, ToolCall, ToolCallFunction}; -/// Build a prompt from messages using the GGUF's Jinja chat template. +/// Build a prompt from messages using the project's Jinja chat template. /// -/// Renders the model's baked-in template via minijinja, passing the message -/// history, optional tool definitions, and generation-prompt switches. +/// Renders the template (see `templates/chat_template.jinja`) via minijinja, +/// passing the message history, optional tool definitions, and the +/// generation-prompt switches. The template owns the `` BOS token, so +/// tokenization must NOT add another one (`AddBos::Never`). pub fn build_prompt( - model: &llama_cpp_2::model::LlamaModel, messages: &[ChatMessage], - _tools: &Option>, + tools: &Option>, ) -> Result { - // Load the GGUF's chat template (embedded at crate build time) let template_str = include_str!("templates/chat_template.jinja"); let mut env = Environment::new(); @@ -31,50 +31,55 @@ pub fn build_prompt( serde_json::to_string(value).unwrap_or_default() }); - let tmpl = env.get_template("chat") + let tmpl = env + .get_template("chat") .map_err(|e| format!("Template get error: {e}"))?; - // Build messages as serde_json::Value for minijinja + // Build messages as serde_json::Value for minijinja. + // Assistant tool calls from history are embedded directly into the content + // as XML (in the exact format the model is told to emit) — the template + // cannot accumulate `set` variables across a loop, so this is built here. let mut msgs_val: Vec = Vec::new(); for msg in messages { let mut m: HashMap = HashMap::new(); m.insert("role".into(), Value::from(msg.role.clone())); - let content = msg.content.clone().unwrap_or_default(); + let mut 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 { - // 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(); - Value::from_serialize(&serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": args, + for tc in tcs { + if tc.call_type == "function" { + content + .push_str(&format!("\n\n\n", tc.function.name)); + let args: serde_json::Value = + serde_json::from_str(&tc.function.arguments).unwrap_or_default(); + if let Some(obj) = args.as_object() { + for (k, v) in obj { + // Strings stay raw (no JSON quotes) so they round-trip + // through parse_tool_calls unchanged. + let rendered = match v { + serde_json::Value::String(s) => s.clone(), + other => other.to_string(), + }; + content + .push_str(&format!("{rendered}\n")); + } } - })) - }).collect(); - m.insert("tool_calls".into(), Value::from(tcs_val)); + content.push_str("\n"); + } + } } } - // 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 { - m.insert("content".into(), Value::from(content)); - } + // NOTE: the template wraps tool-role content in ; do not + // wrap here or it would be double-wrapped. + m.insert("content".into(), Value::from(content)); msgs_val.push(Value::from(m)); } - // BOS token for sentencepiece / unigram models + // BOS token for sentencepiece / unigram models (template-owned). let bos_token: &str = ""; // Build context @@ -84,6 +89,14 @@ pub fn build_prompt( ctx.insert("add_generation_prompt".into(), Value::from(true)); ctx.insert("enable_thinking".into(), Value::from(true)); + // Tool definitions (optional) — previously dead, now actually rendered. + if let Some(tools) = tools { + if !tools.is_empty() { + let tools_val: Vec = tools.iter().map(Value::from_serialize).collect(); + ctx.insert("tools".into(), Value::from(tools_val)); + } + } + // Render let result = tmpl .render(&ctx) @@ -92,6 +105,17 @@ pub fn build_prompt( Ok(result) } +/// Validate that the requested model matches the single served model. +pub fn validate_model(model: &str) -> Result<(), String> { + if model != crate::config::MODEL_ID { + return Err(format!( + "Unknown model '{model}'. Available: {}", + crate::config::MODEL_ID + )); + } + Ok(()) +} + /// Parameters for building a [`LlamaSampler`] chain. pub struct SamplerParams { pub temperature: Option, @@ -174,13 +198,13 @@ pub fn build_sampler(params: &SamplerParams) -> LlamaSampler { /// /// For thinking models, returns (reasoning, cleaned_answer). pub fn clean_text(text: &str) -> (String, String) { - let text = text.replace("<|im_end|>", "") - .replace("<|im_start|>", ""); + 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() + let reasoning = text[..close_idx] + .trim() .trim_start_matches("") .trim() .to_string(); @@ -250,9 +274,9 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec) { for line in lines { let line = line.trim(); - if let Some(param) = - line.strip_prefix("')) + if let Some(param) = line + .strip_prefix("')) { if let Some(p) = current_param.take() { args_map.insert( @@ -275,7 +299,10 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec) { } } if let Some(p) = current_param.take() { - args_map.insert(p, serde_json::Value::String(current_value.trim().to_string())); + args_map.insert( + p, + serde_json::Value::String(current_value.trim().to_string()), + ); } let args_json = serde_json::Value::Object(args_map).to_string(); @@ -298,3 +325,325 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec) { (cleaned, tool_calls) } + +// ═══════════════════════════════════════════════════════════════ +// STREAMING CHUNK SPLITTING +// ═══════════════════════════════════════════════════════════════ + +/// Remove special tokens and tool-call XML markup from a text fragment. +/// +/// Handles both fixed tags (`<|im_end|>`, ``, ``, …) and the +/// attribute-bearing openers used by this model's tool format (``, +/// ``), even when a tag straddles a token boundary. +fn strip_markup(text: &str) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + + while !rest.is_empty() { + let Some(idx) = rest.find('<') else { + out.push_str(rest); + break; + }; + out.push_str(&rest[..idx]); + let tail = &rest[idx..]; + + // Fixed tags (no attribute content). + let fixed = [ + "<|im_start|>", + "<|im_end|>", + "", + "", + "", + "", + "", + "", + ]; + if let Some(tag) = fixed.iter().find(|t| tail.starts_with(**t)) { + rest = &tail[tag.len()..]; + continue; + } + + // Attribute-bearing openers: / . + if let Some(attr) = tail + .strip_prefix("` boundary was already crossed **before** +/// this fragment (i.e. it is not the chunk containing the first ``). +/// * Whitespace **inside** a fragment is preserved — only the whitespace +/// sitting immediately around the `` boundary is trimmed, so +/// reasoning does not end with a dangling newline and content does not begin +/// with one. (Trimming every fragment corrupted inter-word spaces.) +/// * Before the first ``, everything is emitted as `reasoning_content`; +/// after it, as `content`. A stray second `` is stripped, not split. +pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option, String) { + if !think_done { + if let Some(pos) = new_text.find("") { + let mut before = strip_markup(&new_text[..pos]); + let mut after = strip_markup(&new_text[pos + 8..]); + + while before.ends_with(['\n', ' ', '\t']) { + before.pop(); + } + while after.starts_with(['\n', ' ', '\t']) { + after.remove(0); + } + + let reasoning = if before.is_empty() { + None + } else { + Some(before) + }; + (reasoning, after) + } else { + // Still thinking — everything is reasoning. + let cleaned = strip_markup(new_text); + let reasoning = if cleaned.is_empty() { + None + } else { + Some(cleaned) + }; + (reasoning, String::new()) + } + } else { + // Think phase already ended — everything is content; strip stray markup. + (None, strip_markup(new_text)) + } +} + +// ═══════════════════════════════════════════════════════════════ +// TESTS +// ═══════════════════════════════════════════════════════════════ + +#[cfg(test)] +mod tests { + use super::*; + use crate::domain::entity::{ + ChatMessage, ToolCallFunction, ToolCallResponse, ToolDef, ToolFunction, + }; + + fn msg(role: &str, content: &str) -> ChatMessage { + ChatMessage { + role: role.into(), + content: Some(content.into()), + tool_calls: None, + tool_call_id: None, + name: None, + } + } + + fn tool_def() -> ToolDef { + ToolDef { + tool_type: "function".into(), + function: ToolFunction { + name: "get_weather".into(), + description: "Get weather for a city".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"] + }), + }, + } + } + + // ── split_stream_chunk ── + + #[test] + fn split_chunk_before_think_is_reasoning() { + let (reasoning, content) = split_stream_chunk("Hello ", false); + assert_eq!(reasoning.as_deref(), Some("Hello ")); + assert_eq!(content, ""); + } + + #[test] + fn split_chunk_preserves_internal_spaces() { + // Regression: trimming every fragment used to eat inter-word spaces. + let (r1, _) = split_stream_chunk("Hello", false); + let (r2, _) = split_stream_chunk(" world", false); + assert_eq!(format!("{}{}", r1.unwrap(), r2.unwrap()), "Hello world"); + } + + #[test] + fn split_chunk_boundary_trims_only_edges() { + let (reasoning, content) = split_stream_chunk("question\n\n\nAnswer ", false); + assert_eq!(reasoning.as_deref(), Some("question")); + assert_eq!(content, "Answer "); + } + + #[test] + fn split_chunk_after_think_is_content() { + let (reasoning, content) = split_stream_chunk(" answer", true); + assert_eq!(reasoning, None); + assert_eq!(content, " answer"); + } + + #[test] + fn split_chunk_stray_think_tag_is_stripped_not_split() { + // A second (already past the boundary) must not restart + // reasoning classification. + let (reasoning, content) = split_stream_chunk("...more", true); + assert_eq!(reasoning, None); + assert_eq!(content, "...more"); + } + + #[test] + fn split_chunk_strips_special_and_markup() { + let (reasoning, content) = split_stream_chunk( + "<|im_end|>Hello\n", + false, + ); + assert_eq!(reasoning.as_deref(), Some("Hello")); + assert_eq!(content, ""); + } + + // ── clean_text ── + + #[test] + fn clean_text_splits_reasoning_and_answer() { + let (reasoning, answer) = + clean_text("Let me think\nabout it\nThe answer is 42."); + assert_eq!(reasoning, "Let me think\nabout it"); + assert_eq!(answer, "The answer is 42."); + } + + #[test] + fn clean_text_no_think_returns_answer() { + let (reasoning, answer) = clean_text("Just an answer"); + assert_eq!(reasoning, ""); + assert_eq!(answer, "Just an answer"); + } + + // ── parse_tool_calls ── + + #[test] + fn parse_single_tool_call() { + let text = "I'll look that up.\n\n\nJakarta\n\n"; + let (cleaned, calls) = parse_tool_calls(text); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].function.name, "get_weather"); + assert!(calls[0].function.arguments.contains("Jakarta")); + assert!(!cleaned.contains("")); + } + + #[test] + fn parse_multiple_tool_calls() { + let text = "\n\n1\n\n\n\n\n2\n\n"; + let (_cleaned, calls) = parse_tool_calls(text); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].function.name, "a"); + assert_eq!(calls[1].function.name, "b"); + } + + #[test] + fn parse_malformed_tool_call_returns_empty() { + let (cleaned, calls) = parse_tool_calls("no calls here"); + assert!(calls.is_empty()); + assert_eq!(cleaned, "no calls here"); + } + + // ── validate_model ── + + #[test] + fn validate_model_accepts_served_id() { + assert!(validate_model(crate::config::MODEL_ID).is_ok()); + } + + #[test] + fn validate_model_rejects_unknown_id() { + assert!(validate_model("minicpm-v-4.6").is_err()); + } + + // ── build_sampler (smoke — no model required) ── + + #[test] + fn build_sampler_constructs_for_common_params() { + let params = SamplerParams { + temperature: Some(0.8), + top_p: Some(0.9), + top_k: Some(40), + min_p: Some(0.05), + repeat_penalty: Some(1.1), + frequency_penalty: Some(0.0), + presence_penalty: Some(0.0), + seed: Some(42), + }; + let _ = build_sampler(¶ms); + + let greedy = SamplerParams { + temperature: Some(0.0), + ..params + }; + let _ = build_sampler(&greedy); + } + + // ── build_prompt ── + + #[test] + fn build_prompt_renders_messages() { + let messages = vec![msg("system", "You are helpful."), msg("user", "Hi!")]; + let prompt = build_prompt(&messages, &None).unwrap(); + assert!(prompt.starts_with("")); + assert!(prompt.contains("<|im_start|>system\nYou are helpful.")); + assert!(prompt.contains("<|im_start|>user\nHi!")); + assert!(prompt.contains("<|im_start|>assistant\n\n")); + } + + #[test] + fn build_prompt_includes_tool_definitions() { + let messages = vec![msg("user", "What's the weather?")]; + let prompt = build_prompt(&messages, &Some(vec![tool_def()])).unwrap(); + assert!(prompt.contains("")); + assert!(prompt.contains("get_weather")); + assert!(prompt.contains("Tool usage guidelines")); + } + + #[test] + fn build_prompt_wraps_tool_response_once() { + let messages = vec![msg("user", "Weather?"), msg("tool", "Sunny")]; + let prompt = build_prompt(&messages, &None).unwrap(); + assert_eq!(prompt.matches("").count(), 1); + assert_eq!(prompt.matches("").count(), 1); + } + + #[test] + fn build_prompt_renders_assistant_tool_calls_history() { + let assistant = ChatMessage { + role: "assistant".into(), + content: Some("".into()), + tool_calls: Some(vec![ToolCallResponse { + id: "call_1".into(), + call_type: "function".into(), + function: ToolCallFunction { + name: "get_weather".into(), + arguments: r#"{"city":"Jakarta"}"#.into(), + }, + }]), + tool_call_id: None, + name: None, + }; + let prompt = build_prompt(&[assistant], &None).unwrap(); + assert!(prompt.contains("")); + assert!(prompt.contains("Jakarta")); + } +} diff --git a/src/bootstrap/mod.rs b/src/bootstrap/mod.rs index 2710a91..da6b2c4 100644 --- a/src/bootstrap/mod.rs +++ b/src/bootstrap/mod.rs @@ -32,15 +32,12 @@ impl Application { pub async fn build() -> anyhow::Result { // Initialize tracing let env_filter = EnvFilter::new(&CONFIG.log_level); - tracing_subscriber::fmt() - .with_env_filter(env_filter) - .init(); + tracing_subscriber::fmt().with_env_filter(env_filter).init(); tracing::info!("🚀 LLM API starting up..."); // Load model (fail-fast) - let engine = LlamaEngine::load().map_err(|e| { - anyhow::anyhow!("Failed to initialize LLM engine: {e}") - })?; + let engine = LlamaEngine::load() + .map_err(|e| anyhow::anyhow!("Failed to initialize LLM engine: {e}"))?; let engine = Arc::new(engine); let state = Arc::new(AppState { engine }); @@ -51,10 +48,7 @@ impl Application { // Bind listener let addr = format!("0.0.0.0:{}", CONFIG.server_port); let listener = TcpListener::bind(&addr).await?; - tracing::info!( - "Server listening on {}", - listener.local_addr()? - ); + tracing::info!("Server listening on {}", listener.local_addr()?); Ok(Self { port: CONFIG.server_port, diff --git a/src/config/mod.rs b/src/config/mod.rs index ed756fc..628b448 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -7,6 +7,14 @@ use std::sync::LazyLock; const DEFAULT_MODEL_PATH: &str = "/models/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking-Q8_0.gguf"; pub const MODEL_ID: &str = "minicpm5-1b-fable5-v2-thinking"; +/// Read an env var as `T`, falling back to `default` on absence or parse error. +fn env_or(key: &str, default: T) -> T { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + /// Application configuration loaded at startup from environment variables. #[derive(Debug, Clone)] pub struct AppConfig { @@ -44,9 +52,9 @@ impl AppConfig { .and_then(|v| v.parse().ok()) .unwrap_or(4010), log_level: std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string()), - n_ctx: 8192, - n_batch: 512, - n_threads: 4, + n_ctx: env_or("N_CTX", 8192), + n_batch: env_or("N_BATCH", 512), + n_threads: env_or("N_THREADS", 4), } } } diff --git a/src/domain/entity/mod.rs b/src/domain/entity/mod.rs index 8fd985f..14ab11a 100644 --- a/src/domain/entity/mod.rs +++ b/src/domain/entity/mod.rs @@ -104,7 +104,7 @@ pub struct ResponseMessage { pub tool_calls: Option>, } -#[derive(Serialize)] +#[derive(Serialize, Clone)] pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, @@ -139,6 +139,36 @@ pub struct ToolCallResponse { pub function: ToolCallFunction, } +// ═══════════════════════════════════════════════════════════════ +// GENERATION FINISH REASON +// ═══════════════════════════════════════════════════════════════ + +/// Why a generation run stopped. Produced by the engine and rendered as the +/// OpenAI `finish_reason` at the presentation layer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FinishReason { + /// Natural end-of-generation (EOS token) or a stop sequence matched. + Stop, + /// `max_tokens` exhausted. + Length, + /// A complete `` block was emitted. + ToolCalls, + /// Generation aborted early (e.g. client disconnected). + Aborted, +} + +impl FinishReason { + /// Map to the OpenAI-compatible `finish_reason` string. + pub fn as_str(&self) -> &'static str { + match self { + FinishReason::Stop => "stop", + FinishReason::Length => "length", + FinishReason::ToolCalls => "tool_calls", + FinishReason::Aborted => "stop", + } + } +} + // ═══════════════════════════════════════════════════════════════ // SSE (STREAMING) TYPES // ═══════════════════════════════════════════════════════════════ @@ -151,6 +181,59 @@ pub struct SseChunk { pub created: i64, pub model: String, pub choices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, +} + +impl SseChunk { + /// A delta chunk carrying partial reasoning/content/tool-call output. + pub fn delta(id: String, created: i64, model: String, delta: SseDelta) -> Self { + Self { + id, + object: "chat.completion.chunk".into(), + created, + model, + choices: vec![SseChoice { + index: 0, + delta, + finish_reason: None, + }], + usage: None, + } + } + + /// The final chunk carrying the `finish_reason` (no token deltas). + pub fn finish(id: String, created: i64, model: String, finish_reason: &str) -> Self { + Self { + id, + object: "chat.completion.chunk".into(), + created, + model, + choices: vec![SseChoice { + index: 0, + delta: SseDelta { + role: None, + content: None, + tool_calls: None, + reasoning_content: None, + }, + finish_reason: Some(finish_reason.into()), + }], + usage: None, + } + } + + /// A trailing chunk with token usage and empty choices (OpenAI convention). + pub fn usage(id: String, created: i64, model: String, usage: Usage) -> Self { + Self { + id, + object: "chat.completion.chunk".into(), + created, + model, + choices: vec![], + usage: Some(usage), + } + } } #[derive(Serialize)] diff --git a/src/infrastructure/llama/engine.rs b/src/infrastructure/llama/engine.rs index 5fcd007..58ab57a 100644 --- a/src/infrastructure/llama/engine.rs +++ b/src/infrastructure/llama/engine.rs @@ -4,6 +4,7 @@ //! lifetime transmute), tokenization, and generation. use std::num::NonZeroU32; +use std::sync::Mutex; use llama_cpp_2::context::params::LlamaContextParams; use llama_cpp_2::context::LlamaContext; @@ -14,10 +15,10 @@ use llama_cpp_2::model::{AddBos, LlamaModel}; use llama_cpp_2::sampling::LlamaSampler; use llama_cpp_2::token::LlamaToken; use llama_cpp_2::TokenToStringError; -use tokio::sync::Mutex; use tracing::info; use crate::config::CONFIG; +use crate::domain::entity::FinishReason; use crate::domain::LlmError; // ── Thread-safe wrapper for raw llama.cpp context ── @@ -95,17 +96,28 @@ impl CtxInner { // ── LlamaEngine ── +/// Outcome of a generation run. +pub struct GenerationOutcome { + /// Generated tokens (stop-sequence and EOG tokens are excluded). + pub tokens: Vec, + /// Accumulated decoded text (raw, before markup/special-token cleaning). + pub text: String, + /// Why generation stopped. + pub finish: FinishReason, +} + /// Safe interface to a llama.cpp model and inference context. /// /// All access to the underlying context is serialized through a `Mutex`, /// so only one generation can happen at a time. This is intentional — -/// the model is designed for sequential inference. +/// the model is designed for sequential inference. Generation is synchronous +/// and must be invoked from the tokio blocking pool (`spawn_blocking`). pub struct LlamaEngine { /// The loaded model (read-only after load, safe to share). pub model: LlamaModel, /// The inference context (single-threaded access via Mutex). - pub ctx: Mutex, + ctx: Mutex, } impl LlamaEngine { @@ -117,17 +129,17 @@ impl LlamaEngine { /// cannot be created. pub fn load() -> Result { info!("Initializing llama backend..."); - let backend = LlamaBackend::init().map_err(|e| { - LlmError::Model(format!("Backend init failed: {e}")) - })?; + let backend = LlamaBackend::init() + .map_err(|e| LlmError::Model(format!("Backend init failed: {e}")))?; // Backend must outlive model and context. We leak it to achieve 'static // lifetime since the engine lives for the program lifetime. let backend: &'static LlamaBackend = Box::leak(Box::new(backend)); info!("Loading model: {}", CONFIG.model_path); - let model = LlamaModel::load_from_file(backend, &CONFIG.model_path, &LlamaModelParams::default()) - .map_err(|e| LlmError::Model(format!("Failed to load model: {e}")))?; + let model = + LlamaModel::load_from_file(backend, &CONFIG.model_path, &LlamaModelParams::default()) + .map_err(|e| LlmError::Model(format!("Failed to load model: {e}")))?; info!(" Vocab: {}", model.n_vocab()); info!(" Params: {}", model.n_params()); info!(" Layers: {}", model.n_layer()); @@ -156,9 +168,12 @@ impl LlamaEngine { } /// Tokenize a prompt string into tokens. + /// + /// `AddBos::Never`: the chat template already prepends the `` BOS token, + /// so adding another here would produce a double BOS. pub fn tokenize(&self, prompt: &str) -> Result, LlmError> { self.model - .str_to_token(prompt, AddBos::Always) + .str_to_token(prompt, AddBos::Never) .map_err(|e| LlmError::Model(format!("Tokenization failed: {e}"))) } @@ -177,37 +192,28 @@ impl LlamaEngine { String::from_utf8(bytes).unwrap_or_default() } - /// Decode multiple tokens to a single string. - pub fn decode_tokens(&self, tokens: &[LlamaToken]) -> String { - let mut out = String::with_capacity(tokens.len() * 4); - for &token in tokens { - out.push_str(&self.decode_token(token)); - } - out - } - - /// Check if a token is an end-of-generation token. - pub fn is_eog(&self, token: LlamaToken) -> bool { - self.model.is_eog_token(token) - } - - /// Return a reference to the context mutex for advanced operations. - pub fn ctx(&self) -> &Mutex { - &self.ctx - } - - /// Generate tokens (non-streaming) and return output tokens, cleaned text, and tool calls. + /// Generate tokens and invoke `on_token` for each one. /// - /// Locks the context mutex, prefill the prompt, then iterates sampling + decoding - /// until EOG, max_tokens, stop sequence, or tool call completion. - pub async fn generate( + /// Synchronous (CPU-bound) — call from `spawn_blocking`. Locks the context, + /// prefills the prompt, then iterates sampling + decoding until EOG, + /// `max_tokens`, a stop sequence, or a complete `` block. + /// + /// `on_token` is invoked for every generated token (after leading-EOG + /// skipping, before it is decoded into the KV cache) and may return `false` + /// to abort early (e.g. the streaming client disconnected). + pub fn generate( &self, input_tokens: &[LlamaToken], - sampler: &mut SendSampler, + sampler: &mut LlamaSampler, max_tokens: u32, stop: &[String], - ) -> Result<(Vec, String), LlmError> { - let mut inner = self.ctx.lock().await; + enable_tool_detection: bool, + on_token: &mut dyn FnMut(LlamaToken, &str) -> bool, + ) -> Result { + let mut inner = self + .ctx + .lock() + .map_err(|_| LlmError::Internal("context mutex poisoned".into()))?; inner.clear(); inner .prefill(input_tokens) @@ -215,59 +221,66 @@ impl LlamaEngine { let mut output: Vec = Vec::new(); let mut text_buf = String::new(); - let mut stop_now = false; + let mut finish = FinishReason::Length; let mut current = inner.sample(sampler); - // Skip leading EOS tokens (like <|im_end|> as first token) + // Skip leading EOG tokens (like a stray <|im_end|> right after the prompt) while output.is_empty() && self.model.is_eog_token(current) { let pos = input_tokens.len() as i32 + output.len() as i32; - if let Err(e) = inner.decode(current, pos) { - tracing::info!(" Decode error: {e}"); - break; - } + inner + .decode(current, pos) + .map_err(|e| LlmError::Model(format!("Decode: {e}")))?; current = inner.sample(sampler); } for _ in 0..max_tokens { if self.model.is_eog_token(current) { + finish = FinishReason::Stop; + break; + } + + let piece = self.decode_token(current); + + // Check stop sequences *before* committing, so the stop tokens never + // leak into the output text or the stream. + if stop + .iter() + .any(|s| !s.is_empty() && format!("{text_buf}{piece}").contains(s)) + { + finish = FinishReason::Stop; break; } let pos = input_tokens.len() as i32 + output.len() as i32; output.push(current); - - let piece = self.decode_token(current); text_buf.push_str(&piece); - // Check stop sequences - for s in stop { - if text_buf.contains(s) { - stop_now = true; - break; - } - } - if stop_now { - break; - } - - // Check for tool_call block completion - if text_buf.contains("") { - let close_count = text_buf.matches("").count(); - let open_count = text_buf.matches("").count(); - if open_count > 0 && close_count >= open_count { + // Complete block emitted? + if enable_tool_detection && text_buf.contains("") { + let open = text_buf.matches("").count(); + let close = text_buf.matches("").count(); + if close >= open { + finish = FinishReason::ToolCalls; break; } } - if let Err(e) = inner.decode(current, pos) { - tracing::info!(" Decode error: {e}"); + if !on_token(current, &piece) { + finish = FinishReason::Aborted; break; } + inner + .decode(current, pos) + .map_err(|e| LlmError::Model(format!("Decode: {e}")))?; current = inner.sample(sampler); } - Ok((output, text_buf)) + Ok(GenerationOutcome { + tokens: output, + text: text_buf, + finish, + }) } } diff --git a/src/presentation/handler/chat-ui/index.html b/src/presentation/handler/chat-ui/index.html index a04d87c..ad440c1 100644 --- a/src/presentation/handler/chat-ui/index.html +++ b/src/presentation/handler/chat-ui/index.html @@ -168,7 +168,7 @@ async function send() { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ - model: '', + model: 'minicpm5-1b-fable5-v2-thinking', messages: messages.slice(-5), // keep context window manageable stream: true, max_tokens: 1024, diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs index fef7843..d5e5d51 100644 --- a/src/presentation/handler/chat.rs +++ b/src/presentation/handler/chat.rs @@ -1,12 +1,18 @@ //! Chat completions endpoint — streaming and non-streaming. +//! +//! Both paths share the same synchronous generation core +//! ([`LlamaEngine::generate`]), which runs on the tokio blocking pool via +//! `spawn_blocking` so worker threads are not hogged by CPU-bound inference. +//! The streaming path forwards each generated token into an SSE channel. use std::convert::Infallible; use std::sync::Arc; use axum::extract::State; -use axum::response::sse::{Event, KeepAlive, Sse}; use axum::http::HeaderMap; -use axum::{Json, response::{IntoResponse, Response}}; +use axum::response::sse::{Event, KeepAlive, Sse}; +use axum::response::{IntoResponse, Response}; +use axum::Json; use chrono::Utc; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -14,7 +20,7 @@ use tracing::info; use crate::application::chat; use crate::domain::entity::{ - ChatRequest, ChatResponse, Choice, ResponseMessage, SseChunk, SseChoice, SseDelta, Usage, + ChatRequest, ChatResponse, Choice, FinishReason, ResponseMessage, SseChunk, SseDelta, Usage, }; use crate::infrastructure::llama::SendSampler; use crate::presentation::error::AppError; @@ -25,12 +31,14 @@ pub async fn chat_completions( State(state): State>, Json(req): Json, ) -> Result { + // Strict model validation — reject unknown model ids up front. + chat::validate_model(&req.model).map_err(AppError::BadRequest)?; + let max_tokens = req.max_tokens.unwrap_or(256).min(1024); let stop = req.stop.clone().unwrap_or_default(); - let prompt = chat::build_prompt(&state.engine.model, &req.messages, &req.tools) - .map_err(|e| AppError::LlmError(e))?; + let prompt = chat::build_prompt(&req.messages, &req.tools).map_err(AppError::LlmError)?; - // Tokenize + // Tokenize (fast — keep on the async thread). let input_tokens = state .engine .tokenize(&prompt) @@ -38,7 +46,7 @@ pub async fn chat_completions( let prompt_tokens = input_tokens.len() as u32; info!( - " Chat: {} prompt tokens, max_tokens={}, tools={}", + "Chat: {} prompt tokens, max_tokens={}, tools={}", prompt_tokens, max_tokens, req.tools.as_ref().is_some_and(|t| !t.is_empty()) @@ -49,7 +57,7 @@ pub async fn chat_completions( } else { handle_non_streaming(state.clone(), req, max_tokens, stop, input_tokens).await? }; - Ok(response.into_response()) + Ok(response) } // ── Non-streaming path ── @@ -66,28 +74,41 @@ async fn handle_non_streaming( let prompt_tokens = input_tokens.len() as u32; let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty()); let params = chat::SamplerParams::from_request(&req); - let mut sampler = SendSampler(chat::build_sampler(¶ms)); - let (output_tokens, raw_text) = state - .engine - .generate(&input_tokens, &mut sampler, max_tokens, &stop) - .await?; + let engine = state.engine.clone(); + let outcome = tokio::task::spawn_blocking(move || { + let mut sampler = SendSampler(chat::build_sampler(¶ms)); + engine.generate( + &input_tokens, + &mut sampler, + max_tokens, + &stop, + has_tools, + &mut |_token, _piece| true, + ) + }) + .await + .map_err(|e| AppError::Internal(format!("Generation task panicked: {e}")))? + .map_err(AppError::from)?; - 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 = outcome.tokens.len() as u32; info!(" {} generated tokens", completion_tokens); - let finish_reason = if has_tools && !tool_calls.is_empty() { - "tool_calls" - } else if completion_tokens < max_tokens { - "stop" - } else { - "length" + let (reasoning, cleaned) = chat::clean_text(&outcome.text); + let (output_text, tool_calls) = chat::parse_tool_calls(&cleaned); + + let finish_reason = match (outcome.finish, tool_calls.is_empty()) { + (FinishReason::ToolCalls, false) => "tool_calls", + // Model opened a but never completed it — don't claim a call. + (FinishReason::ToolCalls, true) => "stop", + (f, _) => f.as_str(), }; - let reasoning_opt = if reasoning.is_empty() { None } else { Some(reasoning) }; + let reasoning_opt = if reasoning.is_empty() { + None + } else { + Some(reasoning) + }; Ok(Json(ChatResponse { id: chat_id, @@ -130,309 +151,157 @@ async fn handle_streaming( let created = Utc::now().timestamp(); let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty()); let model_name = req.model.clone(); + let prompt_tokens = input_tokens.len() as u32; let params = chat::SamplerParams::from_request(&req); + let (tx, rx) = mpsc::channel::>(64); - tokio::spawn(async move { - // Role chunk - let role_chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: Some("assistant".into()), - content: None, - tool_calls: None, - reasoning_content: None, - }, - finish_reason: None, - }], - }) - .unwrap(); - if tx.send(Ok(Event::default().data(role_chunk))).await.is_err() { - return; - } + // First chunk: announce the assistant role. + let role_chunk = SseChunk::delta( + chat_id.clone(), + created, + model_name.clone(), + SseDelta { + role: Some("assistant".into()), + content: None, + tool_calls: None, + reasoning_content: None, + }, + ); + let role_event = serde_json::to_string(&role_chunk).unwrap(); + let _ = tx.send(Ok(Event::default().data(role_event))).await; - // Build sampler + let engine = state.engine.clone(); + tokio::task::spawn_blocking(move || { let mut sampler = SendSampler(chat::build_sampler(¶ms)); - - // Lock context - let mut inner = state.engine.ctx().lock().await; - inner.clear(); - if let Err(e) = inner.prefill(&input_tokens) { - info!(" Prefill error: {e}"); - return; - } - - let mut count = 0u32; let mut text_buf = String::new(); - let mut sent_len: usize = 0; // how many chars of text_buf have been sent - let mut think_done: bool = false; // true once seen - let mut current = inner.sample(&mut sampler); + let mut sent_len: usize = 0; + let mut think_done = false; - loop { - if count >= max_tokens { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { + let outcome = engine.generate( + &input_tokens, + &mut sampler, + max_tokens, + &stop, + has_tools, + &mut |_token, piece| { + text_buf.push_str(piece); + + // `was_thinking` is passed to split_stream_chunk so the chunk + // containing the first is treated as the boundary. + let was_thinking = !think_done; + if !think_done && text_buf.contains("") { + think_done = true; + } + + let new_text = &text_buf[sent_len..]; + if new_text.is_empty() { + return true; + } + + let (reasoning, content) = chat::split_stream_chunk(new_text, was_thinking); + + if let Some(reasoning) = reasoning { + let chunk = SseChunk::delta( + chat_id.clone(), + created, + model_name.clone(), + SseDelta { role: None, content: None, tool_calls: None, + reasoning_content: Some(reasoning), + }, + ); + let event = serde_json::to_string(&chunk).unwrap(); + if tx.blocking_send(Ok(Event::default().data(event))).is_err() { + return false; + } + } + + if !content.is_empty() { + let chunk = SseChunk::delta( + chat_id.clone(), + created, + model_name.clone(), + SseDelta { + role: None, + content: Some(content), + tool_calls: None, reasoning_content: None, }, - finish_reason: Some("length".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } + ); + let event = serde_json::to_string(&chunk).unwrap(); + if tx.blocking_send(Ok(Event::default().data(event))).is_err() { + return false; + } + } - // Skip leading EOS tokens (like <|im_end|> at start of generation) - if state.engine.is_eog(current) && text_buf.is_empty() { - // safety bound: don't skip more than 10 - if count >= max_tokens || count > 10 { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { + sent_len = text_buf.len(); + true + }, + ); + + match outcome { + Ok(outcome) => { + let completion_tokens = outcome.tokens.len() as u32; + let usage = Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + }; + + // Single-shot tool-calls delta (this model emits whole blocks). + let mut sent_tool_calls = false; + if outcome.finish == FinishReason::ToolCalls { + let (_cleaned, calls) = chat::parse_tool_calls(&outcome.text); + if !calls.is_empty() { + sent_tool_calls = true; + let chunk = SseChunk::delta( + chat_id.clone(), + created, + model_name.clone(), + SseDelta { role: None, content: None, - tool_calls: None, + tool_calls: Some(calls), reasoning_content: None, }, - finish_reason: Some("stop".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; + ); + let event = serde_json::to_string(&chunk).unwrap(); + let _ = tx.blocking_send(Ok(Event::default().data(event))); + } } - count += 1; - let pos = input_tokens.len() as i32 + count as i32; - if let Err(e) = inner.decode(current, pos) { - info!(" Decode error: {e}"); - break; - } - current = inner.sample(&mut sampler); - continue; - } - if state.engine.is_eog(current) { - let reason = if has_tools && text_buf.contains("") { - "tool_calls" - } else { - "stop" + let finish_reason = match (outcome.finish, sent_tool_calls) { + (FinishReason::ToolCalls, true) => "tool_calls", + (FinishReason::ToolCalls, false) => "stop", + (f, _) => f.as_str(), }; - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - reasoning_content: None, - }, - finish_reason: Some(reason.into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; + + let finish_chunk = + SseChunk::finish(chat_id.clone(), created, model_name.clone(), finish_reason); + let event = serde_json::to_string(&finish_chunk).unwrap(); + let _ = tx.blocking_send(Ok(Event::default().data(event))); + + let usage_chunk = SseChunk::usage(chat_id, created, model_name, usage); + let event = serde_json::to_string(&usage_chunk).unwrap(); + let _ = tx.blocking_send(Ok(Event::default().data(event))); } - - let piece = state.engine.decode_token(current); - - // Push into buffer - text_buf.push_str(&piece); - - // Detect transition - if !think_done && text_buf.contains("") { - think_done = true; + Err(e) => { + // Surface the error instead of silently truncating the stream. + let error_body = serde_json::json!({ + "error": { + "message": e.to_string(), + "type": "server_error", + } + }); + let _ = tx.blocking_send(Ok(Event::default().data(error_body.to_string()))); } - - // Find new text since last send - let new_text = &text_buf[sent_len..]; // everything not yet streamed - if new_text.is_empty() { - // Nothing new to send; skip straight to decode - let pos = input_tokens.len() as i32 + count as i32; - if let Err(e) = inner.decode(current, pos) { - info!(" Decode error: {e}"); - break; - } - count += 1; - current = inner.sample(&mut sampler); - continue; - } - - // Strip special tokens from the NEW text chunk - // Split at if present — before goes to reasoning, after to content - let (reasoning_part, content_part) = if let Some(pos) = new_text.find("") { - let before = new_text[..pos] - .replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .trim() - .to_string(); - let after = new_text[pos + 8..] - .replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .trim() - .to_string(); - think_done = true; - (Some(before), after) - } else if think_done { - let cleaned = new_text - .replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .trim() - .to_string(); - (None, cleaned) - } else { - let cleaned = new_text - .replace("<|im_end|>", "") - .replace("<|im_start|>", "") - .replace("", "") - .trim() - .to_string(); - (Some(cleaned), String::new()) - }; - - // Send reasoning part (before , or entire text if still thinking) - if let Some(ref r) = reasoning_part { - if !r.is_empty() { - let delta = SseDelta { - role: None, - content: None, - reasoning_content: Some(r.clone()), - tool_calls: None, - }; - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta, - finish_reason: None, - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - } - } - - // Send content part (after , or never if model doesn't think) - if !content_part.is_empty() { - let delta = SseDelta { - role: None, - content: Some(content_part), - reasoning_content: None, - tool_calls: None, - }; - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta, - finish_reason: None, - }], - }) - .unwrap(); - if tx.send(Ok(Event::default().data(chunk))).await.is_err() { - break; - } - } - - sent_len = text_buf.len(); - - // Check stop sequences - let mut stop_now = false; - for s in &stop { - if text_buf.contains(s) { - stop_now = true; - break; - } - } - if stop_now { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - reasoning_content: None, - }, - finish_reason: Some("stop".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - - // Check tool call completeness - if has_tools && text_buf.contains("") { - let open = text_buf.matches("").count(); - let close = text_buf.matches("").count(); - if close >= open { - let chunk = serde_json::to_string(&SseChunk { - id: chat_id.clone(), - object: "chat.completion.chunk".into(), - created, - model: model_name.clone(), - choices: vec![SseChoice { - index: 0, - delta: SseDelta { - role: None, - content: None, - tool_calls: None, - reasoning_content: None, - }, - finish_reason: Some("tool_calls".into()), - }], - }) - .unwrap(); - let _ = tx.send(Ok(Event::default().data(chunk))).await; - break; - } - } - - let pos = input_tokens.len() as i32 + count as i32; - if let Err(e) = inner.decode(current, pos) { - info!(" Decode error: {e}"); - break; - } - count += 1; - current = inner.sample(&mut sampler); } + + // OpenAI-compatible terminator. + let _ = tx.blocking_send(Ok(Event::default().data("[DONE]"))); }); let stream = ReceiverStream::new(rx); diff --git a/src/presentation/handler/chat_ui.rs b/src/presentation/handler/chat_ui.rs index 451a59f..8da8e0a 100644 --- a/src/presentation/handler/chat_ui.rs +++ b/src/presentation/handler/chat_ui.rs @@ -9,7 +9,10 @@ const HTML: &str = include_str!("chat-ui/index.html"); pub async fn chat_ui() -> Response { ( StatusCode::OK, - [(header::CONTENT_TYPE, HeaderValue::from_static("text/html; charset=utf-8"))], + [( + header::CONTENT_TYPE, + HeaderValue::from_static("text/html; charset=utf-8"), + )], HTML, ) .into_response() diff --git a/src/presentation/handler/health.rs b/src/presentation/handler/health.rs index b1f7e9b..9dfd106 100644 --- a/src/presentation/handler/health.rs +++ b/src/presentation/handler/health.rs @@ -8,6 +8,7 @@ use crate::domain::entity::HealthResponse; pub async fn health_check() -> Json { Json(HealthResponse { status: "ok".into(), - model: format!("{MODEL_ID}-q8_0"), + // Report the exact model id served by /v1/models (no extra suffix). + model: MODEL_ID.into(), }) } diff --git a/src/presentation/middleware/auth.rs b/src/presentation/middleware/auth.rs index d2cd085..766ec56 100644 --- a/src/presentation/middleware/auth.rs +++ b/src/presentation/middleware/auth.rs @@ -4,18 +4,18 @@ //! Only applied to routes that require authentication. use axum::extract::Request; -use axum::http::StatusCode; use axum::middleware::Next; use axum::response::Response; use crate::config::CONFIG; +use crate::presentation::error::AppError; /// Middleware that validates the Bearer token in the Authorization header. /// /// If `API_KEY` is not set (empty), authentication is disabled and /// all requests pass through. If set, the middleware rejects requests /// without a matching token. -pub async fn auth_middleware(request: Request, next: Next) -> Result { +pub async fn auth_middleware(request: Request, next: Next) -> Result { let api_key = &CONFIG.api_key; if api_key.is_empty() { return Ok(next.run(request).await); @@ -32,8 +32,6 @@ pub async fn auth_middleware(request: Request, next: Next) -> Result