llm_models/api_model/
openai.rs

1use super::ApiLlmModel;
2use crate::{tokenizer::LlmTokenizer, LlmModelBase};
3use std::sync::Arc;
4
5impl ApiLlmModel {
6    pub fn openai_model_from_model_id(model_id: &str) -> Self {
7        match model_id {
8            model_id if model_id.starts_with("gpt-4") => Self::gpt_4(),
9            model_id if model_id.starts_with("gpt-3.5-turbo") => Self::gpt_3_5_turbo(),
10            model_id if model_id.starts_with("gpt-4-32k") => Self::gpt_4_32k(),
11            model_id if model_id.starts_with("gpt-4-turbo") => Self::gpt_4_turbo(),
12            model_id if model_id.starts_with("gpt-4o") => Self::gpt_4_o(),
13            model_id if model_id.starts_with("gpt-4o-mini") => Self::gpt_4_o_mini(),
14            model_id if model_id.starts_with("o1") => Self::o1(),
15            model_id if model_id.starts_with("o1-mini") => Self::o1_mini(),
16            _ => panic!("Model ID ({model_id}) not found for ApiLlmModel"),
17        }
18    }
19
20    pub fn gpt_4() -> Self {
21        let model_id = "gpt-4".to_string();
22        let tokenizer = model_tokenizer(&model_id);
23        Self {
24            model_base: LlmModelBase {
25                model_id,
26                model_ctx_size: 8192,
27                inference_ctx_size: 4096,
28                tokenizer,
29            },
30            cost_per_m_in_tokens: 30.00,
31            cost_per_m_out_tokens: 60.00,
32            tokens_per_message: 3,
33            tokens_per_name: Some(1),
34        }
35    }
36
37    pub fn gpt_3_5_turbo() -> Self {
38        let model_id = "gpt-3.5-turbo".to_string();
39        let tokenizer = model_tokenizer(&model_id);
40        Self {
41            model_base: LlmModelBase {
42                model_id,
43                model_ctx_size: 16385,
44                inference_ctx_size: 4096,
45                tokenizer,
46            },
47            cost_per_m_in_tokens: 0.50,
48            cost_per_m_out_tokens: 1.50,
49            tokens_per_message: 4,
50            tokens_per_name: Some(-1),
51        }
52    }
53
54    pub fn gpt_4_32k() -> Self {
55        let model_id = "gpt-4-32k".to_string();
56        let tokenizer = model_tokenizer(&model_id);
57        Self {
58            model_base: LlmModelBase {
59                model_id,
60                model_ctx_size: 32768,
61                inference_ctx_size: 4096,
62                tokenizer,
63            },
64            cost_per_m_in_tokens: 60.00,
65            cost_per_m_out_tokens: 120.00,
66            tokens_per_message: 3,
67            tokens_per_name: Some(1),
68        }
69    }
70
71    pub fn gpt_4_turbo() -> Self {
72        let model_id = "gpt-4-turbo".to_string();
73        let tokenizer = model_tokenizer(&model_id);
74        Self {
75            model_base: LlmModelBase {
76                model_id,
77                model_ctx_size: 128000,
78                inference_ctx_size: 4096,
79                tokenizer,
80            },
81            cost_per_m_in_tokens: 10.00,
82            cost_per_m_out_tokens: 30.00,
83            tokens_per_message: 3,
84            tokens_per_name: Some(1),
85        }
86    }
87
88    pub fn gpt_4_o_mini() -> Self {
89        let model_id = "gpt-4o-mini".to_string();
90        let tokenizer = model_tokenizer(&model_id);
91        Self {
92            model_base: LlmModelBase {
93                model_id,
94                model_ctx_size: 128000,
95                inference_ctx_size: 16384,
96                tokenizer,
97            },
98            cost_per_m_in_tokens: 0.150,
99            cost_per_m_out_tokens: 0.60,
100            tokens_per_message: 3,
101            tokens_per_name: Some(1),
102        }
103    }
104
105    pub fn gpt_4_o() -> Self {
106        let model_id = "gpt-4o".to_string();
107        let tokenizer = model_tokenizer(&model_id);
108        Self {
109            model_base: LlmModelBase {
110                model_id,
111                model_ctx_size: 128000,
112                inference_ctx_size: 4096,
113                tokenizer,
114            },
115            cost_per_m_in_tokens: 5.00,
116            cost_per_m_out_tokens: 15.00,
117            tokens_per_message: 3,
118            tokens_per_name: Some(1),
119        }
120    }
121
122    pub fn o1() -> Self {
123        let model_id = "o1".to_string();
124        let tokenizer = model_tokenizer(&model_id);
125        Self {
126            model_base: LlmModelBase {
127                model_id,
128                model_ctx_size: 200000,
129                inference_ctx_size: 100000,
130                tokenizer,
131            },
132            cost_per_m_in_tokens: 15.00,
133            cost_per_m_out_tokens: 60.00,
134            tokens_per_message: 4,
135            tokens_per_name: Some(-1),
136        }
137    }
138
139    pub fn o1_mini() -> Self {
140        let model_id = "o1-mini".to_string();
141        let tokenizer = model_tokenizer(&model_id);
142        Self {
143            model_base: LlmModelBase {
144                model_id,
145                model_ctx_size: 128000,
146                inference_ctx_size: 65536,
147                tokenizer,
148            },
149            cost_per_m_in_tokens: 3.00,
150            cost_per_m_out_tokens: 12.00,
151            tokens_per_message: 4,
152            tokens_per_name: Some(-1),
153        }
154    }
155}
156
157fn model_tokenizer(model_id: &str) -> Arc<LlmTokenizer> {
158    Arc::new(
159        LlmTokenizer::new_tiktoken(model_id)
160            .unwrap_or_else(|_| panic!("Failed to load tokenizer for {model_id}")),
161    )
162}
163
164pub trait OpenAiModelTrait {
165    fn model(&mut self) -> &mut ApiLlmModel;
166
167    /// Set the model using the model_id string.
168    fn model_id_str(mut self, model_id: &str) -> Self
169    where
170        Self: Sized,
171    {
172        *self.model() = ApiLlmModel::openai_model_from_model_id(model_id);
173        self
174    }
175
176    /// Use gpt-4 as the model for the OpenAI client.
177    fn gpt_4(mut self) -> Self
178    where
179        Self: Sized,
180    {
181        *self.model() = ApiLlmModel::gpt_4();
182        self
183    }
184
185    /// Use gpt-4-32k as the model for the OpenAI client. Limited support for this model from OpenAI.
186    fn gpt_4_32k(mut self) -> Self
187    where
188        Self: Sized,
189    {
190        *self.model() = ApiLlmModel::gpt_4_32k();
191        self
192    }
193
194    /// Use gpt-3.5-turbo as the model for the OpenAI client.
195    fn gpt_3_5_turbo(mut self) -> Self
196    where
197        Self: Sized,
198    {
199        *self.model() = ApiLlmModel::gpt_3_5_turbo();
200        self
201    }
202
203    /// Use gpt-4-turbo as the model for the OpenAI client.
204    fn gpt_4_turbo(mut self) -> Self
205    where
206        Self: Sized,
207    {
208        *self.model() = ApiLlmModel::gpt_4_turbo();
209        self
210    }
211
212    /// Use gpt-4-o as the model for the OpenAI client.
213    fn gpt_4_o(mut self) -> Self
214    where
215        Self: Sized,
216    {
217        *self.model() = ApiLlmModel::gpt_4_o();
218        self
219    }
220
221    /// Use gpt-4o-mini as the model for the OpenAI client.
222    fn gpt_4_o_mini(mut self) -> Self
223    where
224        Self: Sized,
225    {
226        *self.model() = ApiLlmModel::gpt_4_o_mini();
227        self
228    }
229
230    /// Use o1 as the model for the OpenAI client.
231    fn o1(mut self) -> Self
232    where
233        Self: Sized,
234    {
235        *self.model() = ApiLlmModel::o1();
236        self
237    }
238
239    /// Use o1-mini as the model for the OpenAI client.
240    fn o1_mini(mut self) -> Self
241    where
242        Self: Sized,
243    {
244        *self.model() = ApiLlmModel::o1_mini();
245        self
246    }
247}