mrapids 0.1.31

// Embeddings module for semantic search
// Supports local (ONNX) and remote (OpenAI) embedding generation

#![allow(dead_code)]

use crate::core::api::ApiError;
use anyhow::{Context, Result};
use std::path::PathBuf;

/// Embedding configuration
#[derive(Debug, Clone)]
pub enum EmbeddingProvider {
    /// No embeddings - use keyword search only
    None,
    /// Local ONNX model (BGE-small-en-v1.5)
    Local,
    /// OpenAI text-embedding-3-small
    OpenAI { api_key: String },
}

impl Default for EmbeddingProvider {
    fn default() -> Self {
        Self::None
    }
}

/// Embedding dimensions by provider
pub const DIMENSIONS_BGE_SMALL: usize = 384;
pub const DIMENSIONS_OPENAI_SMALL: usize = 1536;

/// Model file info
pub const BGE_MODEL_NAME: &str = "bge-small-en-v1.5";
pub const BGE_MODEL_URL: &str =
    "https://huggingface.co/BAAI/bge-small-en-v1.5/resolve/main/onnx/model.onnx";
pub const BGE_TOKENIZER_URL: &str =
    "https://huggingface.co/BAAI/bge-small-en-v1.5/resolve/main/tokenizer.json";

/// Get the models directory
pub fn get_models_dir() -> Result<PathBuf> {
    let home = dirs::home_dir().context("Could not find home directory")?;
    let models_dir = home.join(".mrapids").join("models");
    std::fs::create_dir_all(&models_dir)?;
    Ok(models_dir)
}

/// Embedding engine trait
pub trait EmbeddingEngine: Send + Sync {
    /// Get the embedding dimensions
    fn dimensions(&self) -> usize;

    /// Embed a single text
    fn embed(&self, text: &str) -> Result<Vec<f32>>;

    /// Embed multiple texts (batch)
    fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        texts.iter().map(|t| self.embed(t)).collect()
    }
}

/// No-op embedding engine (for keyword search only)
pub struct NoOpEmbeddingEngine;

impl EmbeddingEngine for NoOpEmbeddingEngine {
    fn dimensions(&self) -> usize {
        0
    }

    fn embed(&self, _text: &str) -> Result<Vec<f32>> {
        Ok(Vec::new())
    }
}

/// Local ONNX-based embedding engine
/// Uses BGE-small-en-v1.5 (384 dimensions)
#[cfg(feature = "embeddings")]
pub struct LocalEmbeddingEngine {
    session: ort::Session,
    tokenizer: tokenizers::Tokenizer,
}

#[cfg(feature = "embeddings")]
impl LocalEmbeddingEngine {
    pub fn new() -> Result<Self> {
        let models_dir = get_models_dir()?;
        let model_path = models_dir.join("bge-small-en-v1.5.onnx");
        let tokenizer_path = models_dir.join("bge-small-en-v1.5-tokenizer.json");

        // Check if model exists, download if not
        if !model_path.exists() {
            eprintln!("Downloading BGE-small embedding model (~50MB)...");
            download_file(BGE_MODEL_URL, &model_path)?;
        }

        if !tokenizer_path.exists() {
            eprintln!("Downloading tokenizer...");
            download_file(BGE_TOKENIZER_URL, &tokenizer_path)?;
        }

        // Load ONNX session
        let session = ort::Session::builder()?
            .with_optimization_level(ort::GraphOptimizationLevel::Level3)?
            .commit_from_file(&model_path)?;

        // Load tokenizer
        let tokenizer = tokenizers::Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| ApiError::ValidationError(format!("Failed to load tokenizer: {}", e)))?;

        Ok(Self { session, tokenizer })
    }
}

#[cfg(feature = "embeddings")]
impl EmbeddingEngine for LocalEmbeddingEngine {
    fn dimensions(&self) -> usize {
        DIMENSIONS_BGE_SMALL
    }

    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        // Tokenize
        let encoding = self
            .tokenizer
            .encode(text, true)
            .map_err(|e| ApiError::ValidationError(format!("Tokenization failed: {}", e)))?;

        let input_ids: Vec<i64> = encoding.get_ids().iter().map(|&id| id as i64).collect();
        let attention_mask: Vec<i64> = encoding
            .get_attention_mask()
            .iter()
            .map(|&m| m as i64)
            .collect();
        let token_type_ids: Vec<i64> = vec![0i64; input_ids.len()];

        let seq_len = input_ids.len();

        // Create tensors
        let input_ids_tensor =
            ort::Value::from_array(ndarray::Array2::from_shape_vec((1, seq_len), input_ids)?)?;
        let attention_mask_tensor = ort::Value::from_array(ndarray::Array2::from_shape_vec(
            (1, seq_len),
            attention_mask,
        )?)?;
        let token_type_ids_tensor = ort::Value::from_array(ndarray::Array2::from_shape_vec(
            (1, seq_len),
            token_type_ids,
        )?)?;

