Small commit to add hybrid search and autoembedding

milli/src/vector/error.rs

@@ -0,0 +1,229 @@
+use std::path::PathBuf;
+
+use hf_hub::api::sync::ApiError;
+
+use crate::error::FaultSource;
+use crate::vector::openai::OpenAiError;
+
+#[derive(Debug, thiserror::Error)]
+#[error("Error while generating embeddings: {inner}")]
+pub struct Error {
+    pub inner: Box<ErrorKind>,
+}
+
+impl<I: Into<ErrorKind>> From<I> for Error {
+    fn from(value: I) -> Self {
+        Self { inner: Box::new(value.into()) }
+    }
+}
+
+impl Error {
+    pub fn fault(&self) -> FaultSource {
+        match &*self.inner {
+            ErrorKind::NewEmbedderError(inner) => inner.fault,
+            ErrorKind::EmbedError(inner) => inner.fault,
+        }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ErrorKind {
+    #[error(transparent)]
+    NewEmbedderError(#[from] NewEmbedderError),
+    #[error(transparent)]
+    EmbedError(#[from] EmbedError),
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("{fault}: {kind}")]
+pub struct EmbedError {
+    pub kind: EmbedErrorKind,
+    pub fault: FaultSource,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum EmbedErrorKind {
+    #[error("could not tokenize: {0}")]
+    Tokenize(Box<dyn std::error::Error + Send + Sync>),
+    #[error("unexpected tensor shape: {0}")]
+    TensorShape(candle_core::Error),
+    #[error("unexpected tensor value: {0}")]
+    TensorValue(candle_core::Error),
+    #[error("could not run model: {0}")]
+    ModelForward(candle_core::Error),
+    #[error("could not reach OpenAI: {0}")]
+    OpenAiNetwork(reqwest::Error),
+    #[error("unexpected response from OpenAI: {0}")]
+    OpenAiUnexpected(reqwest::Error),
+    #[error("could not authenticate against OpenAI: {0}")]
+    OpenAiAuth(OpenAiError),
+    #[error("sent too many requests to OpenAI: {0}")]
+    OpenAiTooManyRequests(OpenAiError),
+    #[error("received internal error from OpenAI: {0}")]
+    OpenAiInternalServerError(OpenAiError),
+    #[error("sent too many tokens in a request to OpenAI: {0}")]
+    OpenAiTooManyTokens(OpenAiError),
+    #[error("received unhandled HTTP status code {0} from OpenAI")]
+    OpenAiUnhandledStatusCode(u16),
+}
+
+impl EmbedError {
+    pub fn tokenize(inner: Box<dyn std::error::Error + Send + Sync>) -> Self {
+        Self { kind: EmbedErrorKind::Tokenize(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn tensor_shape(inner: candle_core::Error) -> Self {
+        Self { kind: EmbedErrorKind::TensorShape(inner), fault: FaultSource::Bug }
+    }
+
+    pub fn tensor_value(inner: candle_core::Error) -> Self {
+        Self { kind: EmbedErrorKind::TensorValue(inner), fault: FaultSource::Bug }
+    }
+
+    pub fn model_forward(inner: candle_core::Error) -> Self {
+        Self { kind: EmbedErrorKind::ModelForward(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn openai_network(inner: reqwest::Error) -> Self {
+        Self { kind: EmbedErrorKind::OpenAiNetwork(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn openai_unexpected(inner: reqwest::Error) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiUnexpected(inner), fault: FaultSource::Bug }
+    }
+
+    pub(crate) fn openai_auth_error(inner: OpenAiError) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiAuth(inner), fault: FaultSource::User }
+    }
+
+    pub(crate) fn openai_too_many_requests(inner: OpenAiError) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiTooManyRequests(inner), fault: FaultSource::Runtime }
+    }
+
+    pub(crate) fn openai_internal_server_error(inner: OpenAiError) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiInternalServerError(inner), fault: FaultSource::Runtime }
+    }
+
+    pub(crate) fn openai_too_many_tokens(inner: OpenAiError) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiTooManyTokens(inner), fault: FaultSource::Bug }
+    }
+
+    pub(crate) fn openai_unhandled_status_code(code: u16) -> EmbedError {
+        Self { kind: EmbedErrorKind::OpenAiUnhandledStatusCode(code), fault: FaultSource::Bug }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("{fault}: {kind}")]
+pub struct NewEmbedderError {
+    pub kind: NewEmbedderErrorKind,
+    pub fault: FaultSource,
+}
+
+impl NewEmbedderError {
+    pub fn open_config(config_filename: PathBuf, inner: std::io::Error) -> NewEmbedderError {
+        let open_config = OpenConfig { filename: config_filename, inner };
+
+        Self { kind: NewEmbedderErrorKind::OpenConfig(open_config), fault: FaultSource::Runtime }
+    }
+
+    pub fn deserialize_config(
+        config: String,
+        config_filename: PathBuf,
+        inner: serde_json::Error,
+    ) -> NewEmbedderError {
+        let deserialize_config = DeserializeConfig { config, filename: config_filename, inner };
+        Self {
+            kind: NewEmbedderErrorKind::DeserializeConfig(deserialize_config),
+            fault: FaultSource::Runtime,
+        }
+    }
+
+    pub fn open_tokenizer(
+        tokenizer_filename: PathBuf,
+        inner: Box<dyn std::error::Error + Send + Sync>,
+    ) -> NewEmbedderError {
+        let open_tokenizer = OpenTokenizer { filename: tokenizer_filename, inner };
+        Self {
+            kind: NewEmbedderErrorKind::OpenTokenizer(open_tokenizer),
+            fault: FaultSource::Runtime,
+        }
+    }
+
+    pub fn new_api_fail(inner: ApiError) -> Self {
+        Self { kind: NewEmbedderErrorKind::NewApiFail(inner), fault: FaultSource::Bug }
+    }
+
+    pub fn api_get(inner: ApiError) -> Self {
+        Self { kind: NewEmbedderErrorKind::ApiGet(inner), fault: FaultSource::Undecided }
+    }
+
+    pub fn pytorch_weight(inner: candle_core::Error) -> Self {
+        Self { kind: NewEmbedderErrorKind::PytorchWeight(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn safetensor_weight(inner: candle_core::Error) -> Self {
+        Self { kind: NewEmbedderErrorKind::PytorchWeight(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn load_model(inner: candle_core::Error) -> Self {
+        Self { kind: NewEmbedderErrorKind::LoadModel(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn openai_initialize_web_client(inner: reqwest::Error) -> Self {
+        Self { kind: NewEmbedderErrorKind::InitWebClient(inner), fault: FaultSource::Runtime }
+    }
+
+    pub fn openai_invalid_api_key_format(inner: reqwest::header::InvalidHeaderValue) -> Self {
+        Self { kind: NewEmbedderErrorKind::InvalidApiKeyFormat(inner), fault: FaultSource::User }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("could not open config at {filename:?}: {inner}")]
+pub struct OpenConfig {
+    pub filename: PathBuf,
+    pub inner: std::io::Error,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("could not deserialize config at {filename}: {inner}. Config follows:\n{config}")]
+pub struct DeserializeConfig {
+    pub config: String,
+    pub filename: PathBuf,
+    pub inner: serde_json::Error,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("could not open tokenizer at {filename}: {inner}")]
+pub struct OpenTokenizer {
+    pub filename: PathBuf,
+    #[source]
+    pub inner: Box<dyn std::error::Error + Send + Sync>,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum NewEmbedderErrorKind {
+    // hf
+    #[error(transparent)]
+    OpenConfig(OpenConfig),
+    #[error(transparent)]
+    DeserializeConfig(DeserializeConfig),
+    #[error(transparent)]
+    OpenTokenizer(OpenTokenizer),
+    #[error("could not build weights from Pytorch weights: {0}")]
+    PytorchWeight(candle_core::Error),
+    #[error("could not build weights from Safetensor weights: {0}")]
+    SafetensorWeight(candle_core::Error),
+    #[error("could not spawn HG_HUB API client: {0}")]
+    NewApiFail(ApiError),
+    #[error("fetching file from HG_HUB failed: {0}")]
+    ApiGet(ApiError),
+    #[error("loading model failed: {0}")]
+    LoadModel(candle_core::Error),
+    // openai
+    #[error("initializing web client for sending embedding requests failed: {0}")]
+    InitWebClient(reqwest::Error),
+    #[error("The API key passed to Authorization error was in an invalid format: {0}")]
+    InvalidApiKeyFormat(reqwest::header::InvalidHeaderValue),
+}

milli/src/vector/hf.rs

@@ -0,0 +1,192 @@
+use candle_core::Tensor;
+use candle_nn::VarBuilder;
+use candle_transformers::models::bert::{BertModel, Config, DTYPE};
+// FIXME: currently we'll be using the hub to retrieve model, in the future we might want to embed it into Meilisearch itself
+use hf_hub::api::sync::Api;
+use hf_hub::{Repo, RepoType};
+use tokenizers::{PaddingParams, Tokenizer};
+
+pub use super::error::{EmbedError, Error, NewEmbedderError};
+use super::{Embedding, Embeddings};
+
+#[derive(
+    Debug,
+    Clone,
+    Copy,
+    Default,
+    Hash,
+    PartialEq,
+    Eq,
+    serde::Deserialize,
+    serde::Serialize,
+    deserr::Deserr,
+)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub enum WeightSource {
+    #[default]
+    Safetensors,
+    Pytorch,
+}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
+pub struct EmbedderOptions {
+    pub model: String,
+    pub revision: Option<String>,
+    pub weight_source: WeightSource,
+    pub normalize_embeddings: bool,
+}
+
+impl EmbedderOptions {
+    pub fn new() -> Self {
+        Self {
+            //model: "sentence-transformers/all-MiniLM-L6-v2".to_string(),
+            model: "BAAI/bge-base-en-v1.5".to_string(),
+            //revision: Some("refs/pr/21".to_string()),
+            revision: None,
+            //weight_source: Default::default(),
+            weight_source: WeightSource::Pytorch,
+            normalize_embeddings: true,
+        }
+    }
+}
+
+impl Default for EmbedderOptions {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Perform embedding of documents and queries
+pub struct Embedder {
+    model: BertModel,
+    tokenizer: Tokenizer,
+    options: EmbedderOptions,
+}
+
+impl std::fmt::Debug for Embedder {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Embedder")
+            .field("model", &self.options.model)
+            .field("tokenizer", &self.tokenizer)
+            .field("options", &self.options)
+            .finish()
+    }
+}
+
+impl Embedder {
+    pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
+        let device = candle_core::Device::Cpu;
+        let repo = match options.revision.clone() {
+            Some(revision) => Repo::with_revision(options.model.clone(), RepoType::Model, revision),
+            None => Repo::model(options.model.clone()),
+        };
+        let (config_filename, tokenizer_filename, weights_filename) = {
+            let api = Api::new().map_err(NewEmbedderError::new_api_fail)?;
+            let api = api.repo(repo);
+            let config = api.get("config.json").map_err(NewEmbedderError::api_get)?;
+            let tokenizer = api.get("tokenizer.json").map_err(NewEmbedderError::api_get)?;
+            let weights = match options.weight_source {
+                WeightSource::Pytorch => {
+                    api.get("pytorch_model.bin").map_err(NewEmbedderError::api_get)?
+                }
+                WeightSource::Safetensors => {
+                    api.get("model.safetensors").map_err(NewEmbedderError::api_get)?
+                }
+            };
+            (config, tokenizer, weights)
+        };
+
+        let config = std::fs::read_to_string(&config_filename)
+            .map_err(|inner| NewEmbedderError::open_config(config_filename.clone(), inner))?;
+        let config: Config = serde_json::from_str(&config).map_err(|inner| {
+            NewEmbedderError::deserialize_config(config, config_filename, inner)
+        })?;
+        let mut tokenizer = Tokenizer::from_file(&tokenizer_filename)
+            .map_err(|inner| NewEmbedderError::open_tokenizer(tokenizer_filename, inner))?;
+
+        let vb = match options.weight_source {
+            WeightSource::Pytorch => VarBuilder::from_pth(&weights_filename, DTYPE, &device)
+                .map_err(NewEmbedderError::pytorch_weight)?,
+            WeightSource::Safetensors => unsafe {
+                VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)
+                    .map_err(NewEmbedderError::safetensor_weight)?
+            },
+        };
+
+        let model = BertModel::load(vb, &config).map_err(NewEmbedderError::load_model)?;
+
+        if let Some(pp) = tokenizer.get_padding_mut() {
+            pp.strategy = tokenizers::PaddingStrategy::BatchLongest
+        } else {
+            let pp = PaddingParams {
+                strategy: tokenizers::PaddingStrategy::BatchLongest,
+                ..Default::default()
+            };
+            tokenizer.with_padding(Some(pp));
+        }
+
+        Ok(Self { model, tokenizer, options })
+    }
+
+    pub async fn embed(
+        &self,
+        mut texts: Vec<String>,
+    ) -> std::result::Result<Vec<Embeddings<f32>>, EmbedError> {
+        let tokens = match texts.len() {
+            1 => vec![self
+                .tokenizer
+                .encode(texts.pop().unwrap(), true)
+                .map_err(EmbedError::tokenize)?],
+            _ => self.tokenizer.encode_batch(texts, true).map_err(EmbedError::tokenize)?,
+        };
+        let token_ids = tokens
+            .iter()
+            .map(|tokens| {
+                let tokens = tokens.get_ids().to_vec();
+                Tensor::new(tokens.as_slice(), &self.model.device).map_err(EmbedError::tensor_shape)
+            })
+            .collect::<Result<Vec<_>, EmbedError>>()?;
+
+        let token_ids = Tensor::stack(&token_ids, 0).map_err(EmbedError::tensor_shape)?;
+        let token_type_ids = token_ids.zeros_like().map_err(EmbedError::tensor_shape)?;
+        let embeddings =
+            self.model.forward(&token_ids, &token_type_ids).map_err(EmbedError::model_forward)?;
+
+        // Apply some avg-pooling by taking the mean embedding value for all tokens (including padding)
+        let (_n_sentence, n_tokens, _hidden_size) =
+            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
+
+        let embeddings = (embeddings.sum(1).map_err(EmbedError::tensor_value)? / (n_tokens as f64))
+            .map_err(EmbedError::tensor_shape)?;
+
+        let embeddings: Tensor = if self.options.normalize_embeddings {
+            normalize_l2(&embeddings).map_err(EmbedError::tensor_value)?
+        } else {
+            embeddings
+        };
+
+        let embeddings: Vec<Embedding> = embeddings.to_vec2().map_err(EmbedError::tensor_shape)?;
+        Ok(embeddings.into_iter().map(Embeddings::from_single_embedding).collect())
+    }
+
+    pub async fn embed_chunks(
+        &self,
+        text_chunks: Vec<Vec<String>>,
+    ) -> std::result::Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
+        futures::future::try_join_all(text_chunks.into_iter().map(|prompts| self.embed(prompts)))
+            .await
+    }
+
+    pub fn chunk_count_hint(&self) -> usize {
+        1
+    }
+
+    pub fn prompt_count_in_chunk_hint(&self) -> usize {
+        std::thread::available_parallelism().map(|x| x.get()).unwrap_or(8)
+    }
+}
+
+fn normalize_l2(v: &Tensor) -> Result<Tensor, candle_core::Error> {
+    v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)
+}

milli/src/vector/mod.rs

@@ -0,0 +1,142 @@
+use self::error::{EmbedError, NewEmbedderError};
+use crate::prompt::PromptData;
+
+pub mod error;
+pub mod hf;
+pub mod openai;
+pub mod settings;
+
+pub use self::error::Error;
+
+pub type Embedding = Vec<f32>;
+
+pub struct Embeddings<F> {
+    data: Vec<F>,
+    dimension: usize,
+}
+
+impl<F> Embeddings<F> {
+    pub fn new(dimension: usize) -> Self {
+        Self { data: Default::default(), dimension }
+    }
+
+    pub fn from_single_embedding(embedding: Vec<F>) -> Self {
+        Self { dimension: embedding.len(), data: embedding }
+    }
+
+    pub fn from_inner(data: Vec<F>, dimension: usize) -> Result<Self, Vec<F>> {
+        let mut this = Self::new(dimension);
+        this.append(data)?;
+        Ok(this)
+    }
+
+    pub fn dimension(&self) -> usize {
+        self.dimension
+    }
+
+    pub fn into_inner(self) -> Vec<F> {
+        self.data
+    }
+
+    pub fn as_inner(&self) -> &[F] {
+        &self.data
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &'_ [F]> + '_ {
+        self.data.as_slice().chunks_exact(self.dimension)
+    }
+
+    pub fn push(&mut self, mut embedding: Vec<F>) -> Result<(), Vec<F>> {
+        if embedding.len() != self.dimension {
+            return Err(embedding);
+        }
+        self.data.append(&mut embedding);
+        Ok(())
+    }
+
+    pub fn append(&mut self, mut embeddings: Vec<F>) -> Result<(), Vec<F>> {
+        if embeddings.len() % self.dimension != 0 {
+            return Err(embeddings);
+        }
+        self.data.append(&mut embeddings);
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+pub enum Embedder {
+    HuggingFace(hf::Embedder),
+    OpenAi(openai::Embedder),
+}
+
+#[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)]
+pub struct EmbeddingConfig {
+    pub embedder_options: EmbedderOptions,
+    pub prompt: PromptData,
+    // TODO: add metrics and anything needed
+}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
+pub enum EmbedderOptions {
+    HuggingFace(hf::EmbedderOptions),
+    OpenAi(openai::EmbedderOptions),
+}
+
+impl Default for EmbedderOptions {
+    fn default() -> Self {
+        Self::HuggingFace(Default::default())
+    }
+}
+
+impl EmbedderOptions {
+    pub fn huggingface() -> Self {
+        Self::HuggingFace(hf::EmbedderOptions::new())
+    }
+
+    pub fn openai(api_key: String) -> Self {
+        Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key))
+    }
+}
+
+impl Embedder {
+    pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
+        Ok(match options {
+            EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?),
+            EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?),
+        })
+    }
+
+    pub async fn embed(
+        &self,
+        texts: Vec<String>,
+    ) -> std::result::Result<Vec<Embeddings<f32>>, EmbedError> {
+        match self {
+            Embedder::HuggingFace(embedder) => embedder.embed(texts).await,
+            Embedder::OpenAi(embedder) => embedder.embed(texts).await,
+        }
+    }
+
+    pub async fn embed_chunks(
+        &self,
+        text_chunks: Vec<Vec<String>>,
+    ) -> std::result::Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
+        match self {
+            Embedder::HuggingFace(embedder) => embedder.embed_chunks(text_chunks).await,
+            Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks).await,
+        }
+    }
+
+    pub fn chunk_count_hint(&self) -> usize {
+        match self {
+            Embedder::HuggingFace(embedder) => embedder.chunk_count_hint(),
+            Embedder::OpenAi(embedder) => embedder.chunk_count_hint(),
+        }
+    }
+
+    pub fn prompt_count_in_chunk_hint(&self) -> usize {
+        match self {
+            Embedder::HuggingFace(embedder) => embedder.prompt_count_in_chunk_hint(),
+            Embedder::OpenAi(embedder) => embedder.prompt_count_in_chunk_hint(),
+        }
+    }
+}

milli/src/vector/openai.rs

@@ -0,0 +1,416 @@
+use std::fmt::Display;
+
+use reqwest::StatusCode;
+use serde::{Deserialize, Serialize};
+
+use super::error::{EmbedError, NewEmbedderError};
+use super::{Embedding, Embeddings};
+
+#[derive(Debug)]
+pub struct Embedder {
+    client: reqwest::Client,
+    tokenizer: tiktoken_rs::CoreBPE,
+    options: EmbedderOptions,
+}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
+pub struct EmbedderOptions {
+    pub api_key: String,
+    pub embedding_model: EmbeddingModel,
+}
+
+#[derive(
+    Debug,
+    Clone,
+    Copy,
+    Default,
+    Hash,
+    PartialEq,
+    Eq,
+    serde::Serialize,
+    serde::Deserialize,
+    deserr::Deserr,
+)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub enum EmbeddingModel {
+    #[default]
+    TextEmbeddingAda002,
+}
+
+impl EmbeddingModel {
+    pub fn max_token(&self) -> usize {
+        match self {
+            EmbeddingModel::TextEmbeddingAda002 => 8191,
+        }
+    }
+
+    pub fn dimensions(&self) -> usize {
+        match self {
+            EmbeddingModel::TextEmbeddingAda002 => 1536,
+        }
+    }
+
+    pub fn name(&self) -> &'static str {
+        match self {
+            EmbeddingModel::TextEmbeddingAda002 => "text-embedding-ada-002",
+        }
+    }
+
+    pub fn from_name(name: &'static str) -> Option<Self> {
+        match name {
+            "text-embedding-ada-002" => Some(EmbeddingModel::TextEmbeddingAda002),
+            _ => None,
+        }
+    }
+}
+
+pub const OPENAI_EMBEDDINGS_URL: &str = "https://api.openai.com/v1/embeddings";
+
+impl EmbedderOptions {
+    pub fn with_default_model(api_key: String) -> Self {
+        Self { api_key, embedding_model: Default::default() }
+    }
+
+    pub fn with_embedding_model(api_key: String, embedding_model: EmbeddingModel) -> Self {
+        Self { api_key, embedding_model }
+    }
+}
+
+impl Embedder {
+    pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
+        let mut headers = reqwest::header::HeaderMap::new();
+        headers.insert(
+            reqwest::header::AUTHORIZATION,
+            reqwest::header::HeaderValue::from_str(&format!("Bearer {}", &options.api_key))
+                .map_err(NewEmbedderError::openai_invalid_api_key_format)?,
+        );
+        headers.insert(
+            reqwest::header::CONTENT_TYPE,
+            reqwest::header::HeaderValue::from_static("application/json"),
+        );
+        let client = reqwest::ClientBuilder::new()
+            .default_headers(headers)
+            .build()
+            .map_err(NewEmbedderError::openai_initialize_web_client)?;
+
+        // looking at the code it is very unclear that this can actually fail.
+        let tokenizer = tiktoken_rs::cl100k_base().unwrap();
+
+        Ok(Self { options, client, tokenizer })
+    }
+
+    pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Embeddings<f32>>, EmbedError> {
+        let mut tokenized = false;
+
+        for attempt in 0..7 {
+            let result = if tokenized {
+                self.try_embed_tokenized(&texts).await
+            } else {
+                self.try_embed(&texts).await
+            };
+
+            let retry_duration = match result {
+                Ok(embeddings) => return Ok(embeddings),
+                Err(retry) => {
+                    log::warn!("Failed: {}", retry.error);
+                    tokenized |= retry.must_tokenize();
+                    retry.into_duration(attempt)
+                }
+            }?;
+            log::warn!("Attempt #{}, retrying after {}ms.", attempt, retry_duration.as_millis());
+            tokio::time::sleep(retry_duration).await;
+        }
+
+        let result = if tokenized {
+            self.try_embed_tokenized(&texts).await
+        } else {
+            self.try_embed(&texts).await
+        };
+
+        result.map_err(Retry::into_error)
+    }
+
+    async fn check_response(response: reqwest::Response) -> Result<reqwest::Response, Retry> {
+        if !response.status().is_success() {
+            match response.status() {
+                StatusCode::UNAUTHORIZED => {
+                    let error_response: OpenAiErrorResponse = response
+                        .json()
+                        .await
+                        .map_err(EmbedError::openai_unexpected)
+                        .map_err(Retry::retry_later)?;
+
+                    return Err(Retry::give_up(EmbedError::openai_auth_error(
+                        error_response.error,
+                    )));
+                }
+                StatusCode::TOO_MANY_REQUESTS => {
+                    let error_response: OpenAiErrorResponse = response
+                        .json()
+                        .await
+                        .map_err(EmbedError::openai_unexpected)
+                        .map_err(Retry::retry_later)?;
+
+                    return Err(Retry::rate_limited(EmbedError::openai_too_many_requests(
+                        error_response.error,
+                    )));
+                }
+                StatusCode::INTERNAL_SERVER_ERROR => {
+                    let error_response: OpenAiErrorResponse = response
+                        .json()
+                        .await
+                        .map_err(EmbedError::openai_unexpected)
+                        .map_err(Retry::retry_later)?;
+                    return Err(Retry::retry_later(EmbedError::openai_internal_server_error(
+                        error_response.error,
+                    )));
+                }
+                StatusCode::SERVICE_UNAVAILABLE => {
+                    let error_response: OpenAiErrorResponse = response
+                        .json()
+                        .await
+                        .map_err(EmbedError::openai_unexpected)
+                        .map_err(Retry::retry_later)?;
+                    return Err(Retry::retry_later(EmbedError::openai_internal_server_error(
+                        error_response.error,
+                    )));
+                }
+                StatusCode::BAD_REQUEST => {
+                    // Most probably, one text contained too many tokens
+                    let error_response: OpenAiErrorResponse = response
+                        .json()
+                        .await
+                        .map_err(EmbedError::openai_unexpected)
+                        .map_err(Retry::retry_later)?;
+
+                    log::warn!("OpenAI: input was too long, retrying on tokenized version. For best performance, limit the size of your prompt.");
+
+                    return Err(Retry::retry_tokenized(EmbedError::openai_too_many_tokens(
+                        error_response.error,
+                    )));
+                }
+                code => {
+                    return Err(Retry::give_up(EmbedError::openai_unhandled_status_code(
+                        code.as_u16(),
+                    )));
+                }
+            }
+        }
+        Ok(response)
+    }
+
+    async fn try_embed<S: AsRef<str> + serde::Serialize>(
+        &self,
+        texts: &[S],
+    ) -> Result<Vec<Embeddings<f32>>, Retry> {
+        for text in texts {
+            log::trace!("Received prompt: {}", text.as_ref())
+        }
+        let request = OpenAiRequest { model: self.options.embedding_model.name(), input: texts };
+        let response = self
+            .client
+            .post(OPENAI_EMBEDDINGS_URL)
+            .json(&request)
+            .send()
+            .await
+            .map_err(EmbedError::openai_network)
+            .map_err(Retry::retry_later)?;
+
+        let response = Self::check_response(response).await?;
+
+        let response: OpenAiResponse = response
+            .json()
+            .await
+            .map_err(EmbedError::openai_unexpected)
+            .map_err(Retry::retry_later)?;
+
+        log::trace!("response: {:?}", response.data);
+
+        Ok(response
+            .data
+            .into_iter()
+            .map(|data| Embeddings::from_single_embedding(data.embedding))
+            .collect())
+    }
+
+    async fn try_embed_tokenized(&self, text: &[String]) -> Result<Vec<Embeddings<f32>>, Retry> {
+        pub const OVERLAP_SIZE: usize = 200;
+        let mut all_embeddings = Vec::with_capacity(text.len());
+        for text in text {
+            let max_token_count = self.options.embedding_model.max_token();
+            let encoded = self.tokenizer.encode_ordinary(text.as_str());
+            let len = encoded.len();
+            if len < max_token_count {
+                all_embeddings.append(&mut self.try_embed(&[text]).await?);
+                continue;
+            }
+
+            let mut tokens = encoded.as_slice();
+            let mut embeddings_for_prompt =
+                Embeddings::new(self.options.embedding_model.dimensions());
+            while tokens.len() > max_token_count {
+                let window = &tokens[..max_token_count];
+                embeddings_for_prompt.push(self.embed_tokens(window).await?).unwrap();
+
+                tokens = &tokens[max_token_count - OVERLAP_SIZE..];
+            }
+
+            // end of text
+            embeddings_for_prompt.push(self.embed_tokens(tokens).await?).unwrap();
+
+            all_embeddings.push(embeddings_for_prompt);
+        }
+        Ok(all_embeddings)
+    }
+
+    async fn embed_tokens(&self, tokens: &[usize]) -> Result<Embedding, Retry> {
+        for attempt in 0..9 {
+            let duration = match self.try_embed_tokens(tokens).await {
+                Ok(embedding) => return Ok(embedding),
+                Err(retry) => retry.into_duration(attempt),
+            }
+            .map_err(Retry::retry_later)?;
+
+            tokio::time::sleep(duration).await;
+        }
+
+        self.try_embed_tokens(tokens).await.map_err(|retry| Retry::give_up(retry.into_error()))
+    }
+
+    async fn try_embed_tokens(&self, tokens: &[usize]) -> Result<Embedding, Retry> {
+        let request =
+            OpenAiTokensRequest { model: self.options.embedding_model.name(), input: tokens };
+        let response = self
+            .client
+            .post(OPENAI_EMBEDDINGS_URL)
+            .json(&request)
+            .send()
+            .await
+            .map_err(EmbedError::openai_network)
+            .map_err(Retry::retry_later)?;
+
+        let response = Self::check_response(response).await?;
+
+        let mut response: OpenAiResponse = response
+            .json()
+            .await
+            .map_err(EmbedError::openai_unexpected)
+            .map_err(Retry::retry_later)?;
+        Ok(response.data.pop().map(|data| data.embedding).unwrap_or_default())
+    }
+
+    pub async fn embed_chunks(
+        &self,
+        text_chunks: Vec<Vec<String>>,
+    ) -> Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
+        futures::future::try_join_all(text_chunks.into_iter().map(|prompts| self.embed(prompts)))
+            .await
+    }
+
+    pub fn chunk_count_hint(&self) -> usize {
+        10
+    }
+
+    pub fn prompt_count_in_chunk_hint(&self) -> usize {
+        10
+    }
+}
+
+// retrying in case of failure
+
+struct Retry {
+    error: EmbedError,
+    strategy: RetryStrategy,
+}
+
+enum RetryStrategy {
+    GiveUp,
+    Retry,
+    RetryTokenized,
+    RetryAfterRateLimit,
+}
+
+impl Retry {
+    fn give_up(error: EmbedError) -> Self {
+        Self { error, strategy: RetryStrategy::GiveUp }
+    }
+
+    fn retry_later(error: EmbedError) -> Self {
+        Self { error, strategy: RetryStrategy::Retry }
+    }
+
+    fn retry_tokenized(error: EmbedError) -> Self {
+        Self { error, strategy: RetryStrategy::RetryTokenized }
+    }
+
+    fn rate_limited(error: EmbedError) -> Self {
+        Self { error, strategy: RetryStrategy::RetryAfterRateLimit }
+    }
+
+    fn into_duration(self, attempt: u32) -> Result<tokio::time::Duration, EmbedError> {
+        match self.strategy {
+            RetryStrategy::GiveUp => Err(self.error),
+            RetryStrategy::Retry => Ok(tokio::time::Duration::from_millis((10u64).pow(attempt))),
+            RetryStrategy::RetryTokenized => Ok(tokio::time::Duration::from_millis(1)),
+            RetryStrategy::RetryAfterRateLimit => {
+                Ok(tokio::time::Duration::from_millis(100 + 10u64.pow(attempt)))
+            }
+        }
+    }
+
+    fn must_tokenize(&self) -> bool {
+        matches!(self.strategy, RetryStrategy::RetryTokenized)
+    }
+
+    fn into_error(self) -> EmbedError {
+        self.error
+    }
+}
+
+// openai api structs
+
+#[derive(Debug, Serialize)]
+struct OpenAiRequest<'a, S: AsRef<str> + serde::Serialize> {
+    model: &'a str,
+    input: &'a [S],
+}
+
+#[derive(Debug, Serialize)]
+struct OpenAiTokensRequest<'a> {
+    model: &'a str,
+    input: &'a [usize],
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenAiResponse {
+    data: Vec<OpenAiEmbedding>,
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenAiErrorResponse {
+    error: OpenAiError,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct OpenAiError {
+    message: String,
+    // type: String,
+    code: Option<String>,
+}
+
+impl Display for OpenAiError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.code {
+            Some(code) => write!(f, "{} ({})", self.message, code),
+            None => write!(f, "{}", self.message),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenAiEmbedding {
+    embedding: Embedding,
+    // object: String,
+    // index: usize,
+}

milli/src/vector/settings.rs

@@ -0,0 +1,308 @@
+use deserr::Deserr;
+use serde::{Deserialize, Serialize};
+
+use crate::prompt::{PromptData, PromptFallbackStrategy};
+use crate::update::Setting;
+use crate::vector::hf::WeightSource;
+use crate::vector::EmbeddingConfig;
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub struct EmbeddingSettings {
+    #[serde(default, skip_serializing_if = "Setting::is_not_set", rename = "source")]
+    #[deserr(default, rename = "source")]
+    pub embedder_options: Setting<EmbedderSettings>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub prompt: Setting<PromptSettings>,
+}
+
+impl EmbeddingSettings {
+    pub fn apply(&mut self, new: Self) {
+        let EmbeddingSettings { embedder_options, prompt } = new;
+        self.embedder_options.apply(embedder_options);
+        self.prompt.apply(prompt);
+    }
+}
+
+impl From<EmbeddingConfig> for EmbeddingSettings {
+    fn from(value: EmbeddingConfig) -> Self {
+        Self {
+            embedder_options: Setting::Set(value.embedder_options.into()),
+            prompt: Setting::Set(value.prompt.into()),
+        }
+    }
+}
+
+impl From<EmbeddingSettings> for EmbeddingConfig {
+    fn from(value: EmbeddingSettings) -> Self {
+        let mut this = Self::default();
+        let EmbeddingSettings { embedder_options, prompt } = value;
+        if let Some(embedder_options) = embedder_options.set() {
+            this.embedder_options = embedder_options.into();
+        }
+        if let Some(prompt) = prompt.set() {
+            this.prompt = prompt.into();
+        }
+        this
+    }
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub struct PromptSettings {
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub template: Setting<String>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub strategy: Setting<PromptFallbackStrategy>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub fallback: Setting<String>,
+}
+
+impl PromptSettings {
+    pub fn apply(&mut self, new: Self) {
+        let PromptSettings { template, strategy, fallback } = new;
+        self.template.apply(template);
+        self.strategy.apply(strategy);
+        self.fallback.apply(fallback);
+    }
+}
+
+impl From<PromptData> for PromptSettings {
+    fn from(value: PromptData) -> Self {
+        Self {
+            template: Setting::Set(value.template),
+            strategy: Setting::Set(value.strategy),
+            fallback: Setting::Set(value.fallback),
+        }
+    }
+}
+
+impl From<PromptSettings> for PromptData {
+    fn from(value: PromptSettings) -> Self {
+        let mut this = PromptData::default();
+        let PromptSettings { template, strategy, fallback } = value;
+        if let Some(template) = template.set() {
+            this.template = template;
+        }
+        if let Some(strategy) = strategy.set() {
+            this.strategy = strategy;
+        }
+        if let Some(fallback) = fallback.set() {
+            this.fallback = fallback;
+        }
+        this
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+pub enum EmbedderSettings {
+    HuggingFace(Setting<HfEmbedderSettings>),
+    OpenAi(Setting<OpenAiEmbedderSettings>),
+}
+
+impl<E> Deserr<E> for EmbedderSettings
+where
+    E: deserr::DeserializeError,
+{
+    fn deserialize_from_value<V: deserr::IntoValue>(
+        value: deserr::Value<V>,
+        location: deserr::ValuePointerRef,
+    ) -> Result<Self, E> {
+        match value {
+            deserr::Value::Map(map) => {
+                if deserr::Map::len(&map) != 1 {
+                    return Err(deserr::take_cf_content(E::error::<V>(
+                        None,
+                        deserr::ErrorKind::Unexpected {
+                            msg: format!(
+                                "Expected a single field, got {} fields",
+                                deserr::Map::len(&map)
+                            ),
+                        },
+                        location,
+                    )));
+                }
+                let mut it = deserr::Map::into_iter(map);
+                let (k, v) = it.next().unwrap();
+
+                match k.as_str() {
+                    "huggingFace" => Ok(EmbedderSettings::HuggingFace(Setting::Set(
+                        HfEmbedderSettings::deserialize_from_value(
+                            v.into_value(),
+                            location.push_key(&k),
+                        )?,
+                    ))),
+                    "openAi" => Ok(EmbedderSettings::OpenAi(Setting::Set(
+                        OpenAiEmbedderSettings::deserialize_from_value(
+                            v.into_value(),
+                            location.push_key(&k),
+                        )?,
+                    ))),
+                    other => Err(deserr::take_cf_content(E::error::<V>(
+                        None,
+                        deserr::ErrorKind::UnknownKey {
+                            key: other,
+                            accepted: &["huggingFace", "openAi"],
+                        },
+                        location,
+                    ))),
+                }
+            }
+            _ => Err(deserr::take_cf_content(E::error::<V>(
+                None,
+                deserr::ErrorKind::IncorrectValueKind {
+                    actual: value,
+                    accepted: &[deserr::ValueKind::Map],
+                },
+                location,
+            ))),
+        }
+    }
+}
+
+impl Default for EmbedderSettings {
+    fn default() -> Self {
+        Self::HuggingFace(Default::default())
+    }
+}
+
+impl From<crate::vector::EmbedderOptions> for EmbedderSettings {
+    fn from(value: crate::vector::EmbedderOptions) -> Self {
+        match value {
+            crate::vector::EmbedderOptions::HuggingFace(hf) => {
+                Self::HuggingFace(Setting::Set(hf.into()))
+            }
+            crate::vector::EmbedderOptions::OpenAi(openai) => {
+                Self::OpenAi(Setting::Set(openai.into()))
+            }
+        }
+    }
+}
+
+impl From<EmbedderSettings> for crate::vector::EmbedderOptions {
+    fn from(value: EmbedderSettings) -> Self {
+        match value {
+            EmbedderSettings::HuggingFace(Setting::Set(hf)) => Self::HuggingFace(hf.into()),
+            EmbedderSettings::HuggingFace(_setting) => Self::HuggingFace(Default::default()),
+            EmbedderSettings::OpenAi(Setting::Set(ai)) => Self::OpenAi(ai.into()),
+            EmbedderSettings::OpenAi(_setting) => Self::OpenAi(
+                crate::vector::openai::EmbedderOptions::with_default_model(infer_api_key()),
+            ),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub struct HfEmbedderSettings {
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub model: Setting<String>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub revision: Setting<String>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub weight_source: Setting<WeightSource>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub normalize_embeddings: Setting<bool>,
+}
+
+impl HfEmbedderSettings {
+    pub fn apply(&mut self, new: Self) {
+        let HfEmbedderSettings {
+            model,
+            revision,
+            weight_source,
+            normalize_embeddings: normalize_embedding,
+        } = new;
+        self.model.apply(model);
+        self.revision.apply(revision);
+        self.weight_source.apply(weight_source);
+        self.normalize_embeddings.apply(normalize_embedding);
+    }
+}
+
+impl From<crate::vector::hf::EmbedderOptions> for HfEmbedderSettings {
+    fn from(value: crate::vector::hf::EmbedderOptions) -> Self {
+        Self {
+            model: Setting::Set(value.model),
+            revision: value.revision.map(Setting::Set).unwrap_or(Setting::NotSet),
+            weight_source: Setting::Set(value.weight_source),
+            normalize_embeddings: Setting::Set(value.normalize_embeddings),
+        }
+    }
+}
+
+impl From<HfEmbedderSettings> for crate::vector::hf::EmbedderOptions {
+    fn from(value: HfEmbedderSettings) -> Self {
+        let HfEmbedderSettings { model, revision, weight_source, normalize_embeddings } = value;
+        let mut this = Self::default();
+        if let Some(model) = model.set() {
+            this.model = model;
+        }
+        if let Some(revision) = revision.set() {
+            this.revision = Some(revision);
+        }
+        if let Some(weight_source) = weight_source.set() {
+            this.weight_source = weight_source;
+        }
+        if let Some(normalize_embeddings) = normalize_embeddings.set() {
+            this.normalize_embeddings = normalize_embeddings;
+        }
+        this
+    }
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+pub struct OpenAiEmbedderSettings {
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub api_key: Setting<String>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
+    pub embedding_model: Setting<crate::vector::openai::EmbeddingModel>,
+}
+
+impl OpenAiEmbedderSettings {
+    pub fn apply(&mut self, new: Self) {
+        let Self { api_key, embedding_model: embedding_mode } = new;
+        self.api_key.apply(api_key);
+        self.embedding_model.apply(embedding_mode);
+    }
+}
+
+impl From<crate::vector::openai::EmbedderOptions> for OpenAiEmbedderSettings {
+    fn from(value: crate::vector::openai::EmbedderOptions) -> Self {
+        Self {
+            api_key: Setting::Set(value.api_key),
+            embedding_model: Setting::Set(value.embedding_model),
+        }
+    }
+}
+
+impl From<OpenAiEmbedderSettings> for crate::vector::openai::EmbedderOptions {
+    fn from(value: OpenAiEmbedderSettings) -> Self {
+        let OpenAiEmbedderSettings { api_key, embedding_model } = value;
+        Self {
+            api_key: api_key.set().unwrap_or_else(infer_api_key),
+            embedding_model: embedding_model.set().unwrap_or_default(),
+        }
+    }
+}
+
+fn infer_api_key() -> String {
+    /// FIXME: get key from instance options?
+    std::env::var("MEILI_OPENAI_API_KEY").unwrap_or_default()
+}