llm_rs/
model_mode.rs

1use std::str::FromStr;
2/// The modes for the interaction with Open AI
3#[derive(Debug, Clone, PartialEq)]
4pub enum ModelMode {
5    Completions,
6    Chat,
7    Image,
8    ImageEdit,
9    AudioTranscription,
10}
11const MODELS_COMPLETIONS: [&str; 6] = [
12    "text-babbage-001",
13    "text-curie-001",
14    "text-davinci-001",
15    "text-davinci-002",
16    "text-davinci-003",
17    "text-davinci-edit-001",
18];
19const MODELS_CHAT: [&str; 2] = ["gpt-3.5-turbo", "gpt-4"];
20
21const MODELS_AUDIOTRANSCRIPTION: [&str; 1] = ["whisper-1"];
22
23impl std::fmt::Display for ModelMode {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        let str = match self {
26            ModelMode::Completions => "completions",
27            ModelMode::Chat => "chat",
28            ModelMode::Image => "image",
29            ModelMode::ImageEdit => "image_edit",
30            ModelMode::AudioTranscription => "audio_transcription",
31        };
32        write!(f, "{str}")
33    }
34}
35
36#[derive(Debug)]
37pub struct ModelModeParseErr;
38
39impl FromStr for ModelMode {
40    type Err = ModelModeParseErr;
41    fn from_str(mode: &str) -> Result<Self, Self::Err> {
42        match mode {
43            "completions" => Ok(ModelMode::Completions),
44            "chat" => Ok(ModelMode::Chat),
45            "image" => Ok(ModelMode::Image),
46            "image_edit" => Ok(ModelMode::ImageEdit),
47            "audio_transcription" => Ok(ModelMode::AudioTranscription),
48            _ => Err(ModelModeParseErr),
49        }
50    }
51}
52
53impl ModelMode {
54    pub fn models_available(&self) -> Vec<&str> {
55        match self {
56            ModelMode::Completions => MODELS_COMPLETIONS.to_vec(),
57            ModelMode::Chat => MODELS_CHAT.to_vec(),
58            ModelMode::Image => [].to_vec(),
59            ModelMode::ImageEdit => [].to_vec(),
60            ModelMode::AudioTranscription => MODELS_AUDIOTRANSCRIPTION.to_vec(),
61        }
62    }
63}
llm_rs/model_mode.rs

llm_rs/
model_mode.rs
