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 => {