ragrig 0.9.5

//! Embedding backend abstraction.
//!
//! The [`Embedder`] trait decouples text → vector conversion from any
//! specific provider, following the same pattern as
//! [`Generator`](crate::agents::Generator) and
//! [`VectorStore`](crate::store::VectorStore).  Three implementations are
//! provided:
//!
//! - [`OllamaEmbedder`] — delegates to a local Ollama server via rig-core
//! - `FastembedEmbedder` — runs Nomic-Embed-Text-v1.5 on CPU, zero network
//!   (only available with `--features internal-embed`)
//! - [`NoopEmbedder`] — returns empty vectors; used when embeddings are
//!   disabled (pure chat / forgetful mode)

use anyhow::{Result, anyhow};
use async_trait::async_trait;
#[cfg(feature = "internal-embed")]
use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
use rig_core::client::{EmbeddingsClient, Nothing};
use rig_core::embeddings::EmbeddingsBuilder;
use rig_core::providers::ollama;
#[cfg(feature = "internal-embed")]
use std::sync::{Mutex, OnceLock};

// ── Embedder trait ────────────────────────────────────────────────────────

/// Capability: convert text into dense vector representations.
///
/// Methods use `#[async_trait]` which expands to `Pin<Box<dyn Future>>` in
/// the rendered docs — just call them with `.await` as normal.
///
/// # Example
///
/// ```rust,no_run
/// use ragrig::embed::{Embedder, EmbedderSpec};
///
/// # async fn example() -> anyhow::Result<()> {
/// let embedder = EmbedderSpec::Ollama {
///     model: "nomic-embed-text".into(),
/// }.build()?;
/// let results = embedder.embed(vec!["quantum computing".into()]).await?;
/// for (text, vec) in results {
///     println!("{text}: {} dimensions", vec.len());
/// }
/// # Ok(())
/// # }
/// ```
#[async_trait]
pub trait Embedder: Send + Sync + std::fmt::Debug {
    /// Clone this embedder into a new heap allocation.
    fn clone_box(&self) -> Box<dyn Embedder>;

    /// Produce `(text, Vec<f32>)` pairs.  The returned vectors MUST be
    /// in the same order as the input texts.
    async fn embed(&self, texts: Vec<String>) -> Result<Vec<(String, Vec<f32>)>>;

    /// Human-readable backend label, e.g. "Ollama", "Fastembed".
    fn backend_name(&self) -> &'static str;

    /// The specific model in use, e.g. "nomic-embed-text".
    fn model_name(&self) -> &str;

    /// Dimensionality of the vectors produced by this embedder.
    /// Returns 0 when unknown (e.g. NoopEmbedder or not yet initialised).
    fn dimension(&self) -> usize;
}

impl Clone for Box<dyn Embedder> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

// ── Ollama embedder ───────────────────────────────────────────────────────

/// Talks to a local Ollama server for embeddings.
#[derive(Clone, Debug)]
pub struct OllamaEmbedder {
    model_name: String,
}

impl OllamaEmbedder {
    /// Create a new Ollama embedder for the given model.
    pub fn new(model: String) -> Self {
        Self { model_name: model }
    }
}

#[async_trait]
impl Embedder for OllamaEmbedder {
    fn clone_box(&self) -> Box<dyn Embedder> {
        Box::new(self.clone())
    }

    async fn embed(&self, texts: Vec<String>) -> Result<Vec<(String, Vec<f32>)>> {
        let client = ollama::Client::new(Nothing)
            .map_err(|_| anyhow!(crate::RagrigError::OllamaUnreachable {
                context: "embed documents".into(),
            }))?;
        let model = client.embedding_model(&self.model_name);
        let embedded = EmbeddingsBuilder::new(model)
            .documents(texts.clone())?
            .build()
            .await
            .map_err(|e| {
                let msg = e.to_string();
                if msg.contains("not found") || msg.contains("model not found") {
                    anyhow!(crate::RagrigError::EmbedModelNotFound {
                        model: self.model_name.to_string(),
                    })
                } else {
                    anyhow!("Ollama embedder: embedding failed for model '{}': {}", self.model_name, e)
                }
            })?;
        Ok(embedded
            .into_iter()
            .map(|(text, emb)| {
                (
                    text,
                    emb.first().vec.iter().map(|v| *v as f32).collect(),
                )
            })
            .collect())
    }

