Skip to content

Commit d486a94

Browse files
committed
gtp_4_32k support added
1 parent 1009ee9 commit d486a94

File tree

1 file changed

+36
-0
lines changed

1 file changed

+36
-0
lines changed

src/models.rs

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ pub enum Model {
1818
Gpt3_5Turbo,
1919
#[serde(rename = "gpt-4")]
2020
Gpt_4,
21+
#[serde(rename = "gpt-4-32k")]
22+
Gpt_4_32k,
2123
}
2224

2325
/// Implement Display to convert the enum back to a string representation.
@@ -26,6 +28,7 @@ impl Display for Model {
2628
let model_name = match self {
2729
Model::Gpt3_5Turbo => "gpt-3.5-turbo",
2830
Model::Gpt_4 => "gpt-4",
31+
Model::Gpt_4_32k => "gpt-4-32k",
2932
};
3033
write!(f, "{model_name}")
3134
}
@@ -39,6 +42,7 @@ impl FromStr for Model {
3942
match s {
4043
"gpt-3.5-turbo" => Ok(Model::Gpt3_5Turbo),
4144
"gpt-4" => Ok(Model::Gpt_4),
45+
"gpt-4-32k" => Ok(Model::Gpt_4_32k),
4246
_ => Err(()),
4347
}
4448
}
@@ -142,6 +146,38 @@ mod tests {
142146
assert_eq!(deserialized_model, Model::Gpt3_5Turbo);
143147
}
144148

149+
// Test the conversion of a valid model string to a Model enum variant for Gpt_4_32k
150+
#[test]
151+
fn test_from_str_gpt4_32k() {
152+
let input = "gpt-4-32k";
153+
let model: Result<Model, ()> = Model::from_str(input);
154+
assert!(model.is_ok(), "Failed to parse the gpt-4-32k model name");
155+
assert_eq!(model.unwrap(), Model::Gpt_4_32k);
156+
}
157+
// Test the conversion of a Model enum variant to its string representation for Gpt_4_32k:
158+
#[test]
159+
fn test_display_gpt4_32k() {
160+
let model = Model::Gpt_4_32k;
161+
let model_str = format!("{}", model);
162+
assert_eq!(model_str, "gpt-4-32k");
163+
}
164+
165+
// Test the serialization of a Model enum variant to JSON for Gpt_4_32k:
166+
#[test]
167+
fn test_serialize_gpt4_32k() {
168+
let model = Model::Gpt_4_32k;
169+
let serialized_model = serde_json::to_string(&model).unwrap();
170+
assert_eq!(serialized_model, "\"gpt-4-32k\"");
171+
}
172+
173+
// Test the deserialization of a JSON string to a Model enum variant for Gpt_4_32k
174+
#[test]
175+
fn test_deserialize_gpt4_32k() {
176+
let model_json = "\"gpt-4-32k\"";
177+
let deserialized_model: Model = serde_json::from_str(model_json).unwrap();
178+
assert_eq!(deserialized_model, Model::Gpt_4_32k);
179+
}
180+
145181
// Test the deserialization of a JSON string to a `Model` enum variant for Gpt3_5Turbo.
146182
#[test]
147183
fn test_deserialize_gpt4() {

0 commit comments

Comments
 (0)