feat!: add client.embeddings() method

BREAKING CHANGE: Models are now enforced by `Model` & `EmbedModel` enums.
This commit is contained in:
Ivan Gabriele
2024-03-04 03:14:23 +01:00
parent 4e702aa48e
commit f44d951247
9 changed files with 173 additions and 43 deletions

View File

@@ -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);
}