feat: add embeddings support

ollama/ollama-cli/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 anyhow = "1.0.87"
 clap = { version = "4.5.16", features = ["derive", "string"] }
 futures = "0.3.30"
-ollama-rs = { version = "0.2.0", features = ["stream"] }
+ollama-rs = { version = "0.2.1", features = ["stream"] }
 tokio = { version = "1.40.0", features = ["full"] }
 tokio-stream = { version = "0.1.16", features = ["full"] }
 tracing = "0.1.40"

ollama/ollama-cli/src/main.rs

@@ -14,6 +14,7 @@ pub struct Cli {
 #[derive(Subcommand)]
 enum Command {
     Generate(ollama::generate::Request),
+    Embed(ollama::generate::Request),
 }
 
 #[tokio::main]
@@ -26,6 +27,10 @@ async fn main() -> anyhow::Result<()> {
         Command::Generate(request) => {
             client.generate(request).await?;
         }
+        Command::Embed(request) => {
+            let embedding = client.embed(request).await?;
+            println!("{embedding:?}");
+        }
     }
 
     Ok(())

ollama/ollama-cli/src/ollama/embeddings.rs

@@ -0,0 +1,19 @@
+use ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest;
+
+use super::{generate::Request, Client};
+
+impl From<Request> for GenerateEmbeddingsRequest {
+    fn from(value: Request) -> Self {
+        GenerateEmbeddingsRequest::new(value.model.to_string(), value.prompt.into())
+    }
+}
+
+impl Client {
+    pub async fn embed(&self, request: Request) -> anyhow::Result<Vec<Vec<f32>>> {
+        let request: GenerateEmbeddingsRequest = request.into();
+
+        let response = self.client.generate_embeddings(request).await?;
+
+        Ok(response.embeddings)
+    }
+}

ollama/ollama-cli/src/ollama/mod.rs

@@ -4,6 +4,7 @@ use anyhow::anyhow;
 use ollama_rs::Ollama;
 use url::Url;
 
+pub mod embeddings;
 pub mod generate;
 
 pub const DEFAULT_MODEL: &str = "mistral-nemo";