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

    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", %payload, "fragmento SSE ilegible");
                    return true;
                };
                let Some(choice) = event.get("choices").and_then(|c| c.get(0)) else {
                    return true;
                };
                let delta = choice.get("delta").unwrap_or(&Value::Null);

                if let Some(calls) = delta.get("tool_calls").and_then(Value::as_array) {
                    assembler.absorb(calls);
                }
                if let Some(text) = delta.get("content").and_then(Value::as_str) {
                    if !text.is_empty() {
                        if outcome.ttft.is_none() {
                            outcome.ttft = Some(started.elapsed());
                        }
                        outcome.text.push_str(text);
                        if !on_delta(&Delta::Text(text.to_string())) {
                            outcome.stopped_early = true;
                            return false;
                        }
                    }
                }
                // Un `finish_reason` cierra la respuesta aunque el servidor no
                // llegue a mandar el [DONE] (pasa al cortar la conexión).
                if choice.get("finish_reason").is_some_and(|r| !r.is_null()) {
                    finished = true;
                    return false;
                }
                true
            });

        match result {
            Ok(()) => {}
            // Cortar a propósito no es un fallo: `on_delta` devolvió false o
            // se activó la cancelación.
            Err(Error::Cancelled) if finished || outcome.stopped_early || cancel.is_cancelled() => {
            }
            Err(err) => return Err(err),
        }

        outcome.tool_calls = assembler.finish();
        if !outcome.tool_calls.is_empty()
            && !on_delta(&Delta::ToolCalls(outcome.tool_calls.clone()))
        {
            outcome.stopped_early = true;
        }
        Ok(outcome)
    }
}

/// Reensambla las llamadas a herramientas, que llegan repartidas en fragmentos.
///
/// El streaming manda el nombre en un fragmento y los argumentos en varios
/// más, identificados sólo por su `index`; hasta que no termina el flujo no
/// hay una llamada completa que ejecutar.
#[derive(Default)]
struct ToolCallAssembler {
    partial: Vec<PartialCall>,
}

#[derive(Default, Clone)]
struct PartialCall {
    id: String,
    name: String,
    arguments: String,
}

impl ToolCallAssembler {
    fn absorb(&mut self, calls: &[Value]) {
        for call in calls {
            let index = call.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
            if self.partial.len() <= index {
                self.partial.resize(index + 1, PartialCall::default());
            }
            let slot = &mut self.partial[index];
            if let Some(id) = call.get("id").and_then(Value::as_str) {
                if !id.is_empty() {
                    slot.id = id.to_string();
                }
            }
            let Some(function) = call.get("function") else {
                continue;
            };
            if let Some(name) = function.get("name").and_then(Value::as_str) {
                if !name.is_empty() {
                    slot.name = name.to_string();
                }
            }
            if let Some(args) = function.get("arguments").and_then(Value::as_str) {
                slot.arguments.push_str(args);
            }
        }
    }

    fn finish(self) -> Vec<ToolCall> {
        self.partial
            .into_iter()
            .enumerate()
            .filter(|(_, call)| !call.name.is_empty())
            .map(|(i, ca