Merge pull request #15 from raphaelmansuy:feat/ollamaoptions

Feat/ollamaoptions

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hiramu"
-version = "0.1.14"
+version = "0.1.15"
 edition = "2021"
 license = "MIT"
 description = "A Rust AI Engineering Toolbox to Access Ollama, AWS Bedrock"

README.md

@@ -20,7 +20,7 @@ To start using Hiramu in your Rust project, add the following to your `Cargo.tom
 
 ```toml
 [dependencies]
-hiramu = "0.1.14"
+hiramu = "0.1.15"
 ```
 
 ## Examples
@@ -101,6 +101,7 @@ use hiramu::ollama::model::{GenerateRequestBuilder};
 
 async fn generating_text_with_ollama() {
     let client = OllamaClient::new("http://localhost:11434".to_string());
+
     let request = GenerateRequestBuilder::new("mistral".to_string())
         .prompt("Once upon a time".to_string())
         .build();
@@ -118,6 +119,49 @@ async fn generating_text_with_ollama() {
 }
 ```
 
+### Chat with Ollama
+```rust
+
+use futures::TryStreamExt;
+use std::io::{self, Write};
+
+use hiramu::ollama::{ChatRequestBuilder, Message, OllamaClient, OllamaError, OptionsBuilder};
+
+async fn demo_chat_with_ollama_with_stream() -> Result<(), OllamaError> {
+    let client = OllamaClient::new("http://localhost:11434".to_string());
+
+    let messages = vec![Message::new(
+        "user".to_string(),
+        "What is the capital of France?  "
+            .to_string(),
+    )];
+
+    let options = OptionsBuilder::new()
+        .num_predict(100) // Limit the number of predicted tokens
+        .temperature(0.4);
+
+    let request = ChatRequestBuilder::new("mistral".to_string())
+        .messages(messages.to_owned())
+        .options_from_builder(options)
+        .build();
+
+    let response_stream = client.chat(request).await?;
+
+    let result = response_stream
+        .try_for_each(|chunk| async {
+            let message = chunk.message;
+            print!("{}", message.content);
+            // Flush the output to ensure the prompt is displayed.
+            io::stdout().flush().unwrap();
+            Ok(())
+        })
+        .await;
+
+    result
+}
+
+```
+
 ### Chatting with Claude using Bedrock
 
 ```rust
@@ -364,7 +408,8 @@ Here is a table with a description for each example:
 
 | Example | Path | Description |
 |---------|------|--------------|
-| `demo_ollama` | [src/examples/demo_ollama.rs](src/examples/demo_ollama.rs) | A simple example that demonstrates how to use the Ollama API to generate responses to chat messages. |
+| `demo_ollama` | [src/examples/demo_ollama.rs](src/examples/demo_ollama.rs) | A simple example that demonstrates how to use the Ollama API to generate responses. |
+| `demo_chat_with_ollama` | [src/examples/demo_chat_with_ollama.rs](src/examples/demo_chat_with_ollama.rs) | A simple example that demonstrates how to use the Ollama Chat API. |
 | `demo_bedrock_raw_generate` | [src/examples/demo_bedrock_raw_generate.rs](src/examples/demo_bedrock_raw_generate.rs) | Demonstrates how to generate a raw response from the Bedrock service using the `generate_raw` method. |
 | `demo_bedrock_raw_stream` | [src/examples/demo_bedrock_raw_stream.rs](src/examples/demo_bedrock_raw_stream.rs) | Demonstrates how to generate a raw stream of responses from the Bedrock service using the `generate_raw_stream` method. |
 | `demo_bedrock_raw_mistral` | [src/examples/demo_bedrock_raw_mistral.rs](src/examples/demo_bedrock_raw_mistral.rs) | Demonstrates how to generate a raw stream of responses from the Mistral model in the Bedrock service. |

REVISION.md

@@ -4,3 +4,13 @@
 
 ### 0.1.7
 - Added support for the Mistral API in the Bedrock module.
+
+## 0.1.15
+
+Implement new generate_text method in OllamaClient and add options_builder support to GenerateRequestBuilder and ChatRequestBuilder
+
+- Added a new method `generate_text` to the `OllamaClient` that collects the stream of responses from the `generate` method into a single string.
+- Implemented the `merge_options` function to merge the options provided in the request builder and the options builder.
+- Added the `options_builder` field to the `GenerateRequestBuilder` and `ChatRequestBuilder` to allow setting options using an `OptionsBuilder`.
+- Added unit tests for the `merge_options` function.
+- Added a new test case for the `generate_text` method to ensure the maximum number of predictions is respected.

src/examples/demo_chat_with_ollama.rs

@@ -0,0 +1,47 @@
+use futures::TryStreamExt;
+use std::io::{self, Write};
+
+use crate::ollama::{ChatRequestBuilder, Message, OllamaClient, OllamaError, OptionsBuilder};
+
+pub async fn demo_chat_with_ollama_with_stream() -> Result<(), OllamaError> {
+    let client = OllamaClient::new("http://localhost:11434".to_string());
+
+    let messages = vec![Message::new(
+        "user".to_string(),
+        "What is the capital of France?  "
+            .to_string(),
+    )];
+
+    let options = OptionsBuilder::new()
+        .num_predict(100) // Limit the number of predicted tokens
+        .temperature(0.4);
+
+    let request = ChatRequestBuilder::new("mistral".to_string())
+        .messages(messages.to_owned())
+        .options_from_builder(options)
+        .build();
+
+    let response_stream = client.chat(request).await?;
+
+    let result = response_stream
+        .try_for_each(|chunk| async {
+            let message = chunk.message;
+            print!("{}", message.content);
+            // Flush the output to ensure the prompt is displayed.
+            io::stdout().flush().unwrap();
+            Ok(())
+        })
+        .await;
+
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_demo_chat_with_ollama_with_stream() {
+        let _ = demo_chat_with_ollama_with_stream().await;
+    }
+}

src/examples/demo_claude_chat_stream.rs

@@ -71,12 +71,6 @@ fn display_streamresult_data(data: StreamResultData) {
         StreamResultData::MessageDelta(message_delta) => {
             println!("MessageDelta: {:?}", message_delta);
         }
-        StreamResultData::ContentBlockStart(content_block_start) => {
-            println!("ContentBlockStart: {:?}", content_block_start);
-        }
-        StreamResultData::ContentBlockStop(content_block_end) => {
-            println!("ContentBlockEnd: {:?}", content_block_end);
-        }
         StreamResultData::ContentBlockDelta(content_block_delta) => {
             println!("ContentBlockDelta: {:?}", content_block_delta);
         }

src/examples/demo_mistral_stream.rs

@@ -1,6 +1,6 @@
 use futures::TryStreamExt;
 
-use crate::bedrock::model_info::{ModelInfo, ModelName};
+//use crate::bedrock::model_info::{ModelInfo, ModelName};
 use crate::bedrock::models::mistral::MistralClient;
 use crate::bedrock::models::mistral::MistralOptions;
 use crate::bedrock::models::mistral::MistralRequestBuilder;
@@ -55,6 +55,7 @@ pub async fn demo_mistra_with_stream(model_id: &str, prompt: &str) {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::bedrock::model_info::{ModelInfo, ModelName};
 
     #[tokio::test]
     async fn test_demo_chat_mistral_with_stream() {

src/examples/demo_ollama.rs

@@ -4,20 +4,26 @@ use std::u32;
 use futures::stream::TryStream;
 use futures_util::TryStreamExt;
 
+use crate::ollama::options::OptionsBuilder;
 use crate::ollama::{ChatRequestBuilder, ChatResponse, Message};
 use crate::ollama::{GenerateRequestBuilder, GenerateResponse};
 use crate::ollama::OllamaClient;
 use crate::ollama::OllamaError;
 
 
-pub async fn chat_response_loop(max_loop: u32) {
+pub async fn chat_response_loop(max_loop: u32,question: Option<&str>) {
     let client = OllamaClient::new("http://localhost:11434".to_string());
 
     let mut messages = Vec::new();
     let mut counter = 0;
 
     loop {
-        let input = prompt_input("\nUser: ").unwrap();
+        // if question is provided, use it as the prompt
+        let input = match question {
+            Some(q) => q.to_string(),
+            None => prompt_input("\n> ").unwrap(),
+        };
+
         messages.push(Message {
             role: "user".to_string(),
             content: input,
@@ -51,14 +57,24 @@ pub async fn chat_response_loop(max_loop: u32) {
     }
 }
 
-pub async fn generate_response_loop(max_loop: usize) {
+pub async fn generate_response_loop(max_loop: usize, question: Option<&str>) {
     let client = OllamaClient::new("http://localhost:11434".to_string());
 
     let mut counter = 0;
     loop {
-        let input = prompt_input("\n> ").unwrap();
+        // if question is provided, use it as the prompt
+        let input = match question {
+            Some(q) => q.to_string(),
+            None => prompt_input("\n> ").unwrap(),
+        };
         let request = GenerateRequestBuilder::new("mistral".to_string())
             .prompt(input)
+            .options_from_builder(
+                OptionsBuilder::new()
+                    .top_k(50)
+                    .top_p(0.9)
+                    .temperature(0.9)
+            )
             .build();
 
         let response = client.generate(request).await.unwrap();
@@ -119,13 +135,15 @@ pub async fn print_generate_response(
 #[cfg(test)]
 mod tests {
 
+    use super::*;
+
     #[tokio::test]
     async fn test_chat_response_loop() {
-        // chat_response_loop(1).await;
+         chat_response_loop(1, Some("What is the capital of France ?")).await;
     }
 
     #[tokio::test]
     async fn test_generate_response_loop() {
-        //      generate_response_loop(1).await;
+        generate_response_loop(1,Some("Quelle est la couleur du cheval blanc d'Henri IV ?")).await;
     }
 }

src/examples/mod.rs

@@ -1,11 +1,12 @@
 pub mod simple_examples;
 
 pub mod demo_ollama;
+pub mod demo_chat_with_ollama;
+pub mod demo_ollama_embedding;
 pub mod demo_bedrock_raw_generate;
 pub mod demo_bedrock_raw_stream;
 pub mod demo_bedrock_raw_mistral;
 pub mod demo_claude_chat;
 pub mod demo_claude_chat_stream;
 pub mod demo_claude_multimedia;
-pub mod demo_ollama_embedding;
 pub mod demo_mistral_stream;

src/main.rs

@@ -3,5 +3,5 @@ use hiramu::examples::demo_ollama::chat_response_loop;
 #[tokio::main]
 async fn main() {
     // A simple example that demonstrates how to use the Ollama API to generate responses to chat messages.
-    chat_response_loop(u32::max_value()).await;
+    chat_response_loop(u32::max_value(),None).await;
 }

src/ollama/mod.rs

@@ -1,9 +1,11 @@
 pub mod ollama_client;
 pub mod model;
 pub mod error;
+pub mod options;
 
 pub use error::OllamaError;
 pub use ollama_client::OllamaClient;
 pub use model::{ GenerateRequest, GenerateRequestBuilder, GenerateResponse };
 pub use model::{ ChatRequest, ChatRequestBuilder, ChatResponse, Message };
-pub use model::{ EmbeddingsRequest,EmbeddingsResponse, EmbeddingsRequestBuilder};
+pub use model::{ EmbeddingsRequest,EmbeddingsResponse, EmbeddingsRequestBuilder};
+pub use options::OptionsBuilder;