        // Run inference
        let outputs = self.session.run(ort::inputs![
            "input_ids" => input_ids_tensor,
            "attention_mask" => attention_mask_tensor,
            "token_type_ids" => token_type_ids_tensor,
        ]?)?;

        // Extract embeddings (CLS token or mean pooling)
        let output = outputs[0].try_extract_tensor::<f32>()?;
        let embedding: Vec<f32> = output
            .view()
            .iter()
            .take(DIMENSIONS_BGE_SMALL)
            .cloned()
            .collect();

        // L2 normalize
        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        let normalized: Vec<f32> = embedding.iter().map(|x| x / norm).collect();

        Ok(normalized)
    }
}

/// OpenAI embedding engine
pub struct OpenAIEmbeddingEngine {
    api_key: String,
    client: reqwest::blocking::Client,
}

impl OpenAIEmbeddingEngine {
    pub fn new(api_key: String) -> Self {
        Self {
            api_key,
            client: reqwest::blocking::Client::new(),
        }
    }

    pub fn from_env() -> Result<Self> {
        let api_key = std::env::var("OPENAI_API_KEY")
            .context("OPENAI_API_KEY environment variable not set")?;
        Ok(Self::new(api_key))
    }
}

impl EmbeddingEngine for OpenAIEmbeddingEngine {
    fn dimensions(&self) -> usize {
        DIMENSIONS_OPENAI_SMALL
    }

    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let response = self
            .client
            .post("https://api.openai.com/v1/embeddings")
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&serde_json::json!({
                "model": "text-embedding-3-small",
                "input": text,
            }))
            .send()
            .context("Failed to call OpenAI API")?;

        if !response.status().is_success() {
            let error_text = response.text().unwrap_or_default();
            return Err(ApiError::NetworkError(format!("OpenAI API error: {}", error_text)).into());
        }

        let body: serde_json::Value = response.json()?;
        let embedding = body["data"][0]["embedding"]
            .as_array()
            .context("Invalid response format")?
            .iter()
            .filter_map(|v| v.as_f64().map(|f| f as f32))
            .collect();

        Ok(embedding)
    }

    fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        // OpenAI supports batch embedding
        let response = self
            .client
            .post("https://api.openai.com/v1/embeddings")
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&serde_json::json!({
                "model": "text-embedding-3-small",
                "input": texts,
            }))
            .send()
            .context("Failed to call OpenAI API")?;

        if !response.status().is_success() {
            let error_text = response.text().unwrap_or_default();
            return Err(ApiError::NetworkError(format!("OpenAI API error: {}", error_text)).into());
        }

        let body: serde_json::Value = response.json()?;
        let data = body["data"].as_array().context("Invalid response format")?;

        let mut embeddings = Vec::with_capacity(texts.len());
        for item in data {
            let embedding: Vec<f32> = item["embedding"]
                .as_array()
                .context("Missing embedding")?
                .iter()
                .filter_map(|v| v.as_f64().map(|f| f as f32))
                .collect();
            embeddings.push(embedding);
        }

        Ok(embeddings)
    }
}

/// Download a file from URL
#[cfg(feature = "embeddings")]
fn download_file(url: &str, path: &PathBuf) -> Result<()> {
    use std::io::Write;

    let response = reqwest::blocking::get(url).context("Failed to download file")?;

    if !response.status().is_success() {
        return Err(
            ApiError::NetworkError(format!("Download failed: HTTP {}", response.status())).into(),
        );
    }

    let bytes = response.bytes()?;
    let mut file = std::fs::File::create(path)?;
    file.write_all(&bytes)?;

    Ok(())
}

/// Create an embedding engine based on provider
pub fn create_embedding_engine(provider: &EmbeddingProvider) -> Result<Box<dyn EmbeddingEngine>> {
    match provider {
        EmbeddingProvider::None => Ok(Box::new(NoOpEmbeddingEngine)),
        #[cfg(feature = "embeddings")]
        EmbeddingProvider::Local => Ok(Box::new(LocalEmbeddingEngine::new()?)),
        #[cfg(not(feature = "embeddings"))]
        EmbeddingProvider::Local => {
            return Err(ApiError::ValidationError("Local embeddings require the 'embeddings' feature. Rebuild with: cargo build --features embeddings".to_string()).into())
        }
        EmbeddingProvider::OpenAI { api_key } => {
            Ok(Box::new(OpenAIEmbeddingEngine::new(api_key.clone())))
        }
    }
}

/// Compute cosine similarity between two vectors
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() || a.is_empty() {
        return 0.0;
    }

    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm_a == 0.0 || norm_b == 0.0 {
        return 0.0;
    }

    dot / (norm_a * norm_b)
}

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

    #[test]
    fn test_cosine_similarity() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.0001);

        let c = vec![0.0, 1.0, 0.0];
        assert!((cosine_similarity(&a, &c) - 0.0).abs() < 0.0001);
    }

    #[test]
    fn test_noop_engine() {
        let engine = NoOpEmbeddingEngine;
        assert_eq!(engine.dimensions(), 0);
        assert!(engine.embed("test").unwrap().is_empty());
    }
}