From ecd0c3028fdcfab32b867eb1eed86182f5f4ab81 Mon Sep 17 00:00:00 2001
From: Ivan Gabriele <ivan.gabriele@protonmail.com>
Date: Sat, 22 Jun 2024 13:22:57 +0200
Subject: [PATCH] feat(constants): add OpenMixtral8x22b, MistralTiny &
 CodestralLatest to Model enum

---
 src/v1/constants.rs        |  6 +++++
 tests/v1_constants_test.rs | 45 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)
 create mode 100644 tests/v1_constants_test.rs

diff --git a/src/v1/constants.rs b/src/v1/constants.rs
index 06bbceb..e060664 100644
--- a/src/v1/constants.rs
+++ b/src/v1/constants.rs
@@ -8,12 +8,18 @@ pub enum Model {
     OpenMistral7b,
     #[serde(rename = "open-mixtral-8x7b")]
     OpenMixtral8x7b,
+    #[serde(rename = "open-mixtral-8x22b")]
+    OpenMixtral8x22b,
+    #[serde(rename = "mistral-tiny")]
+    MistralTiny,
     #[serde(rename = "mistral-small-latest")]
     MistralSmallLatest,
     #[serde(rename = "mistral-medium-latest")]
     MistralMediumLatest,
     #[serde(rename = "mistral-large-latest")]
     MistralLargeLatest,
+    #[serde(rename = "codestral-latest")]
+    CodestralLatest,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
diff --git a/tests/v1_constants_test.rs b/tests/v1_constants_test.rs
new file mode 100644
index 0000000..f1ced52
--- /dev/null
+++ b/tests/v1_constants_test.rs
@@ -0,0 +1,45 @@
+use jrest::expect;
+use mistralai_client::v1::{
+    chat::{ChatMessage, ChatParams},
+    client::Client,
+    constants::Model,
+};
+
+mod setup;
+
+#[test]
+fn test_model_constant() {
+    setup::setup();
+
+    let models = vec![
+        Model::OpenMistral7b,
+        Model::OpenMixtral8x7b,
+        Model::OpenMixtral8x22b,
+        Model::MistralTiny,
+        Model::MistralSmallLatest,
+        Model::MistralMediumLatest,
+        Model::MistralLargeLatest,
+        Model::CodestralLatest,
+    ];
+
+    let client = Client::new(None, None, None, None).unwrap();
+
+    let messages = vec![ChatMessage::new_user_message("A number between 0 and 100?")];
+    let options = ChatParams {
+        temperature: 0.0,
+        random_seed: Some(42),
+        ..Default::default()
+    };
+
+    for model in models {
+        let response = client
+            .chat(model.clone(), messages.clone(), Some(options.clone()))
+            .unwrap();
+
+        expect!(response.model).to_be(model);
+        expect!(response.object).to_be("chat.completion".to_string());
+        expect!(response.choices.len()).to_be(1);
+        expect!(response.choices[0].index).to_be(0);
+        expect!(response.choices[0].message.content.len()).to_be_greater_than(0);
+    }
+}