From 58972f35cb0d755fc9fc45de110beedc78565602 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:13:21 +0100 Subject: [PATCH 1/3] Allow `url` parameter for ollama embedder --- milli/src/update/settings.rs | 1 - milli/src/vector/mod.rs | 4 ++-- milli/src/vector/ollama.rs | 11 ++++------- milli/src/vector/settings.rs | 13 ++++++++++--- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index e902badc0..62e6b20b7 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1271,7 +1271,6 @@ pub fn validate_embedding_settings( check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&revision, "revision", inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; check_unset(&query, "query", inferred_source, name)?; check_unset(&input_field, "inputField", inferred_source, name)?; check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 65654af4a..8186e4409 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -201,8 +201,8 @@ impl EmbedderOptions { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } - pub fn ollama() -> Self { - Self::Ollama(ollama::EmbedderOptions::with_default_model()) + pub fn ollama(url: Option) -> Self { + Self::Ollama(ollama::EmbedderOptions::with_default_model(url)) } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 9c44e8052..7edfd13b5 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -12,15 +12,12 @@ pub struct Embedder { #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { pub embedding_model: String, + pub url: Option, } impl EmbedderOptions { - pub fn with_default_model() -> Self { - Self { embedding_model: "nomic-embed-text".into() } - } - - pub fn with_embedding_model(embedding_model: String) -> Self { - Self { embedding_model } + pub fn with_default_model(url: Option) -> Self { + Self { embedding_model: "nomic-embed-text".into(), url } } } @@ -31,7 +28,7 @@ impl Embedder { api_key: None, distribution: None, dimensions: None, - url: get_ollama_path(), + url: options.url.unwrap_or_else(get_ollama_path), query: serde_json::json!({ "model": model, }), diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index c5b0d0326..7760573f6 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -124,7 +124,7 @@ impl EmbeddingSettings { EmbedderSource::Ollama, EmbedderSource::Rest, ], - Self::URL => &[EmbedderSource::Rest], + Self::URL => &[EmbedderSource::Ollama, EmbedderSource::Rest], Self::QUERY => &[EmbedderSource::Rest], Self::INPUT_FIELD => &[EmbedderSource::Rest], Self::PATH_TO_EMBEDDINGS => &[EmbedderSource::Rest], @@ -146,7 +146,9 @@ impl EmbeddingSettings { EmbedderSource::HuggingFace => { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } - EmbedderSource::Ollama => &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE], + EmbedderSource::Ollama => { + &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL] + } EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], EmbedderSource::Rest => &[ Self::SOURCE, @@ -387,10 +389,15 @@ impl From for EmbeddingConfig { } EmbedderSource::Ollama => { let mut options: ollama::EmbedderOptions = - super::ollama::EmbedderOptions::with_default_model(); + super::ollama::EmbedderOptions::with_default_model(None); if let Some(model) = model.set() { options.embedding_model = model; } + + if let Some(url) = url.set() { + options.url = Some(url) + } + this.embedder_options = super::EmbedderOptions::Ollama(options); } EmbedderSource::HuggingFace => { From 4136630ea5c62ea2f0dc6d9faa4a3171b24da721 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:39:33 +0100 Subject: [PATCH 2/3] Use constants instead of raw strings in set_*set() --- milli/src/update/settings.rs | 121 ++++++++++++++++++++++++----------- 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 62e6b20b7..e0c559b85 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1225,14 +1225,24 @@ pub fn validate_embedding_settings( }; match inferred_source { EmbedderSource::OpenAi => { - check_unset(&revision, "revision", inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; if let Setting::Set(model) = &model { let model = crate::vector::openai::EmbeddingModel::from_name(model.as_str()) @@ -1266,46 +1276,81 @@ pub fn validate_embedding_settings( } EmbedderSource::Ollama => { // Dimensions get inferred, only model name is required - check_unset(&dimensions, "dimensions", inferred_source, name)?; - check_set(&model, "model", inferred_source, name)?; - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; + check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; + check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::HuggingFace => { - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&dimensions, "dimensions", inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::UserProvided => { - check_unset(&model, "model", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&document_template, "documentTemplate", inferred_source, name)?; - check_set(&dimensions, "dimensions", inferred_source, name)?; + check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset( + &document_template, + EmbeddingSettings::DOCUMENT_TEMPLATE, + inferred_source, + name, + )?; + check_set(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::Rest => { - check_unset(&model, "model", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; - check_set(&url, "url", inferred_source, name)?; + check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; + check_set(&url, EmbeddingSettings::URL, inferred_source, name)?; } } Ok(Setting::Set(EmbeddingSettings { From 817ccc089a1defa4e88db4494e5d9bdfb38d13e8 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:50:00 +0100 Subject: [PATCH 3/3] also allow `api_key` --- milli/src/update/settings.rs | 1 - milli/src/vector/mod.rs | 4 ++-- milli/src/vector/ollama.rs | 7 ++++--- milli/src/vector/settings.rs | 15 ++++++++------- 4 files changed, 14 insertions(+), 13 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index e0c559b85..2b1be9453 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1278,7 +1278,6 @@ pub fn validate_embedding_settings( // Dimensions get inferred, only model name is required check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?; - check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 8186e4409..8b25de56d 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -201,8 +201,8 @@ impl EmbedderOptions { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } - pub fn ollama(url: Option) -> Self { - Self::Ollama(ollama::EmbedderOptions::with_default_model(url)) + pub fn ollama(api_key: Option, url: Option) -> Self { + Self::Ollama(ollama::EmbedderOptions::with_default_model(api_key, url)) } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 7edfd13b5..578b6c8e2 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -13,11 +13,12 @@ pub struct Embedder { pub struct EmbedderOptions { pub embedding_model: String, pub url: Option, + pub api_key: Option, } impl EmbedderOptions { - pub fn with_default_model(url: Option) -> Self { - Self { embedding_model: "nomic-embed-text".into(), url } + pub fn with_default_model(api_key: Option, url: Option) -> Self { + Self { embedding_model: "nomic-embed-text".into(), api_key, url } } } @@ -25,7 +26,7 @@ impl Embedder { pub fn new(options: EmbedderOptions) -> Result { let model = options.embedding_model.as_str(); let rest_embedder = match RestEmbedder::new(RestEmbedderOptions { - api_key: None, + api_key: options.api_key, distribution: None, dimensions: None, url: options.url.unwrap_or_else(get_ollama_path), diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 7760573f6..c277dd0cf 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -114,7 +114,9 @@ impl EmbeddingSettings { &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] } Self::REVISION => &[EmbedderSource::HuggingFace], - Self::API_KEY => &[EmbedderSource::OpenAi, EmbedderSource::Rest], + Self::API_KEY => { + &[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest] + } Self::DIMENSIONS => { &[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest] } @@ -147,7 +149,7 @@ impl EmbeddingSettings { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } EmbedderSource::Ollama => { - &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL] + &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL, Self::API_KEY] } EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], EmbedderSource::Rest => &[ @@ -389,15 +391,14 @@ impl From for EmbeddingConfig { } EmbedderSource::Ollama => { let mut options: ollama::EmbedderOptions = - super::ollama::EmbedderOptions::with_default_model(None); + super::ollama::EmbedderOptions::with_default_model( + api_key.set(), + url.set(), + ); if let Some(model) = model.set() { options.embedding_model = model; } - if let Some(url) = url.set() { - options.url = Some(url) - } - this.embedder_options = super::EmbedderOptions::Ollama(options); } EmbedderSource::HuggingFace => {