Skip to main content

aither_models/
types.rs

1//! Owned model metadata types sourced from `LiteLLM` and `OpenRouter`.
2
3use std::borrow::Cow;
4
5use aither_core::llm::model::Ability;
6
7/// The operational mode of a model.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9#[non_exhaustive]
10pub enum ModelMode {
11    /// Chat / conversational completion.
12    Chat,
13    /// Text embedding / vectorization.
14    Embedding,
15    /// Image generation.
16    ImageGeneration,
17    /// Audio transcription (speech-to-text).
18    AudioTranscription,
19    /// Audio speech (text-to-speech).
20    AudioSpeech,
21    /// Reranking.
22    Rerank,
23    /// Content moderation.
24    Moderation,
25    /// Legacy text completion.
26    Completion,
27    /// Video generation.
28    VideoGeneration,
29    /// Search.
30    Search,
31    /// Optical character recognition.
32    Ocr,
33}
34
35/// Known model providers.
36#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37#[non_exhaustive]
38pub enum Provider {
39    /// `OpenAI`.
40    OpenAI,
41    /// Anthropic.
42    Anthropic,
43    /// Google (Gemini).
44    Google,
45    /// `DeepSeek`.
46    DeepSeek,
47    /// xAI (Grok).
48    XAI,
49    /// Mistral AI.
50    Mistral,
51    /// Meta (`LLaMA`).
52    Meta,
53    /// Alibaba (Qwen / `DashScope`).
54    Alibaba,
55    /// AWS Bedrock.
56    Bedrock,
57    /// Azure `OpenAI`.
58    Azure,
59    /// Google Vertex AI.
60    VertexAI,
61    /// Groq.
62    Groq,
63    /// Together AI.
64    Together,
65    /// Fireworks AI.
66    Fireworks,
67    /// Replicate.
68    Replicate,
69    /// Cohere.
70    Cohere,
71    /// Perplexity.
72    Perplexity,
73    /// GitHub Copilot.
74    Copilot,
75    /// Any other provider not explicitly listed.
76    Other(Cow<'static, str>),
77}
78
79impl Provider {
80    /// Returns the canonical string representation.
81    #[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/// Per-token pricing information (USD).
114// Every field is a per-token rate, so the shared suffix is the point rather
115// than redundant.
116#[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    /// Creates a new pricing entry.
129    #[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    /// Input cost per token.
142    #[must_use]
143    pub const fn input_per_token(&self) -> f64 {
144        self.input_per_token
145    }
146
147    /// Output cost per token.
148    #[must_use]
149    pub const fn output_per_token(&self) -> f64 {
150        self.output_per_token
151    }
152
153    /// Cache read cost per token, if applicable.
154    #[must_use]
155    pub const fn cache_read_per_token(&self) -> Option<f64> {
156        self.cache_read_per_token
157    }
158
159    /// Cache write cost per token, if applicable.
160    #[must_use]
161    pub const fn cache_write_per_token(&self) -> Option<f64> {
162        self.cache_write_per_token
163    }
164
165    /// Reasoning output cost per token, if applicable.
166    #[must_use]
167    pub const fn reasoning_per_token(&self) -> Option<f64> {
168        self.reasoning_per_token
169    }
170
171    /// Image input cost per token, if applicable.
172    #[must_use]
173    pub const fn image_per_token(&self) -> Option<f64> {
174        self.image_per_token
175    }
176
177    /// Sets cache read cost per token.
178    #[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    /// Sets cache write cost per token.
185    #[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    /// Sets reasoning output cost per token.
192    #[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    /// Sets image input cost per token.
199    #[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/// A single model entry from the upstream registry.
207#[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    /// Creates a new model entry with owned strings.
222    #[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    /// Construct from all fields (used by generated code).
243    #[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    /// The canonical model ID (provider prefix stripped).
274    #[must_use]
275    pub fn id(&self) -> &str {
276        &self.canonical_id
277    }
278
279    /// The original `LiteLLM` key (e.g. `"gemini/gemini-2.5-flash"`).
280    #[must_use]
281    pub fn litellm_id(&self) -> &str {
282        &self.litellm_id
283    }
284
285    /// The model provider.
286    #[must_use]
287    pub const fn provider(&self) -> &Provider {
288        &self.provider
289    }
290
291    /// The operational mode.
292    #[must_use]
293    pub const fn mode(&self) -> ModelMode {
294        self.mode
295    }
296
297    /// Maximum input tokens (context window), if known.
298    #[must_use]
299    pub const fn max_input_tokens(&self) -> Option<u32> {
300        self.max_input_tokens
301    }
302
303    /// Maximum output tokens, if known.
304    #[must_use]
305    pub const fn max_output_tokens(&self) -> Option<u32> {
306        self.max_output_tokens
307    }
308
309    /// Pricing information.
310    #[must_use]
311    pub const fn pricing(&self) -> &Pricing {
312        &self.pricing
313    }
314
315    /// Supported abilities.
316    #[must_use]
317    pub fn abilities(&self) -> &[Ability] {
318        &self.abilities
319    }
320
321    /// Whether this model has a specific ability.
322    #[must_use]
323    pub fn has_ability(&self, ability: Ability) -> bool {
324        self.abilities.contains(&ability)
325    }
326
327    /// Deprecation date string (YYYY-MM-DD), if set.
328    #[must_use]
329    pub fn deprecation_date(&self) -> Option<&str> {
330        self.deprecation_date.as_deref()
331    }
332
333    /// Whether this model is past its deprecation date.
334    #[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    /// Sets max input tokens.
349    #[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    /// Sets max output tokens.
356    #[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    /// Sets pricing.
363    #[must_use]
364    pub const fn with_pricing(mut self, pricing: Pricing) -> Self {
365        self.pricing = pricing;
366        self
367    }
368
369    /// Sets abilities from a vec.
370    #[must_use]
371    pub fn with_abilities(mut self, abilities: Vec<Ability>) -> Self {
372        self.abilities = Cow::Owned(abilities);
373        self
374    }
375
376    /// Sets deprecation date.
377    #[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
384/// Convert days since Unix epoch to (year, month, day).
385const fn days_to_ymd(days: u64) -> (u64, u64, u64) {
386    // Civil calendar algorithm from Howard Hinnant
387    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}