From f44d95124767c3a3f14c78c4be3d9c203fac49ad Mon Sep 17 00:00:00 2001 From: Ivan Gabriele Date: Mon, 4 Mar 2024 03:14:23 +0100 Subject: [PATCH] feat!: add client.embeddings() method BREAKING CHANGE: Models are now enforced by `Model` & `EmbedModel` enums. --- README.md | 26 +++++++++++--- src/v1/chat_completion.rs | 8 ++--- src/v1/client.rs | 48 ++++++++++++------------- src/v1/constants.rs | 26 +++++++++++--- src/v1/embedding.rs | 60 ++++++++++++++++++++++++++++++++ src/v1/mod.rs | 1 + tests/v1_chat_completion_test.rs | 8 ++--- tests/v1_embeddings_test.rs | 31 +++++++++++++++++ tests/v1_list_models_test.rs | 8 +++++ 9 files changed, 173 insertions(+), 43 deletions(-) create mode 100644 src/v1/embedding.rs create mode 100644 tests/v1_embeddings_test.rs diff --git a/README.md b/README.md index 0cd2baf..088bee0 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Rust client for the Mistral AI API. - [x] Chat without streaming - [ ] Chat with streaming -- [ ] Embedding +- [x] Embedding - [x] List models - [ ] Function Calling @@ -66,14 +66,14 @@ fn main() { use mistralai_client::v1::{ chat_completion::{ChatCompletionMessage, ChatCompletionMessageRole, ChatCompletionRequestOptions}, client::Client, - constants::OPEN_MISTRAL_7B, + constants::Model, }; fn main() { // This example suppose you have set the `MISTRAL_API_KEY` environment variable. let client = Client::new(None, None, None, None); - let model = OPEN_MISTRAL_7B.to_string(); + let model = Model::OpenMistral7b; let messages = vec![ChatCompletionMessage { role: ChatCompletionMessageRole::user, content: "Just guess the next word: \"Eiffel ...\"?".to_string(), @@ -96,7 +96,25 @@ _In progress._ ### Embeddings -_In progress._ +```rs +use mistralai_client::v1::{client::Client, constants::EmbedModel}; + +fn main() { + // This example suppose you have set the `MISTRAL_API_KEY` environment variable. + let client: Client = Client::new(None, None, None, None); + + let model = EmbedModel::MistralEmbed; + let input = vec!["Embed this sentence.", "As well as this one."] + .iter() + .map(|s| s.to_string()) + .collect(); + let options = None; + + let response = client.embeddings(model, input, options).unwrap(); + println!("Embeddings: {:?}", response.data); + // => "Embeddings: [{...}, {...}]" +} +``` ### List models diff --git a/src/v1/chat_completion.rs b/src/v1/chat_completion.rs index 998677c..6349219 100644 --- a/src/v1/chat_completion.rs +++ b/src/v1/chat_completion.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::v1::common; +use crate::v1::{common, constants}; #[derive(Debug)] pub struct ChatCompletionParams { @@ -29,7 +29,7 @@ impl Default for ChatCompletionParams { #[derive(Debug, Serialize, Deserialize)] pub struct ChatCompletionRequest { pub messages: Vec, - pub model: String, + pub model: constants::Model, #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -51,7 +51,7 @@ pub struct ChatCompletionRequest { } impl ChatCompletionRequest { pub fn new( - model: String, + model: constants::Model, messages: Vec, options: Option, ) -> Self { @@ -85,7 +85,7 @@ pub struct ChatCompletionResponse { pub object: String, /// Unix timestamp (in seconds). pub created: u32, - pub model: String, + pub model: constants::Model, pub choices: Vec, pub usage: common::ResponseUsage, } diff --git a/src/v1/client.rs b/src/v1/client.rs index 4de4e3c..f454bd2 100644 --- a/src/v1/client.rs +++ b/src/v1/client.rs @@ -5,7 +5,8 @@ use crate::v1::{ chat_completion::{ ChatCompletionMessage, ChatCompletionParams, ChatCompletionRequest, ChatCompletionResponse, }, - constants::API_URL_BASE, + constants::{EmbedModel, Model, API_URL_BASE}, + embedding::{EmbeddingRequest, EmbeddingRequestOptions, EmbeddingResponse}, model_list::ModelListResponse, }; @@ -59,6 +60,8 @@ impl Client { let result = request.send(); match result { Ok(response) => { + print!("{:?}", response.as_str().unwrap()); + if (200..=299).contains(&response.status_code) { Ok(response) } else { @@ -88,7 +91,7 @@ impl Client { let result = request.with_json(params).unwrap().send(); match result { Ok(response) => { - // print!("{:?}", response.as_str().unwrap()); + print!("{:?}", response.as_str().unwrap()); if (200..=299).contains(&response.status_code) { Ok(response) @@ -106,32 +109,9 @@ impl Client { } } - pub fn delete(&self, path: &str) -> Result { - let url = format!("{}{}", self.endpoint, path); - let request = self.build_request(minreq::post(url)); - - let result = request.send(); - match result { - Ok(response) => { - if (200..=299).contains(&response.status_code) { - Ok(response) - } else { - Err(APIError { - message: format!( - "{}: {}", - response.status_code, - response.as_str().unwrap() - ), - }) - } - } - Err(error) => Err(self.new_error(error)), - } - } - pub fn chat( &self, - model: String, + model: Model, messages: Vec, options: Option, ) -> Result { @@ -145,6 +125,22 @@ impl Client { } } + pub fn embeddings( + &self, + model: EmbedModel, + input: Vec, + options: Option, + ) -> Result { + let request = EmbeddingRequest::new(model, input, options); + + let response = self.post("/embeddings", &request)?; + let result = response.json::(); + match result { + Ok(response) => Ok(response), + Err(error) => Err(self.new_error(error)), + } + } + pub fn list_models(&self) -> Result { let response = self.get("/models")?; let result = response.json::(); diff --git a/src/v1/constants.rs b/src/v1/constants.rs index 1cd4c85..c0bd577 100644 --- a/src/v1/constants.rs +++ b/src/v1/constants.rs @@ -1,7 +1,23 @@ +use serde::{Deserialize, Serialize}; + pub const API_URL_BASE: &str = "https://api.mistral.ai/v1"; -pub const OPEN_MISTRAL_7B: &str = "open-mistral-7b"; -pub const OPEN_MISTRAL_8X7B: &str = "open-mixtral-8x7b"; -pub const MISTRAL_SMALL_LATEST: &str = "mistral-small-latest"; -pub const MISTRAL_MEDIUM_LATEST: &str = "mistral-medium-latest"; -pub const MISTRAL_LARGE_LATEST: &str = "mistral-large-latest"; +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)] +pub enum Model { + #[serde(rename = "open-mistral-7b")] + OpenMistral7b, + #[serde(rename = "open-mistral-8x7b")] + OpenMistral8x7b, + #[serde(rename = "mistral-small-latest")] + MistralSmallLatest, + #[serde(rename = "mistral-medium-latest")] + MistralMediumLatest, + #[serde(rename = "mistral-large-latest")] + MistralLargeLatest, +} + +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)] +pub enum EmbedModel { + #[serde(rename = "mistral-embed")] + MistralEmbed, +} diff --git a/src/v1/embedding.rs b/src/v1/embedding.rs new file mode 100644 index 0000000..7edc719 --- /dev/null +++ b/src/v1/embedding.rs @@ -0,0 +1,60 @@ +use serde::{Deserialize, Serialize}; + +use crate::v1::{common, constants}; + +#[derive(Debug)] +pub struct EmbeddingRequestOptions { + pub encoding_format: Option, +} +impl Default for EmbeddingRequestOptions { + fn default() -> Self { + Self { + encoding_format: None, + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct EmbeddingRequest { + pub model: constants::EmbedModel, + pub input: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub encoding_format: Option, +} +impl EmbeddingRequest { + pub fn new( + model: constants::EmbedModel, + input: Vec, + options: Option, + ) -> Self { + let EmbeddingRequestOptions { encoding_format } = options.unwrap_or_default(); + + Self { + model, + input, + encoding_format, + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)] +#[allow(non_camel_case_types)] +pub enum EmbeddingRequestEncodingFormat { + float, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct EmbeddingResponse { + pub id: String, + pub object: String, + pub model: constants::EmbedModel, + pub data: Vec, + pub usage: common::ResponseUsage, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct EmbeddingResponseDataItem { + pub index: u32, + pub embedding: Vec, + pub object: String, +} diff --git a/src/v1/mod.rs b/src/v1/mod.rs index 8127a48..8f29a01 100644 --- a/src/v1/mod.rs +++ b/src/v1/mod.rs @@ -2,5 +2,6 @@ pub mod chat_completion; pub mod client; pub mod common; pub mod constants; +pub mod embedding; pub mod error; pub mod model_list; diff --git a/tests/v1_chat_completion_test.rs b/tests/v1_chat_completion_test.rs index e4eab28..60133f6 100644 --- a/tests/v1_chat_completion_test.rs +++ b/tests/v1_chat_completion_test.rs @@ -2,7 +2,7 @@ use jrest::expect; use mistralai_client::v1::{ chat_completion::{ChatCompletionMessage, ChatCompletionMessageRole, ChatCompletionParams}, client::Client, - constants::OPEN_MISTRAL_7B, + constants::Model, }; #[test] @@ -14,7 +14,7 @@ fn test_chat_completion() { let client = Client::new(None, None, None, None); - let model = OPEN_MISTRAL_7B.to_string(); + let model = Model::OpenMistral7b; let messages = vec![ChatCompletionMessage { role: ChatCompletionMessageRole::user, content: "Just guess the next word: \"Eiffel ...\"?".to_string(), @@ -27,7 +27,7 @@ fn test_chat_completion() { let response = client.chat(model, messages, Some(options)).unwrap(); - expect!(response.model).to_be("open-mistral-7b".to_string()); + expect!(response.model).to_be(Model::OpenMistral7b); expect!(response.object).to_be("chat.completion".to_string()); expect!(response.choices.len()).to_be(1); expect!(response.choices[0].index).to_be(0); @@ -36,5 +36,5 @@ fn test_chat_completion() { .to_be("Tower. The Eiffel Tower is a famous landmark in Paris, France.".to_string()); expect!(response.usage.prompt_tokens).to_be_greater_than(0); expect!(response.usage.completion_tokens).to_be_greater_than(0); - expect!(response.usage.total_tokens).to_be_greater_than(21); + expect!(response.usage.total_tokens).to_be_greater_than(0); } diff --git a/tests/v1_embeddings_test.rs b/tests/v1_embeddings_test.rs new file mode 100644 index 0000000..a0be427 --- /dev/null +++ b/tests/v1_embeddings_test.rs @@ -0,0 +1,31 @@ +use jrest::expect; +use mistralai_client::v1::{client::Client, constants::EmbedModel}; + +#[test] +fn test_embeddings() { + extern crate dotenv; + + use dotenv::dotenv; + dotenv().ok(); + + let client: Client = Client::new(None, None, None, None); + + let model = EmbedModel::MistralEmbed; + let input = vec!["Embed this sentence.", "As well as this one."] + .iter() + .map(|s| s.to_string()) + .collect(); + let options = None; + + let response = client.embeddings(model, input, options).unwrap(); + + expect!(response.model).to_be(EmbedModel::MistralEmbed); + expect!(response.object).to_be("list".to_string()); + expect!(response.data.len()).to_be(2); + expect!(response.data[0].index).to_be(0); + expect!(response.data[0].object.clone()).to_be("embedding".to_string()); + expect!(response.data[0].embedding.len()).to_be_greater_than(0); + expect!(response.usage.prompt_tokens).to_be_greater_than(0); + expect!(response.usage.completion_tokens).to_be(0); + expect!(response.usage.total_tokens).to_be_greater_than(0); +} diff --git a/tests/v1_list_models_test.rs b/tests/v1_list_models_test.rs index 13ec9fb..b2ea09e 100644 --- a/tests/v1_list_models_test.rs +++ b/tests/v1_list_models_test.rs @@ -14,4 +14,12 @@ fn test_list_models() { expect!(response.object).to_be("list".to_string()); expect!(response.data.len()).to_be_greater_than(0); + + // let open_mistral_7b_data_item = response + // .data + // .iter() + // .find(|item| item.id == "open-mistral-7b") + // .unwrap(); + + // expect!(open_mistral_7b_data_item.id).to_be("open-mistral-7b".to_string()); }