Small commit to add hybrid search and autoembedding

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()
+}