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 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 fn gpt_4(mut self) -> Self
178 where
179 Self: Sized,
180 {
181 *self.model() = ApiLlmModel::gpt_4();
182 self
183 }
184
185 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 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 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 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 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 fn o1(mut self) -> Self
232 where
233 Self: Sized,
234 {
235 *self.model() = ApiLlmModel::o1();
236 self
237 }
238
239 fn o1_mini(mut self) -> Self
241 where
242 Self: Sized,
243 {
244 *self.model() = ApiLlmModel::o1_mini();
245 self
246 }
247}