piperust/src/llm.rs

//! 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 crate::tools::ToolRegistry;
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() }
    }

    /// A tool result handed back to the model after a tool call.
    pub fn tool(content: impl Into<String>) -> Self {
        Self { role: "tool".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,
    /// Optional tool registry. When set, every request advertises the tools and
    /// the orchestrator runs the call → execute → re-prompt loop. See
    /// [`crate::tools`].
    tools: Option<Arc<ToolRegistry>>,
}

impl LlmClient {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        base_url: &str,
        api_key: &str,
        model: &str,
        max_tokens: u32,
        temperature: f32,
        enable_thinking: Option<bool>,
        strip_reasoning: bool,
        tools: Option<Arc<ToolRegistry>>,
    ) -> 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,
            tools,
        }
    }

    /// 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
    }

    /// The tool registry, if tool calling is configured. Used by the
    /// orchestrator to execute parsed calls.
    pub fn tools(&self) -> Option<&Arc<ToolRegistry>> {
        self.tools.as_ref()
    }

    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 });
        }
        // Advertise the configured tools so the model can call them. The
        // server has no tool parser, so the call comes back as text in the
        // native `<tool_call>` format — parsed client-side in `crate::tools`.
        if let Some(reg) = &self.tools {
            if !reg.is_empty() {
                body["tools"] = reg.openai_tools();
                body["tool_choice"] = json!("auto");
            }
        }
        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)))
    }

    /// One-off non-streaming completion with an explicit system/user pair, its
    /// own token budget and temperature, and **no tools advertised**.
    ///
    /// Used by auxiliary calls that must not go through the conversational
    /// prompt or the tool loop — language detection and phrase-list generation
    /// (see [`crate::lang`]).
    pub async fn simple_chat(
        &self,
        system: &str,
        user: &str,
        max_tokens: u32,
        temperature: f32,
        cancel: &CancellationToken,
    ) -> Result<String> {
        let mut body = json!({
            "model": self.model,
            "messages": [ChatMessage::system(system), ChatMessage::user(user)],
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": false,
        });
        if let Some(enable) = self.enable_thinking {
            body["chat_template_kwargs"] = json!({ "enable_thinking": enable });
        }
        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,
        load_tools(config),
    )))
}

/// Load the tool registry from `LLM_TOOLS_FILE` when tool calling is enabled.
/// A missing or invalid file is logged and tool calling is disabled — the
/// pipeline keeps working without tools.
fn load_tools(config: &Config) -> Option<Arc<ToolRegistry>> {
    if !config.llm_tools_enabled || config.llm_tools_file.trim().is_empty() {
        return None;
    }
    let path = std::path::Path::new(config.llm_tools_file.trim());
    match ToolRegistry::load(path, config.llm_tool_timeout_secs, config.llm_tool_result_max_chars) {
        Ok(reg) if reg.is_empty() => {
            warn!("tools file present but empty — tool calling disabled");
            None
        }
        Ok(reg) => Some(Arc::new(reg)),
        Err(e) => {
            warn!(error = %e, "failed to load tools file — tool calling disabled");
            None
        }
    }
}

// --- 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, None);
        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, None);
        let body = off.request_body(&msgs, false);
        assert_eq!(body["chat_template_kwargs"]["enable_thinking"], false);
        assert_eq!(body["stream"], false);
        // Without a registry no `tools` field is advertised.
        assert!(body.get("tools").is_none());
    }

    #[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.");
    }
}