1use std::{collections::VecDeque, str::FromStr};
2
3use derive_more::derive::Display;
4use serde::{Deserialize, Serialize};
5
6pub mod llm;
7
8#[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#[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}