aura-gui 0.2.2

//! Autocompletado FIM («fill in the middle»): se envía el texto antes y después del cursor a un modelo hecho para
//! rellenar código, sin chat ni razonamiento. Un solo adaptador para Codestral (Mistral), Mercury Edit (Inception)
//! y modelos locales con Ollama: todos reciben `{model, prompt, suffix, max_tokens}` y solo cambia dónde ponen el
//! texto en la respuesta.

use crate::ai::http;
use crate::ai::provider::Decoder;
use crate::ai::sse::SseEvent;
use crate::ai::types::{AiError, Event};
use crate::core::CancelToken;
use serde_json::{Value, json};

/// Servicio FIM y cómo se le pregunta.
#[derive(Clone, Debug, PartialEq)]
pub struct FimProvider {
    /// Dirección completa del endpoint.
    pub url: String,
    pub api_key: Option<String>,
    pub model: String,
    /// `false` en los servicios que no documentan streaming: responden de una vez.
    pub stream: bool,
}

/// Proveedores de predicción que se piden por FIM, con su modelo por defecto.
pub const FIM_PROVIDERS: &[(&str, &str)] =
    &[("codestral", "codestral-latest"), ("mercury", "mercury-edit-2"), ("ollama", "qwen2.5-coder:1.5b")];

/// Modelos que se ofrecen en Ajustes para cada proveedor FIM (se puede escribir otro).
pub fn suggested_models(provider: &str) -> &'static [&'static str] {
    match provider {
        "codestral" => &["codestral-latest", "codestral-2508"],
        "mercury" => &["mercury-edit-2"],
        "ollama" => &["qwen2.5-coder:1.5b", "qwen2.5-coder:0.5b", "qwen2.5-coder:3b", "starcoder2:3b"],
        _ => &[],
    }
}

/// Modelo por defecto de un proveedor FIM; `None` si el proveedor no es FIM.
pub fn default_model(provider: &str) -> Option<&'static str> {
    FIM_PROVIDERS.iter().find(|(id, _)| *id == provider).map(|(_, model)| *model)
}

impl FimProvider {
    /// Codestral con la clave gratuita para editores (endpoint `codestral.mistral.ai`).
    pub fn codestral(api_key: String) -> Self {
        let url = "https://codestral.mistral.ai/v1/fim/completions".into();
        Self { url, api_key: Some(api_key), model: "codestral-latest".into(), stream: true }
    }

    /// Mercury Edit de Inception: muy rápido; de pago por uso.
    pub fn mercury(api_key: String) -> Self {
        let url = "https://api.inceptionlabs.ai/v1/fim/completions".into();
        Self { url, api_key: Some(api_key), model: "mercury-edit-2".into(), stream: false }
    }

    /// Modelo local servido por Ollama (`suffix` en su API compatible con OpenAI); `base` sin `/v1`.
    pub fn ollama(base: Option<&str>) -> Self {
        let base = base.unwrap_or("http://127.0.0.1:11434").trim_end_matches('/').trim_end_matches("/v1");
        Self { url: format!("{base}/v1/completions"), api_key: None, model: "qwen2.5-coder:1.5b".into(), stream: true }
    }

    pub fn body(&self, prefix: &str, suffix: &str, max_tokens: u32) -> Value {
        json!({
            "model": self.model,
            "prompt": prefix,
            "suffix": suffix,
            "max_tokens": max_tokens,
            "temperature": 0,
            "stream": self.stream,
        })
    }

    fn headers(&self) -> Vec<(String, String)> {
        self.api_key.iter().map(|key| ("authorization".to_owned(), format!("Bearer {key}"))).collect()
    }

    /// Pide el relleno entre `prefix` y `suffix`; el texto llega a `emit` según se genera (o de una vez).
    pub fn complete(
        &self,
        prefix: &str,
        suffix: &str,
        max_tokens: u32,
        cancel: &CancelToken,
        emit: &mut dyn FnMut(Event),
    ) -> Result<(), AiError> {
        let body = self.body(prefix, suffix, max_tokens);
        if self.stream {
            let policy = http::RetryPolicy { retries: 0, ..Default::default() };
            return http::stream_sse_with(
                policy,
                &self.url,
                &self.headers(),
                &body,
                cancel,
                &mut FimDecoder::default(),
                emit,
            );
        }
        let response = http::post_json(&self.url, &self.headers(), &body, cancel)?;
        if let Some(text) = text_of(&response).filter(|text| !text.is_empty()) {
            emit(Event::Text(text.to_owned()));
        }
        Ok(())
    }
}

/// Texto de una respuesta o de un trozo del stream: `choices[0].text` (Mercury, Ollama), `.delta.content` (stream de
/// Codestral) o `.message.content` (respuesta entera de Codestral).
pub fn text_of(value: &Value) -> Option<&str> {
    let choice = &value["choices"][0];
    choice["text"]
        .as_str()
        .or_else(|| choice["delta"]["content"].as_str())
        .or_else(|| choice["message"]["content"].as_str())
}

/// Lee el stream SSE de cualquiera de los tres formatos.
#[derive(Default)]
pub struct FimDecoder {
    complete: bool,
}

impl Decoder for FimDecoder {
    fn feed(&mut self, event: &SseEvent, out: &mut Vec<Event>) -> Result<(), AiError> {
        if event.data.trim() == "[DONE]" {
            self.complete = true;
            return Ok(());
        }
        let Ok(value) = serde_json::from_str::<Value>(&event.data) else { return Ok(()) };
        if let Some(message) = value["error"]["message"]
            .as_str()
            .or_else(|| value["message"].as_str().filter(|_| value["object"] == "error"))
        {
            return Err(AiError::Api(message.to_owned()));
        }
        if let Some(text) = text_of(&value).filter(|text| !text.is_empty()) {
            out.push(Event::Text(text.to_owned()));
        }
        if !value["choices"][0]["finish_reason"].is_null() {
            self.complete = true;
        }
        Ok(())
    }

    fn finish(&mut self, _out: &mut Vec<Event>) {}

    fn is_complete(&self) -> bool {
        self.complete
    }
}

#[cfg(test)]
#[path = "fim_tests.rs"]
mod tests;