Skip to main content

openai_models/
lib.rs

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// General models, note might alias to a specific model
14#[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// USD per 1M tokens
119// From https://openai.com/api/pricing/
120#[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,   // TODO: 4.00 for > 200k tokens
215                output_tokens: 12.00, // TODO: 18.00 for > 200k tokens
216                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,   // 2.50
225                output_tokens: 10.00, // 15.00
226                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}