refactor(chat): unify generation flow and fix tool/streaming bugs
- Unified synchronous generate() core with callback; both streaming and non-streaming paths run it via spawn_blocking (context: std Mutex). - build_prompt now passes tool definitions to the template (was dead) and embeds assistant tool-call history as XML matching the parser format; fixes double <tool_response> wrap and template set-scoping bug. - Tokenize with AddBos::Never (template owns <s>) to remove double BOS. - Streaming: preserve inter-word spaces (per-chunk trim removed), add [DONE] + usage chunk, emit error events, single-shot tool_calls delta. - Strict model validation (400 on unknown model); health/UI/README aligned to minicpm5-1b-fable5-v2-thinking; auth returns JSON errors; n_ctx/ n_batch/n_threads env-configurable. - Added 18 unit tests; cargo check/clippy/fmt clean. - scripts/smoke-test.sh for post-deploy verification on the VPS. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
asepharyana
co-authored by
Claude Opus 5
Claude Opus 5
parent
344bc195fa
commit
b636496497
@@ -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 |
|
||||
|
||||
Executable
+104
@@ -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. <tool_call> / <|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 ]]
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
{{- "\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>' }}
|
||||
{{- '\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, wrap each call in <tool_call> tags:\n<tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>' }}
|
||||
{%- 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<tool_call>\n<function=" + tool_call.function.name + ">\n" + (tool_call.function.arguments | tojson) + "\n</function>\n</tool_call>" %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{#- Tool-call XML from history is pre-embedded in content by build_prompt. #}
|
||||
{{- '<|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 %}
|
||||
|
||||
@@ -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 `<s>` 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<Vec<crate::domain::entity::ToolDef>>,
|
||||
tools: &Option<Vec<crate::domain::entity::ToolDef>>,
|
||||
) -> Result<String, String> {
|
||||
// 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<Value> = Vec::new();
|
||||
for msg in messages {
|
||||
let mut m: HashMap<String, Value> = 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<Value> = tcs.iter().map(|tc| {
|
||||
for tc in tcs {
|
||||
if tc.call_type == "function" {
|
||||
content
|
||||
.push_str(&format!("\n<tool_call>\n<function={}>\n", tc.function.name));
|
||||
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,
|
||||
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!("<parameter={k}>{rendered}</parameter>\n"));
|
||||
}
|
||||
}
|
||||
content.push_str("</function>\n</tool_call>");
|
||||
}
|
||||
}
|
||||
}))
|
||||
}).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 {
|
||||
// NOTE: the template wraps tool-role content in <tool_response>; 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 = "<s>";
|
||||
|
||||
// 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<Value> = 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<f32>,
|
||||
@@ -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 <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()
|
||||
let reasoning = text[..close_idx]
|
||||
.trim()
|
||||
.trim_start_matches("<think>")
|
||||
.trim()
|
||||
.to_string();
|
||||
@@ -250,8 +274,8 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
|
||||
|
||||
for line in lines {
|
||||
let line = line.trim();
|
||||
if let Some(param) =
|
||||
line.strip_prefix("<parameter=")
|
||||
if let Some(param) = line
|
||||
.strip_prefix("<parameter=")
|
||||
.and_then(|s| s.strip_suffix('>'))
|
||||
{
|
||||
if let Some(p) = current_param.take() {
|
||||
@@ -275,7 +299,10 @@ pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
|
||||
}
|
||||
}
|
||||
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<ToolCall>) {
|
||||
|
||||
(cleaned, tool_calls)
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// STREAMING CHUNK SPLITTING
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
|
||||
/// Remove special tokens and tool-call XML markup from a text fragment.
|
||||
///
|
||||
/// Handles both fixed tags (`<|im_end|>`, `<think>`, `<tool_call>`, …) and the
|
||||
/// attribute-bearing openers used by this model's tool format (`<function=…>`,
|
||||
/// `<parameter=…>`), 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|>",
|
||||
"<think>",
|
||||
"</think>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"</function>",
|
||||
"</parameter>",
|
||||
];
|
||||
if let Some(tag) = fixed.iter().find(|t| tail.starts_with(**t)) {
|
||||
rest = &tail[tag.len()..];
|
||||
continue;
|
||||
}
|
||||
|
||||
// Attribute-bearing openers: <function=…> / <parameter=…>.
|
||||
if let Some(attr) = tail
|
||||
.strip_prefix("<function=")
|
||||
.or_else(|| tail.strip_prefix("<parameter="))
|
||||
{
|
||||
if let Some(end) = attr.find('>') {
|
||||
rest = &attr[end + 1..];
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Not a known tag — keep this char and advance one UTF-8 char.
|
||||
let ch = tail.chars().next().expect("non-empty tail");
|
||||
out.push(ch);
|
||||
rest = &tail[ch.len_utf8()..];
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Split an incremental streamed text fragment into `(reasoning, content)`
|
||||
/// deltas for SSE.
|
||||
///
|
||||
/// * `think_done` means the `</think>` boundary was already crossed **before**
|
||||
/// this fragment (i.e. it is not the chunk containing the first `</think>`).
|
||||
/// * Whitespace **inside** a fragment is preserved — only the whitespace
|
||||
/// sitting immediately around the `</think>` 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 `</think>`, everything is emitted as `reasoning_content`;
|
||||
/// after it, as `content`. A stray second `</think>` is stripped, not split.
|
||||
pub fn split_stream_chunk(new_text: &str, think_done: bool) -> (Option<String>, String) {
|
||||
if !think_done {
|
||||
if let Some(pos) = new_text.find("</think>") {
|
||||
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</think>\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 </think> (already past the boundary) must not restart
|
||||
// reasoning classification.
|
||||
let (reasoning, content) = split_stream_chunk("...</think>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|><think>Hello</think>\n<tool_call><function=get_weather>",
|
||||
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("<think>Let me think\nabout it</think>\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<tool_call>\n<function=get_weather>\n<parameter=city>Jakarta</parameter>\n</function>\n</tool_call>";
|
||||
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("<tool_call>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_multiple_tool_calls() {
|
||||
let text = "<tool_call>\n<function=a>\n<parameter=x>1</parameter>\n</function>\n</tool_call>\n<tool_call>\n<function=b>\n<parameter=y>2</parameter>\n</function>\n</tool_call>";
|
||||
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("<s>"));
|
||||
assert!(prompt.contains("<|im_start|>system\nYou are helpful."));
|
||||
assert!(prompt.contains("<|im_start|>user\nHi!"));
|
||||
assert!(prompt.contains("<|im_start|>assistant\n<think>\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("<tools>"));
|
||||
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("<tool_response>").count(), 1);
|
||||
assert_eq!(prompt.matches("</tool_response>").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("<function=get_weather>"));
|
||||
assert!(prompt.contains("<parameter=city>Jakarta</parameter>"));
|
||||
}
|
||||
}
|
||||
|
||||
+4
-10
@@ -32,15 +32,12 @@ impl Application {
|
||||
pub async fn build() -> anyhow::Result<Self> {
|
||||
// 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,
|
||||
|
||||
+11
-3
@@ -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<T: std::str::FromStr>(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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -104,7 +104,7 @@ pub struct ResponseMessage {
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
}
|
||||
|
||||
#[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 `<tool_call>` 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<SseChoice>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
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)]
|
||||
|
||||
@@ -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<LlamaToken>,
|
||||
/// 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<CtxInner>,
|
||||
ctx: Mutex<CtxInner>,
|
||||
}
|
||||
|
||||
impl LlamaEngine {
|
||||
@@ -117,16 +129,16 @@ impl LlamaEngine {
|
||||
/// cannot be created.
|
||||
pub fn load() -> Result<Self, LlmError> {
|
||||
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())
|
||||
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());
|
||||
@@ -156,9 +168,12 @@ impl LlamaEngine {
|
||||
}
|
||||
|
||||
/// Tokenize a prompt string into tokens.
|
||||
///
|
||||
/// `AddBos::Never`: the chat template already prepends the `<s>` BOS token,
|
||||
/// so adding another here would produce a double BOS.
|
||||
pub fn tokenize(&self, prompt: &str) -> Result<Vec<LlamaToken>, 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<CtxInner> {
|
||||
&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 `<tool_call>` 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<LlamaToken>, String), LlmError> {
|
||||
let mut inner = self.ctx.lock().await;
|
||||
enable_tool_detection: bool,
|
||||
on_token: &mut dyn FnMut(LlamaToken, &str) -> bool,
|
||||
) -> Result<GenerationOutcome, LlmError> {
|
||||
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<LlamaToken> = 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("<tool_call>") {
|
||||
let close_count = text_buf.matches("</tool_call>").count();
|
||||
let open_count = text_buf.matches("<tool_call>").count();
|
||||
if open_count > 0 && close_count >= open_count {
|
||||
// Complete <tool_call> block emitted?
|
||||
if enable_tool_detection && text_buf.contains("<tool_call>") {
|
||||
let open = text_buf.matches("<tool_call>").count();
|
||||
let close = text_buf.matches("</tool_call>").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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
+151
-282
@@ -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<Arc<AppState>>,
|
||||
Json(req): Json<ChatRequest>,
|
||||
) -> Result<Response, AppError> {
|
||||
// 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)
|
||||
@@ -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 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 (output_tokens, raw_text) = state
|
||||
.engine
|
||||
.generate(&input_tokens, &mut sampler, max_tokens, &stop)
|
||||
.await?;
|
||||
|
||||
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 <tool_call> 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,310 +151,158 @@ 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::<Result<Event, Infallible>>(64);
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Role chunk
|
||||
let role_chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
// First chunk: announce the assistant role.
|
||||
let role_chunk = SseChunk::delta(
|
||||
chat_id.clone(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
model_name.clone(),
|
||||
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;
|
||||
}
|
||||
);
|
||||
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 </think> 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 {
|
||||
role: None,
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
finish_reason: Some("length".into()),
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||
break;
|
||||
}
|
||||
let outcome = engine.generate(
|
||||
&input_tokens,
|
||||
&mut sampler,
|
||||
max_tokens,
|
||||
&stop,
|
||||
has_tools,
|
||||
&mut |_token, piece| {
|
||||
text_buf.push_str(piece);
|
||||
|
||||
// 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 {
|
||||
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;
|
||||
}
|
||||
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_call>") {
|
||||
"tool_calls"
|
||||
} else {
|
||||
"stop"
|
||||
};
|
||||
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 piece = state.engine.decode_token(current);
|
||||
|
||||
// Push into buffer
|
||||
text_buf.push_str(&piece);
|
||||
|
||||
// Detect </think> transition
|
||||
// `was_thinking` is passed to split_stream_chunk so the chunk
|
||||
// containing the first </think> is treated as the boundary.
|
||||
let was_thinking = !think_done;
|
||||
if !think_done && text_buf.contains("</think>") {
|
||||
think_done = true;
|
||||
}
|
||||
|
||||
// Find new text since last send
|
||||
let new_text = &text_buf[sent_len..]; // everything not yet streamed
|
||||
let new_text = &text_buf[sent_len..];
|
||||
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;
|
||||
return true;
|
||||
}
|
||||
|
||||
// Strip special tokens from the NEW text chunk
|
||||
// Split at </think> if present — before goes to reasoning, after to content
|
||||
let (reasoning_part, content_part) = if let Some(pos) = new_text.find("</think>") {
|
||||
let before = new_text[..pos]
|
||||
.replace("<|im_end|>", "")
|
||||
.replace("<|im_start|>", "")
|
||||
.replace("<think>", "")
|
||||
.trim()
|
||||
.to_string();
|
||||
let after = new_text[pos + 8..]
|
||||
.replace("<|im_end|>", "")
|
||||
.replace("<|im_start|>", "")
|
||||
.replace("<think>", "")
|
||||
.trim()
|
||||
.to_string();
|
||||
think_done = true;
|
||||
(Some(before), after)
|
||||
} else if think_done {
|
||||
let cleaned = new_text
|
||||
.replace("<|im_end|>", "")
|
||||
.replace("<|im_start|>", "")
|
||||
.replace("<think>", "")
|
||||
.trim()
|
||||
.to_string();
|
||||
(None, cleaned)
|
||||
} else {
|
||||
let cleaned = new_text
|
||||
.replace("<|im_end|>", "")
|
||||
.replace("<|im_start|>", "")
|
||||
.replace("<think>", "")
|
||||
.trim()
|
||||
.to_string();
|
||||
(Some(cleaned), String::new())
|
||||
};
|
||||
let (reasoning, content) = chat::split_stream_chunk(new_text, was_thinking);
|
||||
|
||||
// Send reasoning part (before </think>, or entire text if still thinking)
|
||||
if let Some(ref r) = reasoning_part {
|
||||
if !r.is_empty() {
|
||||
let delta = SseDelta {
|
||||
if let Some(reasoning) = reasoning {
|
||||
let chunk = SseChunk::delta(
|
||||
chat_id.clone(),
|
||||
created,
|
||||
model_name.clone(),
|
||||
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;
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// Send content part (after </think>, 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(),
|
||||
if !content.is_empty() {
|
||||
let chunk = SseChunk::delta(
|
||||
chat_id.clone(),
|
||||
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;
|
||||
model_name.clone(),
|
||||
SseDelta {
|
||||
role: None,
|
||||
content: Some(content),
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
);
|
||||
let event = serde_json::to_string(&chunk).unwrap();
|
||||
if tx.blocking_send(Ok(Event::default().data(event))).is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
sent_len = text_buf.len();
|
||||
true
|
||||
},
|
||||
);
|
||||
|
||||
// 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(),
|
||||
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: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
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;
|
||||
}
|
||||
|
||||
// Check tool call completeness
|
||||
if has_tools && text_buf.contains("<tool_call>") {
|
||||
let open = text_buf.matches("<tool_call>").count();
|
||||
let close = text_buf.matches("</tool_call>").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 event = serde_json::to_string(&chunk).unwrap();
|
||||
let _ = tx.blocking_send(Ok(Event::default().data(event)));
|
||||
}
|
||||
}
|
||||
|
||||
let pos = input_tokens.len() as i32 + count as i32;
|
||||
if let Err(e) = inner.decode(current, pos) {
|
||||
info!(" Decode error: {e}");
|
||||
break;
|
||||
let finish_reason = match (outcome.finish, sent_tool_calls) {
|
||||
(FinishReason::ToolCalls, true) => "tool_calls",
|
||||
(FinishReason::ToolCalls, false) => "stop",
|
||||
(f, _) => f.as_str(),
|
||||
};
|
||||
|
||||
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)));
|
||||
}
|
||||
count += 1;
|
||||
current = inner.sample(&mut sampler);
|
||||
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())));
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAI-compatible terminator.
|
||||
let _ = tx.blocking_send(Ok(Event::default().data("[DONE]")));
|
||||
});
|
||||
|
||||
let stream = ReceiverStream::new(rx);
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -8,6 +8,7 @@ use crate::domain::entity::HealthResponse;
|
||||
pub async fn health_check() -> Json<HealthResponse> {
|
||||
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(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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<Response, (StatusCode, String)> {
|
||||
pub async fn auth_middleware(request: Request, next: Next) -> Result<Response, AppError> {
|
||||
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<Response, (
|
||||
return Ok(next.run(request).await);
|
||||
}
|
||||
|
||||
Err((
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"{\"error\":\"unauthorized\",\"message\":\"Invalid API key\"}".into(),
|
||||
))
|
||||
// AppError renders a JSON body with the correct Content-Type.
|
||||
Err(AppError::Unauthorized)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user