From 817ccc089a1defa4e88db4494e5d9bdfb38d13e8 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Mon, 25 Mar 2024 11:50:00 +0100
Subject: [PATCH] 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<String>) -> Self {
-        Self::Ollama(ollama::EmbedderOptions::with_default_model(url))
+    pub fn ollama(api_key: Option<String>, url: Option<String>) -> 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<String>,
+    pub api_key: Option<String>,
 }
 
 impl EmbedderOptions {
-    pub fn with_default_model(url: Option<String>) -> Self {
-        Self { embedding_model: "nomic-embed-text".into(), url }
+    pub fn with_default_model(api_key: Option<String>, url: Option<String>) -> Self {
+        Self { embedding_model: "nomic-embed-text".into(), api_key, url }
     }
 }
 
@@ -25,7 +26,7 @@ impl Embedder {
     pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
         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<EmbeddingSettings> 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 => {