    fn backend_name(&self) -> &'static str {
        "Ollama"
    }

    fn model_name(&self) -> &str {
        &self.model_name
    }

    fn dimension(&self) -> usize {
        // Ollama nomic-embed-text outputs 768-d vectors.
        768
    }
}

// ── Fastembed embedder ────────────────────────────────────────────────────

/// Runs Nomic-Embed-Text-v1.5 directly on the CPU.  Zero network overhead.
/// Only available when the `internal-embed` feature is enabled.
#[cfg(feature = "internal-embed")]
#[derive(Clone, Debug, Default)]
pub struct FastembedEmbedder;

#[cfg(feature = "internal-embed")]
static FASTEMBED: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();

#[cfg(feature = "internal-embed")]
fn get_fastembed() -> &'static Mutex<TextEmbedding> {
    FASTEMBED.get_or_init(|| {
        log::info!("Initializing fastembed (Nomic-Embed-Text-v1.5) on CPU …");
        let model = TextEmbedding::try_new(TextInitOptions::new(
            EmbeddingModel::NomicEmbedTextV15,
        ))
        .expect("Failed to initialize fastembed model");
        Mutex::new(model)
    })
}

#[cfg(feature = "internal-embed")]
#[async_trait]
impl Embedder for FastembedEmbedder {
    fn clone_box(&self) -> Box<dyn Embedder> {
        Box::new(self.clone())
    }

    async fn embed(&self, texts: Vec<String>) -> Result<Vec<(String, Vec<f32>)>> {
        let texts_for_blocking = texts.clone();
        let vectors = tokio::task::spawn_blocking(move || {
            let mutex = get_fastembed();
            let mut model = mutex.lock().unwrap();
            model
                .embed(texts_for_blocking, None)
                .map_err(|e| anyhow!("fastembed: {}", e))
        })
        .await??;
        Ok(texts.into_iter().zip(vectors.into_iter()).collect())
    }

    fn backend_name(&self) -> &'static str {
        "Fastembed"
    }

    fn model_name(&self) -> &str {
        "Nomic-Embed-Text-v1.5"
    }

    fn dimension(&self) -> usize {
        768
    }
}

// ── No-op embedder ────────────────────────────────────────────────────────

/// Returns zero-vectors.  Useful for pure-chat / forgetful sessions
/// where document search is not needed.
#[derive(Clone, Debug, Default)]
pub struct NoopEmbedder;

#[async_trait]
impl Embedder for NoopEmbedder {
    fn clone_box(&self) -> Box<dyn Embedder> {
        Box::new(self.clone())
    }

    async fn embed(&self, texts: Vec<String>) -> Result<Vec<(String, Vec<f32>)>> {
        // Return zero vectors — the store will receive them but similarity
        // search will produce meaningless results, which is fine since the
        // caller should not be querying the store in this mode.
        Ok(texts
            .into_iter()
            .map(|t| (t, vec![0.0f32; 768]))
            .collect())
    }

    fn backend_name(&self) -> &'static str {
        "None"
    }

    fn model_name(&self) -> &str {
        "(disabled)"
    }

    fn dimension(&self) -> usize {
        0
    }
}

// ── Builder / Config ───────────────────────────────────────────────────────

/// A parsed `/embed` command payload.
#[derive(Clone, Debug)]
pub enum EmbedderSpec {
    /// Ollama embedding server.
    Ollama { /// Embedding model name (e.g. `"nomic-embed-text"`).
        model: String },
    #[cfg(feature = "internal-embed")]
    /// CPU-only embeddings via the `fastembed` crate.
    Fastembed,
    /// Embeddings disabled — search will not work.
    None,
}

impl EmbedderSpec {
    /// Convenience constructor: `Ollama` variant.
    pub fn ollama(model: impl Into<String>) -> Self {
        Self::Ollama { model: model.into() }
    }

    /// Convenience constructor: `None` variant (embeddings disabled).
    pub fn none() -> Self {
        Self::None
    }

    /// Parse a backend name and optional model into an `EmbedderSpec`.
    pub fn parse(backend: &str, model: Option<&str>) -> Result<Self> {
        match backend.to_lowercase().as_str() {
            "ollama" => {
                let model = model.unwrap_or("nomic-embed-text").to_string();
                Ok(Self::Ollama { model })
            }
            #[cfg(feature = "internal-embed")]
            "fastembed" => Ok(Self::Fastembed),
            "none" | "off" => Ok(Self::None),
            other => Err(anyhow!(
                "Unknown embedding backend: '{}'. Available: {}",
                other,
                Self::available_backends().join(", ")
            )),
        }
    }

