drfoo/
model.rs

1use derive_more::Display;
2use serde::{Deserialize, Serialize};
3
4/// `OpenAIChatModel` is an enum of the OpenAI models available for chat.
5#[allow(non_camel_case_types)]
6#[derive(Debug, Clone, Deserialize, Serialize, Display)]
7pub enum OpenAIChatModel {
8    #[serde(rename = "gpt-4")]
9    Gpt_4,
10    #[serde(rename = "gpt-4-32k")]
11    Gpt_4_32k,
12    #[serde(rename = "gpt-3.5-turbo")]
13    Gpt_3dot5_turbo,
14}
15
16/// `ChatModel` is an enum of the models available for chat.
17#[derive(Debug, Clone, Deserialize, Serialize, Display)]
18#[serde(untagged)]
19pub enum ChatModel {
20    OpenAI(OpenAIChatModel),
21
22    /// If a model is not available in the enum, you can use a custom model.
23    Custom(String),
24}
25
26impl Default for ChatModel {
27    fn default() -> Self {
28        Self::OpenAI(OpenAIChatModel::Gpt_4)
29    }
30}
31
32impl<T: Into<String>> From<T> for ChatModel {
33    fn from(model: T) -> Self {
34        Self::Custom(model.into())
35    }
36}
37
38impl From<OpenAIChatModel> for ChatModel {
39    fn from(model: OpenAIChatModel) -> Self {
40        Self::OpenAI(model)
41    }
42}
43
44/// `OpenAICompletionModel` is an enum of the OpenAI models available for completion.
45#[allow(non_camel_case_types)]
46#[derive(Debug, Clone, Deserialize, Serialize, Display)]
47pub enum OpenAICompletionModel {
48    #[serde(rename = "text-davinci-003")]
49    Text_Davinci_003,
50    #[serde(rename = "text-davinci-002")]
51    Text_Davinci_002,
52    #[serde(rename = "text-davinci-001")]
53    Text_Davinci_001,
54    #[serde(rename = "text-curie-001")]
55    Text_Curie_001,
56    #[serde(rename = "text-babbage-001")]
57    Text_Babbage_001,
58    #[serde(rename = "text-ada-001")]
59    Text_Ada_001,
60}
61
62#[derive(Debug, Clone, Deserialize, Serialize, Display)]
63#[serde(untagged)]
64pub enum CompletionModel {
65    OpenAI(OpenAICompletionModel),
66
67    /// If a model is not available in the enum, you can use a custom model.
68    Custom(String),
69}
70
71impl Default for CompletionModel {
72    fn default() -> Self {
73        Self::OpenAI(OpenAICompletionModel::Text_Davinci_003)
74    }
75}
76
77impl<T: Into<String>> From<T> for CompletionModel {
78    fn from(model: T) -> Self {
79        Self::Custom(model.into())
80    }
81}
82
83impl From<OpenAICompletionModel> for CompletionModel {
84    fn from(model: OpenAICompletionModel) -> Self {
85        Self::OpenAI(model)
86    }
87}