first
This commit is contained in:
valenti
+451
@@ -0,0 +1,451 @@
|
||||
//! Minimal OpenAI-compatible LLM client (chat completions).
|
||||
//!
|
||||
//! This is the "reasoning" stage of the speech-to-speech pipeline: the user
|
||||
//! transcript (from STT or typed in the chat box) is turned into a reply that
|
||||
//! the TTS engine then speaks. It works with any OpenAI-compatible server
|
||||
//! (vLLM, SGLang, llama.cpp, Ollama, OpenAI itself) thanks to the `/v1`
|
||||
//! convention: `POST {LLM_BASE_URL}/chat/completions`.
|
||||
//!
|
||||
//! Reasoning models (e.g. Qwen3) may leak their chain-of-thought either in a
|
||||
//! separate `reasoning_content` delta field (vLLM) or inline between think
|
||||
//! markers (llama.cpp / other servers), or — when the server runs no reasoning
|
||||
//! parser at all — as plain narration inside `content`. None of it should be
|
||||
//! displayed in the chat box nor read aloud, so three independent mechanisms
|
||||
//! are available, each switched from the environment:
|
||||
//! - `LLM_ENABLE_THINKING` — send `chat_template_kwargs.enable_thinking`
|
||||
//! (`false` turns the thinking mode off server-side; unset omits the field).
|
||||
//! - `LLM_STRIP_REASONING` — drop `reasoning_content` / inline think blocks.
|
||||
//! - `LLM_DIRECT_ANSWERS` — strengthen the system prompt (see
|
||||
//! [`crate::config::DIRECT_ANSWER_DIRECTIVE`]).
|
||||
//!
|
||||
//! Two entry points are provided, selectable at runtime via `LLM_STREAM`:
|
||||
//! - [`LlmClient::chat_stream`] — SSE token streaming (lowest time-to-first-audio).
|
||||
//! - [`LlmClient::chat`] — single blocking request (simplest, no partial output).
|
||||
|
||||
use crate::config::Config;
|
||||
use anyhow::{anyhow, Result};
|
||||
use futures_util::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
/// One entry of the OpenAI `messages` array.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
Self { role: "system".into(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn user(content: impl Into<String>) -> Self {
|
||||
Self { role: "user".into(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn assistant(content: impl Into<String>) -> Self {
|
||||
Self { role: "assistant".into(), content: content.into() }
|
||||
}
|
||||
}
|
||||
|
||||
/// Client for an OpenAI-compatible `/chat/completions` endpoint.
|
||||
#[derive(Clone)]
|
||||
pub struct LlmClient {
|
||||
http: Client,
|
||||
/// Fully-qualified chat-completions URL (`{base}/chat/completions`).
|
||||
endpoint: String,
|
||||
model: String,
|
||||
api_key: String,
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
/// `Some(v)` sends `chat_template_kwargs: {"enable_thinking": v}` — the
|
||||
/// vLLM/SGLang switch for Qwen3-style hybrid reasoning models. `None`
|
||||
/// omits the field so strict OpenAI-compatible servers keep working.
|
||||
enable_thinking: Option<bool>,
|
||||
/// Strip the chain-of-thought (`reasoning_content` / inline think blocks)
|
||||
/// from the reply before it is spoken and shown in the chat box.
|
||||
strip_reasoning: bool,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
pub fn new(
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
model: &str,
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
enable_thinking: Option<bool>,
|
||||
strip_reasoning: bool,
|
||||
) -> Self {
|
||||
let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/'));
|
||||
// No global request timeout: an SSE stream can legitimately stay open
|
||||
// for a while. Only the TCP connect is bounded.
|
||||
let http = Client::builder()
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
http,
|
||||
endpoint,
|
||||
model: model.to_string(),
|
||||
api_key: api_key.to_string(),
|
||||
max_tokens,
|
||||
temperature,
|
||||
enable_thinking,
|
||||
strip_reasoning,
|
||||
}
|
||||
}
|
||||
|
||||
/// Endpoint URL, for logging.
|
||||
pub fn endpoint(&self) -> &str {
|
||||
&self.endpoint
|
||||
}
|
||||
|
||||
/// Model identifier, for logging.
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
/// Requested thinking mode, for logging (`None` = field not sent).
|
||||
pub fn enable_thinking(&self) -> Option<bool> {
|
||||
self.enable_thinking
|
||||
}
|
||||
|
||||
/// Whether chain-of-thought stripping is active, for logging.
|
||||
pub fn strip_reasoning(&self) -> bool {
|
||||
self.strip_reasoning
|
||||
}
|
||||
|
||||
fn request_body(&self, messages: &[ChatMessage], stream: bool) -> serde_json::Value {
|
||||
let mut body = json!({
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"max_tokens": self.max_tokens,
|
||||
"temperature": self.temperature,
|
||||
"stream": stream,
|
||||
});
|
||||
if let Some(enable) = self.enable_thinking {
|
||||
body["chat_template_kwargs"] = json!({ "enable_thinking": enable });
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
/// The text a delta carries: the assistant `content` plus — only when
|
||||
/// stripping is off — the separate `reasoning_content` field. With
|
||||
/// stripping on, `reasoning_content` is dropped whole.
|
||||
fn delta_text(&self, delta: &serde_json::Value) -> String {
|
||||
let content = delta["content"].as_str().unwrap_or("");
|
||||
if self.strip_reasoning {
|
||||
return content.to_string();
|
||||
}
|
||||
let reasoning = delta["reasoning_content"].as_str().unwrap_or("");
|
||||
let mut out = String::with_capacity(content.len() + reasoning.len());
|
||||
out.push_str(content);
|
||||
out.push_str(reasoning);
|
||||
out
|
||||
}
|
||||
|
||||
/// Apply (or skip) chain-of-thought stripping to a complete reply.
|
||||
fn clean(&self, text: &str) -> String {
|
||||
if self.strip_reasoning {
|
||||
strip_think(text)
|
||||
} else {
|
||||
text.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Single blocking request: returns the full reply once complete.
|
||||
pub async fn chat(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<String> {
|
||||
let body = self.request_body(messages, false);
|
||||
let req = self
|
||||
.http
|
||||
.post(&self.endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.json(&body)
|
||||
.send();
|
||||
|
||||
let resp = tokio::select! {
|
||||
r = req => r?,
|
||||
_ = cancel.cancelled() => return Ok(String::new()),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("LLM HTTP {}: {}", status, text.trim()));
|
||||
}
|
||||
let value: serde_json::Value = resp.json().await?;
|
||||
let message = &value["choices"][0]["message"];
|
||||
Ok(self.clean(&self.delta_text(message)))
|
||||
}
|
||||
|
||||
/// Streaming request. Every clean (think-stripped) text delta is pushed to
|
||||
/// `delta_tx` as it arrives; the full clean reply is returned at the end.
|
||||
///
|
||||
/// When `cancel` fires, the stream is abandoned and the partial reply
|
||||
/// returned — used for barge-in.
|
||||
pub async fn chat_stream(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
cancel: &CancellationToken,
|
||||
delta_tx: mpsc::UnboundedSender<String>,
|
||||
) -> Result<String> {
|
||||
let body = self.request_body(messages, true);
|
||||
let resp = self
|
||||
.http
|
||||
.post(&self.endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("LLM HTTP {}: {}", status, text.trim()));
|
||||
}
|
||||
|
||||
let mut stream = resp.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
// Only allocate the stripper when chain-of-thought filtering is on.
|
||||
let mut strip = self.strip_reasoning.then(ThinkStripper::new);
|
||||
let mut full = String::new();
|
||||
let mut done = false;
|
||||
|
||||
'outer: while !done {
|
||||
// Race the next chunk against barge-in: a cancelled turn drops the
|
||||
// HTTP response immediately instead of waiting for the next token.
|
||||
let chunk = tokio::select! {
|
||||
next = stream.next() => match next {
|
||||
Some(Ok(c)) => c,
|
||||
Some(Err(e)) => {
|
||||
warn!(error = %e, "LLM stream read error");
|
||||
break;
|
||||
}
|
||||
None => break,
|
||||
},
|
||||
_ = cancel.cancelled() => break,
|
||||
};
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// SSE frames are newline-delimited; keep the trailing partial line.
|
||||
while let Some(pos) = buffer.find('\n') {
|
||||
let line: String = buffer.drain(..=pos).collect();
|
||||
let line = line.trim();
|
||||
let Some(data) = line.strip_prefix("data:") else {
|
||||
continue;
|
||||
};
|
||||
let data = data.trim();
|
||||
if data == "[DONE]" {
|
||||
done = true;
|
||||
break;
|
||||
}
|
||||
let Ok(value) = serde_json::from_str::<serde_json::Value>(data) else {
|
||||
continue;
|
||||
};
|
||||
// A separate `reasoning_content` delta (vLLM) is dropped unless
|
||||
// stripping is disabled — see `delta_text`.
|
||||
let content = self.delta_text(&value["choices"][0]["delta"]);
|
||||
if content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let clean = match strip.as_mut() {
|
||||
Some(s) => s.push(&content),
|
||||
None => content,
|
||||
};
|
||||
if !clean.is_empty() {
|
||||
full.push_str(&clean);
|
||||
if delta_tx.send(clean).is_err() {
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(s) = strip.as_mut() {
|
||||
let tail = s.finish();
|
||||
if !tail.is_empty() {
|
||||
full.push_str(&tail);
|
||||
let _ = delta_tx.send(tail);
|
||||
}
|
||||
}
|
||||
Ok(full)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the LLM client from the environment, or `None` when disabled /
|
||||
/// misconfigured (`LLM_ENABLED=false`, or no base URL / model).
|
||||
pub fn create_llm_client(config: &Config) -> Option<Arc<LlmClient>> {
|
||||
if !config.llm_enabled {
|
||||
return None;
|
||||
}
|
||||
let base = config.llm_base_url.trim();
|
||||
let model = config.llm_model.trim();
|
||||
if base.is_empty() || model.is_empty() {
|
||||
warn!("LLM enabled but LLM_BASE_URL / LLM_MODEL is empty — LLM disabled");
|
||||
return None;
|
||||
}
|
||||
Some(Arc::new(LlmClient::new(
|
||||
base,
|
||||
&config.llm_api_key,
|
||||
model,
|
||||
config.llm_max_tokens,
|
||||
config.llm_temperature,
|
||||
config.llm_enable_thinking,
|
||||
config.llm_strip_reasoning,
|
||||
)))
|
||||
}
|
||||
|
||||
// --- Chain-of-thought stripping --------------------------------------------
|
||||
|
||||
// Opening / closing chain-of-thought markers emitted by Qwen3-style reasoning
|
||||
// models. Written with `\u{..}` escapes so the literal tags never appear in
|
||||
// this source file as markup.
|
||||
const THINK_OPEN: &str = "\u{3C}think\u{3E}";
|
||||
const THINK_CLOSE: &str = "\u{3C}/think\u{3E}";
|
||||
|
||||
/// Remove every ` thinking…<｜end▁of▁thinking｜>` block from `text`.
|
||||
///
|
||||
/// An unterminated opening tag drops the remainder of the text (the model was
|
||||
/// still reasoning when the stream ended).
|
||||
fn strip_think(text: &str) -> String {
|
||||
let mut out = String::with_capacity(text.len());
|
||||
let mut rest = text;
|
||||
loop {
|
||||
match rest.find(THINK_OPEN) {
|
||||
Some(i) => {
|
||||
out.push_str(&rest[..i]);
|
||||
let after = &rest[i + THINK_OPEN.len()..];
|
||||
match after.find(THINK_CLOSE) {
|
||||
Some(j) => rest = &after[j + THINK_CLOSE.len()..],
|
||||
None => return out,
|
||||
}
|
||||
}
|
||||
None => {
|
||||
out.push_str(rest);
|
||||
return out;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Incremental think-stripper for a stream of deltas.
|
||||
///
|
||||
/// Recomputes the clean text from the accumulated raw text on every delta, so
|
||||
/// a ` thinking` block that spans several deltas is removed as a whole. A short
|
||||
/// tail is held back because the last few characters may belong to an opening
|
||||
/// tag split across deltas.
|
||||
struct ThinkStripper {
|
||||
raw: String,
|
||||
/// Bytes of clean text already returned to the caller.
|
||||
emitted: usize,
|
||||
}
|
||||
|
||||
impl ThinkStripper {
|
||||
fn new() -> Self {
|
||||
Self { raw: String::new(), emitted: 0 }
|
||||
}
|
||||
|
||||
fn push(&mut self, delta: &str) -> String {
|
||||
self.raw.push_str(delta);
|
||||
let clean = strip_think(&self.raw);
|
||||
let safe = floor_char_boundary(&clean, clean.len().saturating_sub(THINK_OPEN.len()));
|
||||
if safe <= self.emitted {
|
||||
return String::new();
|
||||
}
|
||||
let out = clean[self.emitted..safe].to_string();
|
||||
self.emitted = safe;
|
||||
out
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> String {
|
||||
let clean = strip_think(&self.raw);
|
||||
if clean.len() <= self.emitted {
|
||||
return String::new();
|
||||
}
|
||||
let out = clean[self.emitted..].to_string();
|
||||
self.emitted = clean.len();
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// Largest char boundary `<= i`.
|
||||
fn floor_char_boundary(s: &str, mut i: usize) -> usize {
|
||||
if i >= s.len() {
|
||||
return s.len();
|
||||
}
|
||||
while i > 0 && !s.is_char_boundary(i) {
|
||||
i -= 1;
|
||||
}
|
||||
i
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn strips_closed_think_block() {
|
||||
let text = format!("a{}b c{}d", THINK_OPEN, THINK_CLOSE);
|
||||
assert_eq!(strip_think(&text), "ad");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_unterminated_think_block() {
|
||||
// No opening tag yet: the partial text is kept verbatim.
|
||||
assert_eq!(strip_think("hello"), "hello");
|
||||
// Opening tag without a closing tag: everything after it is dropped.
|
||||
let text = format!("hello{}still going", THINK_OPEN);
|
||||
assert_eq!(strip_think(&text), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_body_thinking_is_opt_in() {
|
||||
let msgs = vec![ChatMessage::user("hi")];
|
||||
// Unset: the vendor-specific field must not be sent.
|
||||
let plain = LlmClient::new("http://x/v1", "k", "m", 16, 0.5, None, true);
|
||||
assert!(plain.request_body(&msgs, true).get("chat_template_kwargs").is_none());
|
||||
// Explicit: forwarded verbatim.
|
||||
let off = LlmClient::new("http://x/v1", "k", "m", 16, 0.5, Some(false), true);
|
||||
let body = off.request_body(&msgs, false);
|
||||
assert_eq!(body["chat_template_kwargs"]["enable_thinking"], false);
|
||||
assert_eq!(body["stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stripper_is_stable_across_deltas() {
|
||||
// Split both markers in the middle to exercise the tail-holdback logic.
|
||||
let open = THINK_OPEN;
|
||||
let close = THINK_CLOSE;
|
||||
let parts = vec![
|
||||
"Ho".to_string(),
|
||||
"la ".to_string(),
|
||||
open[..4].to_string(),
|
||||
open[4..].to_string(),
|
||||
"plan".to_string(),
|
||||
"ning".to_string(),
|
||||
close[..4].to_string(),
|
||||
close[4..].to_string(),
|
||||
" answer".to_string(),
|
||||
" done.".to_string(),
|
||||
];
|
||||
let mut s = ThinkStripper::new();
|
||||
let mut out = String::new();
|
||||
for part in &parts {
|
||||
out.push_str(&s.push(part));
|
||||
}
|
||||
out.push_str(&s.finish());
|
||||
assert_eq!(out, "Hola answer done.");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user