    /// List of backend names supported by this build.
    pub fn available_backends() -> &'static [&'static str] {
        &[
            "ollama",
            #[cfg(feature = "internal-embed")]
            "fastembed",
            "none",
        ]
    }

    /// Build a boxed [`Embedder`] from this spec.
    pub fn build(&self) -> Result<Box<dyn Embedder>> {
        match self {
            Self::Ollama { model } => Ok(Box::new(OllamaEmbedder::new(model.clone()))),
            #[cfg(feature = "internal-embed")]
            Self::Fastembed => Ok(Box::new(FastembedEmbedder)),
            Self::None => Ok(Box::new(NoopEmbedder)),
        }
    }
}

#[cfg(feature = "internal-embed")]
impl EmbedderSpec {
    /// Convenience constructor: `Fastembed` variant.
    pub fn fastembed() -> Self {
        Self::Fastembed
    }
}

impl TryFrom<EmbedderSpec> for Box<dyn Embedder> {
    type Error = anyhow::Error;
    fn try_from(spec: EmbedderSpec) -> Result<Self, Self::Error> {
        spec.build()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_ollama_default_model() {
        let spec = EmbedderSpec::parse("ollama", None).unwrap();
        match spec {
            EmbedderSpec::Ollama { model } => assert_eq!(model, "nomic-embed-text"),
            _ => panic!("expected Ollama"),
        }
    }

    #[test]
    fn parse_ollama_custom_model() {
        let spec = EmbedderSpec::parse("ollama", Some("custom-model")).unwrap();
        match spec {
            EmbedderSpec::Ollama { model } => assert_eq!(model, "custom-model"),
            _ => panic!("expected Ollama"),
        }
    }

    #[test]
    fn parse_none() {
        let spec = EmbedderSpec::parse("none", None).unwrap();
        assert!(matches!(spec, EmbedderSpec::None));
    }

    #[test]
    fn parse_off_is_none() {
        let spec = EmbedderSpec::parse("off", None).unwrap();
        assert!(matches!(spec, EmbedderSpec::None));
    }

    #[test]
    fn parse_unknown_is_error() {
        assert!(EmbedderSpec::parse("openai", None).is_err());
    }

    #[test]
    fn available_backends_contains_ollama() {
        let backs = EmbedderSpec::available_backends();
        assert!(backs.contains(&"ollama"));
        assert!(backs.contains(&"none"));
    }

    #[test]
    fn build_none_returns_noop() {
        let embedder = EmbedderSpec::None.build().unwrap();
        assert_eq!(embedder.backend_name(), "None");
        assert_eq!(embedder.model_name(), "(disabled)");
        assert_eq!(embedder.dimension(), 0);
    }

    #[test]
    fn ollama_embedder_dimension() {
        let e = OllamaEmbedder::new("nomic-embed-text".into());
        assert_eq!(e.dimension(), 768);
    }

    #[tokio::test]
    async fn noop_embedder_returns_zero_vectors() {
        let e = NoopEmbedder;
        let result = e.embed(vec!["hello".into(), "world".into()]).await.unwrap();
        assert_eq!(result.len(), 2);
        for (_, v) in &result {
            assert_eq!(v.len(), 768);
            assert!(v.iter().all(|&x| x == 0.0));
        }
    }

    // ── TryFrom<EmbedderSpec> for Box<dyn Embedder> ─────────────────

    #[test]
    fn try_from_ollama_spec_succeeds() {
        use std::convert::TryFrom;
        let spec = EmbedderSpec::Ollama {
            model: "nomic-embed-text".into(),
        };
        let embedder = Box::<dyn Embedder>::try_from(spec).unwrap();
        assert_eq!(embedder.backend_name(), "Ollama");
        assert_eq!(embedder.model_name(), "nomic-embed-text");
        assert_eq!(embedder.dimension(), 768);
    }

    #[test]
    fn try_from_none_spec_succeeds() {
        use std::convert::TryFrom;
        let spec = EmbedderSpec::None;
        let embedder = Box::<dyn Embedder>::try_from(spec).unwrap();
        assert_eq!(embedder.backend_name(), "None");
        assert_eq!(embedder.model_name(), "(disabled)");
        assert_eq!(embedder.dimension(), 0);
    }

    #[test]
    fn try_from_unknown_spec_is_error() {
        let spec = EmbedderSpec::parse("bogus", None);
        assert!(spec.is_err());
    }
}