Adds chat and chat example

src/lib.rs

@@ -10,6 +10,7 @@ use tracing::info;
 use crate::{
     error::{OllamaError, OllamaResult},
     types::{
+        chat::{ChatRequest, ChatResponse},
         generate::{GenerateRequest, GenerateResponse},
         ps::RunningModel,
         tags::Model,
@@ -113,4 +114,44 @@ impl OllamaClient {
             }
         })
     }
+
+    /// Generate the next chat message in a conversation between a user and an assistant.
+    pub async fn chat(
+        &self,
+        request: ChatRequest,
+    ) -> impl Stream<Item = OllamaResult<ChatResponse>> {
+        let request_address = format!("{}/api/chat", self.server_address);
+        let client = reqwest::Client::new();
+
+        // The stream macro creates an asynchronous generator
+        Box::pin(stream! {
+            let response = client
+                .post(request_address)
+                .json(&request)
+                .send()
+                .await
+                .map_err(|e| OllamaError::from(e))?; // Adjust based on your error type
+
+            let bytes_stream = response.bytes_stream();
+
+            let body_reader = StreamReader::new(
+                bytes_stream.map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))),
+            );
+
+            let mut lines_stream = FramedRead::new(body_reader, LinesCodec::new());
+
+            while let Some(line_result) = lines_stream.next().await {
+                match line_result {
+                    Ok(line_content) => {
+                        if let Ok(parsed) = serde_json::from_str::<ChatResponse>(&line_content) {
+                            let done = parsed.done;
+                            yield Ok(parsed);
+                            if done { break; }
+                        }
+                    }
+                    Err(e) => yield Err(OllamaError::from(e)),
+                }
+            }
+        })
+    }
 }

src/types/chat.rs

@@ -0,0 +1,61 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    User,
+    System,
+    Assistant,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Message {
+    pub content: String,
+    pub role: Role,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ChatRequest {
+    pub model: String,
+    pub messages: Vec<Message>,
+    pub stream: Option<bool>,
+}
+
+impl ChatRequest {
+    pub fn builder<M: Into<String>>(model: M) -> ChatRequestBuilder {
+        ChatRequestBuilder::new(model)
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ChatResponse {
+    pub model: String,
+    pub created_at: String,
+    pub message: Message,
+    pub done: bool,
+}
+
+pub struct ChatRequestBuilder {
+    chat_request: ChatRequest,
+}
+
+impl ChatRequestBuilder {
+    fn new<M: Into<String>>(model: M) -> Self {
+        Self {
+            chat_request: ChatRequest {
+                model: model.into(),
+                messages: vec![],
+                stream: None,
+            },
+        }
+    }
+
+    pub fn messages(mut self, messages: Vec<Message>) -> Self {
+        self.chat_request.messages = messages;
+        self
+    }
+
+    pub fn build(self) -> ChatRequest {
+        self.chat_request
+    }
+}

src/types/generate.rs

@@ -17,6 +17,10 @@ pub struct GenerateRequest {
     /// System prompt for the model to generate a response from
     #[serde(skip_serializing_if = "Option::is_none")]
     pub system: Option<String>,
+
+    /// When true, returns a stream of partial responses
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream: Option<bool>,
 }
 
 impl GenerateRequest {
@@ -37,6 +41,7 @@ impl GenerateRequestBuilder {
                 prompt: None,
                 suffix: None,
                 system: None,
+                stream: None,
             },
         }
     }
@@ -51,6 +56,11 @@ impl GenerateRequestBuilder {
         self
     }
 
+    pub fn stream(mut self, stream: bool) -> Self {
+        self.generate_request.stream = Some(stream);
+        self
+    }
+
     pub fn build(self) -> GenerateRequest {
         self.generate_request
     }