Merge pull request #10 from pepperoni21/images-support

Added images support for multimodal models

Cargo.toml

@@ -24,3 +24,4 @@ rustls = ["reqwest/rustls-tls"]
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }
 ollama-rs = { path = ".", features = ["stream"] }
+base64 = "0.21.5"

src/generation.rs

@@ -2,4 +2,5 @@ pub mod chat;
 pub mod completion;
 pub mod embeddings;
 pub mod format;
+pub mod images;
 pub mod options;

src/generation/chat/mod.rs

@@ -6,6 +6,8 @@ pub mod request;
 
 use request::ChatMessageRequest;
 
+use super::images::Image;
+
 #[cfg(feature = "stream")]
 /// A stream of `ChatMessageResponse` objects
 pub type ChatMessageResponseStream =
@@ -123,11 +125,16 @@ pub struct ChatMessageFinalResponseData {
 pub struct ChatMessage {
     pub role: MessageRole,
     pub content: String,
+    pub images: Option<Vec<Image>>,
 }
 
 impl ChatMessage {
     pub fn new(role: MessageRole, content: String) -> Self {
-        Self { role, content }
+        Self {
+            role,
+            content,
+            images: None,
+        }
     }
 
     pub fn user(content: String) -> Self {
@@ -141,6 +148,20 @@ impl ChatMessage {
     pub fn system(content: String) -> Self {
         Self::new(MessageRole::System, content)
     }
+
+    pub fn with_images(mut self, images: Vec<Image>) -> Self {
+        self.images = Some(images);
+        self
+    }
+
+    pub fn add_image(mut self, image: Image) -> Self {
+        if let Some(images) = self.images.as_mut() {
+            images.push(image);
+        } else {
+            self.images = Some(vec![image]);
+        }
+        self
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

src/generation/completion/request.rs

@@ -1,6 +1,6 @@
 use serde::Serialize;
 
-use crate::generation::{format::FormatType, options::GenerationOptions};
+use crate::generation::{format::FormatType, images::Image, options::GenerationOptions};
 
 use super::GenerationContext;
 
@@ -10,6 +10,7 @@ pub struct GenerationRequest {
     #[serde(rename = "model")]
     pub model_name: String,
     pub prompt: String,
+    pub images: Vec<Image>,
     pub options: Option<GenerationOptions>,
     pub system: Option<String>,
     pub template: Option<String>,
@@ -23,6 +24,7 @@ impl GenerationRequest {
         Self {
             model_name,
             prompt,
+            images: Vec::new(),
             options: None,
             system: None,
             template: None,
@@ -33,6 +35,18 @@ impl GenerationRequest {
         }
     }
 
+    /// A list of images to be used with the prompt
+    pub fn images(mut self, images: Vec<Image>) -> Self {
+        self.images = images;
+        self
+    }
+
+    /// Add an image to be used with the prompt
+    pub fn add_image(mut self, image: Image) -> Self {
+        self.images.push(image);
+        self
+    }
+
     /// Additional model parameters listed in the documentation for the Modelfile
     pub fn options(mut self, options: GenerationOptions) -> Self {
         self.options = Some(options);

src/generation/images.rs

@@ -0,0 +1,10 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Image(String);
+
+impl Image {
+    pub fn from_base64(base64: &str) -> Self {
+        Self(base64.to_string())
+    }
+}

tests/generation.rs

@@ -1,7 +1,10 @@
 #![allow(unused_imports)]
-
+use base64::Engine;
 use ollama_rs::{
-    generation::completion::{request::GenerationRequest, GenerationResponseStream},
+    generation::{
+        completion::{request::GenerationRequest, GenerationResponseStream},
+        images::Image,
+    },
     Ollama,
 };
 use tokio::io::AsyncWriteExt;
@@ -48,3 +51,31 @@ async fn test_generation() {
         .unwrap();
     dbg!(res);
 }
+
+const IMAGE_URL: &str = "https://images.pexels.com/photos/1054655/pexels-photo-1054655.jpeg";
+
+#[tokio::test]
+async fn test_generation_with_images() {
+    let ollama = Ollama::default();
+
+    let bytes = reqwest::get(IMAGE_URL)
+        .await
+        .unwrap()
+        .bytes()
+        .await
+        .unwrap();
+    let base64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
+    let image = Image::from_base64(&base64);
+
+    let res = ollama
+        .generate(
+            GenerationRequest::new(
+                "llava:latest".to_string(),
+                "What can we see in this image?".to_string(),
+            )
+            .add_image(image),
+        )
+        .await
+        .unwrap();
+    dbg!(res);
+}

tests/send_chat_messages.rs

@@ -1,5 +1,9 @@
+use base64::Engine;
 use ollama_rs::{
-    generation::chat::{request::ChatMessageRequest, ChatMessage},
+    generation::{
+        chat::{request::ChatMessageRequest, ChatMessage},
+        images::Image,
+    },
     Ollama,
 };
 use tokio_stream::StreamExt;
@@ -49,3 +53,32 @@ async fn test_send_chat_messages() {
 
     assert!(res.done);
 }
+
+const IMAGE_URL: &str = "https://images.pexels.com/photos/1054655/pexels-photo-1054655.jpeg";
+
+#[tokio::test]
+async fn test_send_chat_messages_with_images() {
+    let ollama = Ollama::default();
+
+    let bytes = reqwest::get(IMAGE_URL)
+        .await
+        .unwrap()
+        .bytes()
+        .await
+        .unwrap();
+    let base64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
+    let image = Image::from_base64(&base64);
+
+    let messages =
+        vec![ChatMessage::user("What can we see in this image?".to_string()).add_image(image)];
+    let res = ollama
+        .send_chat_messages(ChatMessageRequest::new(
+            "llava:latest".to_string(),
+            messages,
+        ))
+        .await
+        .unwrap();
+    dbg!(&res);
+
+    assert!(res.done);
+}