From 58972f35cb0d755fc9fc45de110beedc78565602 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Mon, 25 Mar 2024 11:13:21 +0100
Subject: [PATCH] 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<String>) -> 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<String>,
 }
 
 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<String>) -> 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<EmbeddingSettings> 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 => {