Update embeddings generation to /api/embed endpoint and allow for batch embedding

src/generation/embeddings.rs → src/generation/embeddings/mod.rs

@@ -1,26 +1,20 @@
-use serde::{Deserialize, Serialize};
+use serde::Deserialize;
 
 use crate::Ollama;
 
-use super::options::GenerationOptions;
+use self::request::GenerateEmbeddingsRequest;
+
+pub mod request;
 
 impl Ollama {
     /// Generate embeddings from a model
     /// * `model_name` - Name of model to generate embeddings from
     /// * `prompt` - Prompt to generate embeddings for
     pub async fn generate_embeddings(
         &self,
-        model_name: String,
-        prompt: String,
-        options: Option<GenerationOptions>,
+        request: GenerateEmbeddingsRequest,
     ) -> crate::error::Result<GenerateEmbeddingsResponse> {
-        let request = GenerateEmbeddingsRequest {
-            model_name,
-            prompt,
-            options,
-        };
-
-        let url = format!("{}api/embeddings", self.url_str());
+        let url = format!("{}api/embed", self.url_str());
         let serialized = serde_json::to_string(&request).map_err(|e| e.to_string())?;
         let res = self
             .reqwest_client
@@ -42,19 +36,9 @@ impl Ollama {
     }
 }
 
-/// An embeddings generation request to Ollama.
-#[derive(Debug, Serialize)]
-struct GenerateEmbeddingsRequest {
-    #[serde(rename = "model")]
-    model_name: String,
-    prompt: String,
-    options: Option<GenerationOptions>,
-}
-
 /// An embeddings generation response from Ollama.
 #[derive(Debug, Deserialize, Clone)]
 pub struct GenerateEmbeddingsResponse {
-    #[serde(rename = "embedding")]
     #[allow(dead_code)]
-    pub embeddings: Vec<f64>,
+    pub embeddings: Vec<Vec<f64>>,
 }

src/generation/embeddings/request.rs

@@ -0,0 +1,84 @@
+use serde::{Serialize, Serializer};
+
+use crate::generation::{options::GenerationOptions, parameters::KeepAlive};
+
+#[derive(Debug)]
+pub enum EmbeddingsInput {
+    Single(String),
+    Multiple(Vec<String>),
+}
+
+impl Default for EmbeddingsInput {
+    fn default() -> Self {
+        Self::Single(String::default())
+    }
+}
+
+impl From<String> for EmbeddingsInput {
+    fn from(s: String) -> Self {
+        Self::Single(s)
+    }
+}
+
+impl From<&str> for EmbeddingsInput {
+    fn from(s: &str) -> Self {
+        Self::Single(s.to_string())
+    }
+}
+
+impl From<Vec<String>> for EmbeddingsInput {
+    fn from(v: Vec<String>) -> Self {
+        Self::Multiple(v)
+    }
+}
+
+impl From<Vec<&str>> for EmbeddingsInput {
+    fn from(v: Vec<&str>) -> Self {
+        Self::Multiple(v.iter().map(|s| s.to_string()).collect())
+    }
+}
+
+impl Serialize for EmbeddingsInput {
+    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+        match self {
+            EmbeddingsInput::Single(s) => s.serialize(serializer),
+            EmbeddingsInput::Multiple(v) => v.serialize(serializer),
+        }
+    }
+}
+
+/// An embeddings generation request to Ollama.
+#[derive(Debug, Serialize, Default)]
+pub struct GenerateEmbeddingsRequest {
+    #[serde(rename = "model")]
+    model_name: String,
+    input: EmbeddingsInput,
+    truncate: Option<bool>,
+    options: Option<GenerationOptions>,
+    keep_alive: Option<KeepAlive>,
+}
+
+impl GenerateEmbeddingsRequest {
+    pub fn new(model_name: String, input: EmbeddingsInput) -> Self {
+        Self {
+            model_name,
+            input,
+            ..Default::default()
+        }
+    }
+
+    pub fn options(mut self, options: GenerationOptions) -> Self {
+        self.options = Some(options);
+        self
+    }
+
+    pub fn keep_alive(mut self, keep_alive: KeepAlive) -> Self {
+        self.keep_alive = Some(keep_alive);
+        self
+    }
+
+    pub fn truncate(mut self, truncate: bool) -> Self {
+        self.truncate = Some(truncate);
+        self
+    }
+}

tests/embeddings_generation.rs

@@ -1,13 +1,29 @@
-use ollama_rs::Ollama;
+use ollama_rs::{generation::embeddings::request::GenerateEmbeddingsRequest, Ollama};
 
 #[tokio::test]
 async fn test_embeddings_generation() {
     let ollama = Ollama::default();
 
-    let prompt = "Why is the sky blue?".to_string();
+    let res = ollama
+        .generate_embeddings(GenerateEmbeddingsRequest::new(
+            "llama2:latest".to_string(),
+            "Why is the sky blue".into(),
+        ))
+        .await
+        .unwrap();
+
+    dbg!(res);
+}
+
+#[tokio::test]
+async fn test_batch_embeddings_generation() {
+    let ollama = Ollama::default();
 
     let res = ollama
-        .generate_embeddings("llama2:latest".to_string(), prompt, None)
+        .generate_embeddings(GenerateEmbeddingsRequest::new(
+            "llama2:latest".to_string(),
+            vec!["Why is the sky blue?", "Why is the sky red?"].into(),
+        ))
         .await
         .unwrap();