1use std::{collections::VecDeque, str::FromStr};
2
3use derive_more::derive::Display;
4use serde::{Deserialize, Serialize};
5
6pub mod error;
7pub mod llm;
8
9pub mod openai {
10 pub use async_openai::*;
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize, Display)]
15pub enum OpenAIModel {
16 #[display("gpt-4o")]
17 GPT4O,
18 #[display("gpt-4o-mini")]
19 GPT4OMINI,
20 #[display("o1")]
21 O1,
22 #[display("o1-mini")]
23 O1MINI,
24 #[display("gpt-3.5-turbo")]
25 GPT35TURBO,
26 #[display("gpt-4")]
27 GPT4,
28 #[display("gpt-4-turbo")]
29 GPT4TURBO,
30 #[display("gpt-5-mini")]
31 GPT5MINI,
32 #[display("gpt-5-nano")]
33 GPT5NANO,
34 #[display("gpt-5.2")]
35 GPT52,
36 #[display("gpt-5")]
37 GPT5,
38 #[display("gemini-3-pro-preview")]
39 GEMINI3PRO,
40 #[display("gemini-3-flash-preview")]
41 GEMINI3FLASH,
42 #[display("gemini-2.5-pro")]
43 GEMINI25PRO,
44 #[display("gemini-2.5-flash")]
45 GEMINI25FLASH,
46 #[display("{_0}")]
47 Other(String, PricingInfo),
48}
49
50impl FromStr for OpenAIModel {
51 type Err = String;
52 fn from_str(s: &str) -> Result<Self, Self::Err> {
53 match s {
54 "gpt-4o" | "gpt4o" => Ok(Self::GPT4O),
55 "gpt-4" | "gpt" => Ok(Self::GPT4),
56 "gpt-4-turbo" | "gpt4turbo" => Ok(Self::GPT4TURBO),
57 "gpt-4o-mini" | "gpt4omini" => Ok(Self::GPT4OMINI),
58 "o1" => Ok(Self::O1),
59 "o1-mini" => Ok(Self::O1MINI),
60 "gpt-3.5-turbo" | "gpt3.5turbo" => Ok(Self::GPT35TURBO),
61 "gpt-5.2" => Ok(Self::GPT52),
62 "gpt-5-mini" | "gpt-5mini" => Ok(Self::GPT5MINI),
63 "gpt-5" => Ok(Self::GPT5),
64 "gpt-5-nano" | "gpt-5nano" => Ok(Self::GPT5NANO),
65 "gemini-3-pro-preview" | "gemini-3-pro" => Ok(Self::GEMINI3PRO),
66 "gemini-3-flash-preview" | "gemini-3-flash" => Ok(Self::GEMINI3FLASH),
67 "gemini-2.5-pro" => Ok(Self::GEMINI25PRO),
68 "gemini-2.5-flash" => Ok(Self::GEMINI25FLASH),
69 _ => {
70 if !s.contains(",") {
71 return Ok(Self::Other(
72 s.to_string(),
73 PricingInfo {
74 input_tokens: 0.0f64,
75 output_tokens: 0.0f64,
76 cached_input_tokens: None,
77 },
78 ));
79 }
80 let mut tks = s
81 .split(",")
82 .map(|t| t.to_string())
83 .collect::<VecDeque<String>>();
84
85 if tks.len() >= 2 {
86 let model = tks.pop_front().unwrap();
87 let tks = tks
88 .into_iter()
89 .map(|t| f64::from_str(&t))
90 .collect::<Result<Vec<f64>, _>>()
91 .map_err(|e| e.to_string())?;
92
93 let pricing = if tks.len() == 2 {
94 PricingInfo {
95 input_tokens: tks[0],
96 output_tokens: tks[1],
97 cached_input_tokens: None,
98 }
99 } else if tks.len() == 3 {
100 PricingInfo {
101 input_tokens: tks[0],
102 output_tokens: tks[1],
103 cached_input_tokens: Some(tks[2]),
104 }
105 } else {
106 return Err("fail to parse pricing".to_string());
107 };
108
109 Ok(Self::Other(model, pricing))
110 } else {
111 Err("unreconigized model".to_string())
112 }
113 }
114 }
115 }
116}
117
118#[derive(Copy, Debug, Clone, Serialize, Deserialize)]
121pub struct PricingInfo {
122 pub input_tokens: f64,
123 pub output_tokens: f64,
124 pub cached_input_tokens: Option<f64>,
125}
126
127impl FromStr for PricingInfo {
128 type Err = String;
129
130 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
131 let tks = s
132 .split(",")
133 .map(f64::from_str)
134 .collect::<Result<Vec<f64>, _>>()
135 .map_err(|e| e.to_string())?;
136
137 if tks.len() == 2 {
138 Ok(PricingInfo {
139 input_tokens: tks[0],
140 output_tokens: tks[1],
141 cached_input_tokens: None,
142 })
143 } else if tks.len() == 3 {
144 Ok(PricingInfo {
145 input_tokens: tks[0],
146 output_tokens: tks[1],
147 cached_input_tokens: Some(tks[2]),
148 })
149 } else {
150 Err("fail to parse pricing".to_string())
151 }
152 }
153}
154
155impl OpenAIModel {
156 pub fn pricing(&self) -> PricingInfo {
157 match self {
158 Self::GPT4O => PricingInfo {
159 input_tokens: 2.5,
160 output_tokens: 10.00,
161 cached_input_tokens: Some(1.25),
162 },
163 Self::GPT4OMINI => PricingInfo {
164 input_tokens: 0.15,
165 cached_input_tokens: Some(0.075),
166 output_tokens: 0.6,
167 },
168 Self::O1 => PricingInfo {
169 input_tokens: 15.00,
170 cached_input_tokens: Some(7.5),
171 output_tokens: 60.00,
172 },
173 Self::O1MINI => PricingInfo {
174 input_tokens: 3.0,
175 cached_input_tokens: Some(1.5),
176 output_tokens: 12.00,
177 },
178 Self::GPT35TURBO => PricingInfo {
179 input_tokens: 3.0,
180 cached_input_tokens: None,
181 output_tokens: 6.0,
182 },
183 Self::GPT4 => PricingInfo {
184 input_tokens: 30.0,
185 output_tokens: 60.0,
186 cached_input_tokens: None,
187 },
188 Self::GPT4TURBO => PricingInfo {
189 input_tokens: 10.0,
190 output_tokens: 30.0,
191 cached_input_tokens: None,
192 },
193 Self::GPT52 => PricingInfo {
194 input_tokens: 1.75,
195 output_tokens: 14.00,
196 cached_input_tokens: Some(0.175),
197 },
198 Self::GPT5MINI => PricingInfo {
199 input_tokens: 0.25,
200 output_tokens: 2.00,
201 cached_input_tokens: Some(0.025),
202 },
203 Self::GPT5NANO => PricingInfo {
204 input_tokens: 0.05,
205 output_tokens: 0.40,
206 cached_input_tokens: Some(0.005),
207 },
208 Self::GPT5 => PricingInfo {
209 input_tokens: 1.25,
210 output_tokens: 0.125,
211 cached_input_tokens: Some(10.00),
212 },
213 Self::GEMINI3PRO => PricingInfo {
214 input_tokens: 2.00, output_tokens: 12.00, cached_input_tokens: None,
217 },
218 Self::GEMINI3FLASH => PricingInfo {
219 input_tokens: 0.50,
220 output_tokens: 3.0,
221 cached_input_tokens: None,
222 },
223 Self::GEMINI25PRO => PricingInfo {
224 input_tokens: 1.25, output_tokens: 10.00, cached_input_tokens: None,
227 },
228 Self::GEMINI25FLASH => PricingInfo {
229 input_tokens: 0.30,
230 output_tokens: 2.50,
231 cached_input_tokens: None,
232 },
233 Self::Other(_, pricing) => *pricing,
234 }
235 }
236
237 pub fn batch_pricing(&self) -> Option<PricingInfo> {
238 match self {
239 Self::GPT4O => Some(PricingInfo {
240 input_tokens: 1.25,
241 output_tokens: 5.00,
242 cached_input_tokens: None,
243 }),
244 Self::GPT4OMINI => Some(PricingInfo {
245 input_tokens: 0.075,
246 output_tokens: 0.3,
247 cached_input_tokens: None,
248 }),
249 _ => None,
250 }
251 }
252}