openai_models/
lib.rs

1use std::{collections::VecDeque, str::FromStr};
2
3use derive_more::derive::Display;
4use serde::{Deserialize, Serialize};
5
6pub mod llm;
7
8// General models, note might alias to a specific model
9#[derive(Debug, Clone, Serialize, Deserialize, Display)]
10pub enum OpenAIModel {
11    #[display("gpt-4o")]
12    GPT4O,
13    #[display("gpt-4o-mini")]
14    GPT4OMINI,
15    #[display("o1")]
16    O1,
17    #[display("o1-mini")]
18    O1MINI,
19    #[display("gpt-3.5-turbo")]
20    GPT35TURBO,
21    #[display("gpt-4")]
22    GPT4,
23    #[display("gpt-4-turbo")]
24    GPT4TURBO,
25    #[display("{_0}")]
26    Other(String, PricingInfo),
27}
28
29impl FromStr for OpenAIModel {
30    type Err = String;
31    fn from_str(s: &str) -> Result<Self, Self::Err> {
32        match s {
33            "gpt-4o" | "gpt4o" => Ok(Self::GPT4O),
34            "gpt-4" | "gpt" => Ok(Self::GPT4),
35            "gpt-4-turbo" | "gpt4turbo" => Ok(Self::GPT4TURBO),
36            "gpt-4o-mini" | "gpt4omini" => Ok(Self::GPT4OMINI),
37            "o1" => Ok(Self::O1),
38            "o1-mini" => Ok(Self::O1MINI),
39            "gpt-3.5-turbo" | "gpt3.5turbo" => Ok(Self::GPT35TURBO),
40            _ => {
41                let mut tks = s
42                    .split(",")
43                    .map(|t| t.to_string())
44                    .collect::<VecDeque<String>>();
45
46                if tks.len() >= 2 {
47                    let model = tks.pop_front().unwrap();
48                    let tks = tks
49                        .into_iter()
50                        .map(|t| f64::from_str(&t))
51                        .collect::<Result<Vec<f64>, _>>()
52                        .map_err(|e| e.to_string())?;
53
54                    let pricing = if tks.len() == 2 {
55                        PricingInfo {
56                            input_tokens: tks[0],
57                            output_tokens: tks[1],
58                            cached_input_tokens: None,
59                        }
60                    } else if tks.len() == 3 {
61                        PricingInfo {
62                            input_tokens: tks[0],
63                            output_tokens: tks[1],
64                            cached_input_tokens: Some(tks[2]),
65                        }
66                    } else {
67                        return Err("fail to parse pricing".to_string());
68                    };
69
70                    Ok(Self::Other(model, pricing))
71                } else {
72                    Err("unreconigized model".to_string())
73                }
74            }
75        }
76    }
77}
78
79// USD per 1M tokens
80// From https://openai.com/api/pricing/
81#[derive(Copy, Debug, Clone, Serialize, Deserialize)]
82pub struct PricingInfo {
83    pub input_tokens: f64,
84    pub output_tokens: f64,
85    pub cached_input_tokens: Option<f64>,
86}
87
88impl FromStr for PricingInfo {
89    type Err = String;
90
91    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
92        let tks = s
93            .split(",")
94            .map(|t| f64::from_str(t))
95            .collect::<Result<Vec<f64>, _>>()
96            .map_err(|e| e.to_string())?;
97
98        if tks.len() == 2 {
99            Ok(PricingInfo {
100                input_tokens: tks[0],
101                output_tokens: tks[1],
102                cached_input_tokens: None,
103            })
104        } else if tks.len() == 3 {
105            Ok(PricingInfo {
106                input_tokens: tks[0],
107                output_tokens: tks[1],
108                cached_input_tokens: Some(tks[2]),
109            })
110        } else {
111            Err("fail to parse pricing".to_string())
112        }
113    }
114}
115
116impl OpenAIModel {
117    pub fn pricing(&self) -> PricingInfo {
118        match self {
119            Self::GPT4O => PricingInfo {
120                input_tokens: 2.5,
121                output_tokens: 10.00,
122                cached_input_tokens: Some(1.25),
123            },
124            Self::GPT4OMINI => PricingInfo {
125                input_tokens: 0.15,
126                cached_input_tokens: Some(0.075),
127                output_tokens: 0.6,
128            },
129            Self::O1 => PricingInfo {
130                input_tokens: 15.00,
131                cached_input_tokens: Some(7.5),
132                output_tokens: 60.00,
133            },
134            Self::O1MINI => PricingInfo {
135                input_tokens: 3.0,
136                cached_input_tokens: Some(1.5),
137                output_tokens: 12.00,
138            },
139            Self::GPT35TURBO => PricingInfo {
140                input_tokens: 3.0,
141                cached_input_tokens: None,
142                output_tokens: 6.0,
143            },
144            Self::GPT4 => PricingInfo {
145                input_tokens: 30.0,
146                output_tokens: 60.0,
147                cached_input_tokens: None,
148            },
149            Self::GPT4TURBO => PricingInfo {
150                input_tokens: 10.0,
151                output_tokens: 30.0,
152                cached_input_tokens: None,
153            },
154            Self::Other(_, pricing) => *pricing,
155        }
156    }
157
158    pub fn batch_pricing(&self) -> Option<PricingInfo> {
159        match self {
160            Self::GPT4O => Some(PricingInfo {
161                input_tokens: 1.25,
162                output_tokens: 5.00,
163                cached_input_tokens: None,
164            }),
165            Self::GPT4OMINI => Some(PricingInfo {
166                input_tokens: 0.075,
167                output_tokens: 0.3,
168                cached_input_tokens: None,
169            }),
170            _ => None,
171        }
172    }
173}