1use std::borrow::Cow;
4
5use aither_core::llm::model::Ability;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9#[non_exhaustive]
10pub enum ModelMode {
11 Chat,
13 Embedding,
15 ImageGeneration,
17 AudioTranscription,
19 AudioSpeech,
21 Rerank,
23 Moderation,
25 Completion,
27 VideoGeneration,
29 Search,
31 Ocr,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37#[non_exhaustive]
38pub enum Provider {
39 OpenAI,
41 Anthropic,
43 Google,
45 DeepSeek,
47 XAI,
49 Mistral,
51 Meta,
53 Alibaba,
55 Bedrock,
57 Azure,
59 VertexAI,
61 Groq,
63 Together,
65 Fireworks,
67 Replicate,
69 Cohere,
71 Perplexity,
73 Copilot,
75 Other(Cow<'static, str>),
77}
78
79impl Provider {
80 #[must_use]
82 pub fn as_str(&self) -> &str {
83 match self {
84 Self::OpenAI => "openai",
85 Self::Anthropic => "anthropic",
86 Self::Google => "gemini",
87 Self::DeepSeek => "deepseek",
88 Self::XAI => "xai",
89 Self::Mistral => "mistral",
90 Self::Meta => "meta",
91 Self::Alibaba => "alibaba",
92 Self::Bedrock => "bedrock",
93 Self::Azure => "azure",
94 Self::VertexAI => "vertex_ai",
95 Self::Groq => "groq",
96 Self::Together => "together_ai",
97 Self::Fireworks => "fireworks_ai",
98 Self::Replicate => "replicate",
99 Self::Cohere => "cohere",
100 Self::Perplexity => "perplexity",
101 Self::Copilot => "github_copilot",
102 Self::Other(s) => s,
103 }
104 }
105}
106
107impl core::fmt::Display for Provider {
108 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
109 f.write_str(self.as_str())
110 }
111}
112
113#[allow(clippy::struct_field_names)]
117#[derive(Debug, Clone, PartialEq, Default)]
118pub struct Pricing {
119 input_per_token: f64,
120 output_per_token: f64,
121 cache_read_per_token: Option<f64>,
122 cache_write_per_token: Option<f64>,
123 reasoning_per_token: Option<f64>,
124 image_per_token: Option<f64>,
125}
126
127impl Pricing {
128 #[must_use]
130 pub const fn new(input_per_token: f64, output_per_token: f64) -> Self {
131 Self {
132 input_per_token,
133 output_per_token,
134 cache_read_per_token: None,
135 cache_write_per_token: None,
136 reasoning_per_token: None,
137 image_per_token: None,
138 }
139 }
140
141 #[must_use]
143 pub const fn input_per_token(&self) -> f64 {
144 self.input_per_token
145 }
146
147 #[must_use]
149 pub const fn output_per_token(&self) -> f64 {
150 self.output_per_token
151 }
152
153 #[must_use]
155 pub const fn cache_read_per_token(&self) -> Option<f64> {
156 self.cache_read_per_token
157 }
158
159 #[must_use]
161 pub const fn cache_write_per_token(&self) -> Option<f64> {
162 self.cache_write_per_token
163 }
164
165 #[must_use]
167 pub const fn reasoning_per_token(&self) -> Option<f64> {
168 self.reasoning_per_token
169 }
170
171 #[must_use]
173 pub const fn image_per_token(&self) -> Option<f64> {
174 self.image_per_token
175 }
176
177 #[must_use]
179 pub const fn with_cache_read(mut self, cost: f64) -> Self {
180 self.cache_read_per_token = Some(cost);
181 self
182 }
183
184 #[must_use]
186 pub const fn with_cache_write(mut self, cost: f64) -> Self {
187 self.cache_write_per_token = Some(cost);
188 self
189 }
190
191 #[must_use]
193 pub const fn with_reasoning(mut self, cost: f64) -> Self {
194 self.reasoning_per_token = Some(cost);
195 self
196 }
197
198 #[must_use]
200 pub const fn with_image(mut self, cost: f64) -> Self {
201 self.image_per_token = Some(cost);
202 self
203 }
204}
205
206#[derive(Debug, Clone, PartialEq)]
208pub struct ModelEntry {
209 litellm_id: Cow<'static, str>,
210 canonical_id: Cow<'static, str>,
211 provider: Provider,
212 mode: ModelMode,
213 max_input_tokens: Option<u32>,
214 max_output_tokens: Option<u32>,
215 pricing: Pricing,
216 abilities: Cow<'static, [Ability]>,
217 deprecation_date: Option<Cow<'static, str>>,
218}
219
220impl ModelEntry {
221 #[must_use]
223 pub fn new(
224 litellm_id: String,
225 canonical_id: String,
226 provider: Provider,
227 mode: ModelMode,
228 ) -> Self {
229 Self {
230 litellm_id: Cow::Owned(litellm_id),
231 canonical_id: Cow::Owned(canonical_id),
232 provider,
233 mode,
234 max_input_tokens: None,
235 max_output_tokens: None,
236 pricing: Pricing::default(),
237 abilities: Cow::Borrowed(&[]),
238 deprecation_date: None,
239 }
240 }
241
242 #[must_use]
244 #[doc(hidden)]
245 #[allow(clippy::too_many_arguments)]
246 pub const fn new_static(
247 litellm_id: &'static str,
248 canonical_id: &'static str,
249 provider: Provider,
250 mode: ModelMode,
251 max_input_tokens: Option<u32>,
252 max_output_tokens: Option<u32>,
253 pricing: Pricing,
254 abilities: &'static [Ability],
255 deprecation_date: Option<&'static str>,
256 ) -> Self {
257 Self {
258 litellm_id: Cow::Borrowed(litellm_id),
259 canonical_id: Cow::Borrowed(canonical_id),
260 provider,
261 mode,
262 max_input_tokens,
263 max_output_tokens,
264 pricing,
265 abilities: Cow::Borrowed(abilities),
266 deprecation_date: match deprecation_date {
267 Some(s) => Some(Cow::Borrowed(s)),
268 None => None,
269 },
270 }
271 }
272
273 #[must_use]
275 pub fn id(&self) -> &str {
276 &self.canonical_id
277 }
278
279 #[must_use]
281 pub fn litellm_id(&self) -> &str {
282 &self.litellm_id
283 }
284
285 #[must_use]
287 pub const fn provider(&self) -> &Provider {
288 &self.provider
289 }
290
291 #[must_use]
293 pub const fn mode(&self) -> ModelMode {
294 self.mode
295 }
296
297 #[must_use]
299 pub const fn max_input_tokens(&self) -> Option<u32> {
300 self.max_input_tokens
301 }
302
303 #[must_use]
305 pub const fn max_output_tokens(&self) -> Option<u32> {
306 self.max_output_tokens
307 }
308
309 #[must_use]
311 pub const fn pricing(&self) -> &Pricing {
312 &self.pricing
313 }
314
315 #[must_use]
317 pub fn abilities(&self) -> &[Ability] {
318 &self.abilities
319 }
320
321 #[must_use]
323 pub fn has_ability(&self, ability: Ability) -> bool {
324 self.abilities.contains(&ability)
325 }
326
327 #[must_use]
329 pub fn deprecation_date(&self) -> Option<&str> {
330 self.deprecation_date.as_deref()
331 }
332
333 #[must_use]
335 pub fn is_deprecated(&self) -> bool {
336 self.deprecation_date.as_deref().is_some_and(|date| {
337 let now = std::time::SystemTime::now()
338 .duration_since(std::time::UNIX_EPOCH)
339 .unwrap_or_default()
340 .as_secs();
341 let days = now / 86400;
342 let (year, month, day) = days_to_ymd(days);
343 let today = format!("{year:04}-{month:02}-{day:02}");
344 date <= today.as_str()
345 })
346 }
347
348 #[must_use]
350 pub const fn with_max_input_tokens(mut self, tokens: u32) -> Self {
351 self.max_input_tokens = Some(tokens);
352 self
353 }
354
355 #[must_use]
357 pub const fn with_max_output_tokens(mut self, tokens: u32) -> Self {
358 self.max_output_tokens = Some(tokens);
359 self
360 }
361
362 #[must_use]
364 pub const fn with_pricing(mut self, pricing: Pricing) -> Self {
365 self.pricing = pricing;
366 self
367 }
368
369 #[must_use]
371 pub fn with_abilities(mut self, abilities: Vec<Ability>) -> Self {
372 self.abilities = Cow::Owned(abilities);
373 self
374 }
375
376 #[must_use]
378 pub fn with_deprecation_date(mut self, date: String) -> Self {
379 self.deprecation_date = Some(Cow::Owned(date));
380 self
381 }
382}
383
384const fn days_to_ymd(days: u64) -> (u64, u64, u64) {
386 let z = days + 719_468;
388 let era = z / 146_097;
389 let doe = z - era * 146_097;
390 let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146_096) / 365;
391 let y = yoe + era * 400;
392 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
393 let mp = (5 * doy + 2) / 153;
394 let d = doy - (153 * mp + 2) / 5 + 1;
395 let m = if mp < 10 { mp + 3 } else { mp - 9 };
396 let y = if m <= 2 { y + 1 } else { y };
397 (y, m, d)
398}