feat!: add client.embeddings() method
BREAKING CHANGE: Models are now enforced by `Model` & `EmbedModel` enums.
This commit is contained in:
@@ -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);
|
||||
}
|
||||
|
||||
31
tests/v1_embeddings_test.rs
Normal file
31
tests/v1_embeddings_test.rs
Normal file
@@ -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);
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user