second-brain-core 0.5.1

use std::sync::Mutex;

use anyhow::{Context, Result};
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};

pub struct Embedder {
    model: Mutex<TextEmbedding>,
}

pub fn query_prompt(text: &str) -> String {
    format!("Represent this sentence for searching relevant passages: {text}")
}

impl Embedder {
    pub fn new() -> Result<Self> {
        let model = TextEmbedding::try_new(
            InitOptions::new(EmbeddingModel::BGESmallENV15).with_show_download_progress(true),
        )
        .context("initializing embedding model")?;

        Ok(Self {
            model: Mutex::new(model),
        })
    }

    pub fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let mut model = self.model.lock().unwrap();
        let results = model
            .embed(vec![text], None)
            .context("generating embedding")?;

        results
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("no embedding returned"))
    }

    // BGE retrieval is asymmetric: the instruction prefix goes on the query only,
    // documents stay bare, so embed() is left unchanged.
    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        self.embed(&query_prompt(text))
    }

    pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let mut model = self.model.lock().unwrap();
        let owned: Vec<String> = texts.iter().map(|t| t.to_string()).collect();
        let refs: Vec<&str> = owned.iter().map(|s| s.as_str()).collect();
        model
            .embed(refs, None)
            .context("generating batch embeddings")
    }

    pub fn dimension(&self) -> usize {
        384
    }
}