path: root/crates/asist-llm/src/client.rs

//! Transporte hacia llama-server.

use std::time::{Duration, Instant};

use serde_json::{json, Value};

use asist_core::config::LlmConfig;
use asist_core::error::{Error, Result};
use asist_core::http::{Cancel, HttpClient};
use asist_core::tools::{ToolCall, ToolRegistry};

use crate::chat::Conversation;

/// Un trozo de respuesta según sale del modelo.
#[derive(Debug, Clone)]
pub enum Delta {
    /// Texto para el usuario.
    Text(String),
    /// El modelo ha decidido llamar a herramientas; no habrá más texto.
    ToolCalls(Vec<ToolCall>),
}

/// Cómo terminó una respuesta en streaming.
#[derive(Debug, Clone, Default)]
pub struct StreamOutcome {
    pub text: String,
    pub tool_calls: Vec<ToolCall>,
    /// Tiempo hasta el primer fragmento con contenido.
    pub ttft: Option<Duration>,
    pub stopped_early: bool,
}

pub struct LlmClient {
    http: HttpClient,
    config: LlmConfig,
}

impl LlmClient {
    pub fn new(authority: String, config: &LlmConfig) -> Self {
        Self {
            http: HttpClient::new(authority)
                .with_read_timeout(Duration::from_secs(config.request_timeout_secs)),
            config: config.clone(),
        }
    }

    pub fn healthy(&self) -> bool {
        self.http.healthy("/health")
    }

    /// Espera a que el servidor termine de cargar el modelo.
    pub fn wait_ready(&self, timeout: Duration) -> Result<()> {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if self.healthy() {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(500));
        }
        Err(Error::Llm(format!(
            "{} no respondió a /health en {} s",
            self.http.authority,
            timeout.as_secs()
        )))
    }

    /// Comprueba que la plantilla de chat no arranca un bloque de
    /// razonamiento sin cerrar.
    ///
    /// Merece la pena avisar: con la plantilla original de Qwen3.5 el modelo
    /// se pasa entre 7 y 9 segundos «pensando» antes de la primera palabra
    /// audible, y desde fuera parece que el asistente se ha colgado.
    pub fn warn_if_thinking_template(&self) {
        let Ok(body) = self.http.get("/props") else {
            return;
        };
        let Ok(props) = serde_json::from_slice::<Value>(&body) else {
            return;
        };
        let Some(template) = props.get("chat_template").and_then(Value::as_str) else {
            return;
        };

        let opens = template.matches("<think>").count();
        let closes = template.matches("</think>").count();
        if opens > closes {
            tracing::warn!(
                target: "llm",
                "la plantilla de chat abre <think> sin cerrarlo: el modelo razonará \
                 varios segundos antes de contestar. Arranca llama-server con \
                 --chat-template-file config/qwen35-no-think.jinja"
            );
        }
    }

    /// Pregunta al modelo por una imagen, en una petición aparte de la
    /// conversación.
    ///
    /// El servidor tiene cargado el proyector multimodal (`--mmproj`), así que
    /// acepta partes `image_url` con la imagen en base64. Se hace fuera del
    /// historial a propósito: una imagen ocupa cientos de tokens de contexto y
    /// arrastrarla turno tras turno saldría carísimo para lo poco que aporta
    /// una vez descrita. Lo que vuelve a la conversación es el texto.
    ///
    /// Medido en esta máquina, el coste depende mucho del tamaño: 1,3 s a
    /// 320x240, 2,9 s a 640x480 y 7,8 s a 1280x720.
    pub fn look(&self, image_jpeg: &[u8], question: &str, cancel: &Cancel) -> Result<String> {
        use base64::Engine;
        let encoded = base64::engine::general_purpose::STANDARD.encode(image_jpeg);

        let mut body = json!({
            "messages": [{
                "role": "user",
                "content": [
                    { "type": "text", "text": question },
                    {
                        "type": "image_url",
                        "image_url": { "url": format!("data:image/jpeg;base64,{encoded}") }
                    }
                ]
            }],
            "stream": true,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
        });
        if !self.config.model.is_empty() {
            body["model"] = json!(self.config.model);
        }

        // En streaming aunque no se use el texto parcial: una petición de
        // visión tarda segundos y sin flujo el socket puede quedarse callado
        // hasta pasado el plazo de lectura.
        let started = Instant::now();
        let mut text = String::new();
        let mut finished = false;
        let result = self
            .http
            .post_json_lines("/v1/chat/completions", &body, cancel, |line| {
                let Some(payload) = line.strip_prefix("data: ") else {
                    return true;
                };
                if payload.trim() == "[DONE]" {
                    finished = true;
                    return false;
                }
                let Ok(event) = serde_json::from_str::<Value>(payload) else {
                    return true;
                };
                let Some(choice) = event.get("choices").and_then(|c| c.get(0)) else {
                    return true;
                };
                if let Some(chunk) = choice
                    .get("delta")
                    .and_then(|d| d.get("content"))
                    .and_then(Value::as_str)
                {
                    text.push_str(chunk);
                }
                if choice.get("finish_reason").is_some_and(|r| !r.is_null()) {
                    finished = true;
                    return false;
                }
                true
            });
        match result {
            Ok(()) => {}
            Err(Error::Cancelled) if finished || cancel.is_cancelled() => {}
            Err(err) => return Err(err),
        }

        tracing::debug!(
            target: "llm",
            bytes = image_jpeg.len(),
            ms = started.elapsed().as_millis(),
            "visión"
        );
        let text = text.trim().to_string();
        if text.is_empty() {
            return Err(Error::Llm("el modelo no describió la imagen".into()));
        }
        Ok(text)
    }

    fn request_body(&self, chat: &Conversation, tools: Option<&ToolRegistry>) -> Value {
        let mut body = json!({
            "messages": chat.to_json(),
            "stream": true,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "top_k": self.config.top_k,
            "min_p": self.config.min_p,
            "repeat_penalty": self.config.repeat_penalty,
            "max_tokens": self.config.max_tokens,
        });
        if !self.config.model.is_empty() {
            body["model"] = json!(self.config.model);
        }
        if let Some(tools) = tools.filter(|t| !t.is_empty()) {
            body["tools"] = tools.schema();
            body["tool_choice"] = json!("auto");
        }
        body
    }

    /// Envía la conversación y entrega la respuesta a trozos.
    ///
    /// `on_delta` devuelve `false` para cortar —lo que hace el orquestador
    /// cuando el usuario interrumpe—; `cancel` hace lo mismo desde otro hilo.
    pub fn stream(
        &self,
        chat: &Conversation,
        tools: Option<&ToolRegistry>,
        cancel: &Cancel,
        mut on_delta: impl FnMut(&Delta) -> bool,
    ) -> Result<StreamOutcome> {
        let body = self.request_body(chat, tools);
        let started = Instant::now();
        let mut outcome = StreamOutcome::default();
        let mut assembler = ToolCallAssembler::default();
        let mut finished = false;

        let result = self
            .http
            .post_json_lines("/v1/chat/completions", &body, cancel, |line| {
                let Some(payload) = line.strip_prefix("data: ") else {
                    return true;
                };
                if payload.trim() == "[DONE]" {
                    finished = true;
                    return false;
                }
                let Ok(event) = serde_json::from_str::<Value>(payload) else {
                    tracing::debug!(target: "llm"