Skip to main content

aither_models/
convert.rs

1//! Shared conversion logic for `LiteLLM` and `OpenRouter` JSON.
2//!
3//! This module is included via `#[path]` in build.rs so the same parsing
4//! logic runs at compile time and runtime. It depends only on `serde`,
5//! `serde_json`, and sibling types.
6
7use serde::Deserialize;
8
9/// Raw `LiteLLM` entry as it appears in the JSON.
10#[derive(Debug, Deserialize)]
11pub struct RawLiteLLMEntry {
12    /// Provider slug as `LiteLLM` reports it.
13    #[serde(default)]
14    pub litellm_provider: Option<String>,
15    /// What the model does: `chat`, `embedding`, `image_generation`, and so on.
16    #[serde(default)]
17    pub mode: Option<String>,
18    /// Largest prompt the model accepts, in tokens.
19    #[serde(default)]
20    pub max_input_tokens: Option<u32>,
21    /// Largest completion the model will produce, in tokens.
22    #[serde(default)]
23    pub max_output_tokens: Option<u32>,
24    /// Combined limit, used when the input/output split is unknown.
25    #[serde(default)]
26    pub max_tokens: Option<u32>,
27    /// USD charged per input token.
28    #[serde(default)]
29    pub input_cost_per_token: Option<f64>,
30    /// USD charged per output token.
31    #[serde(default)]
32    pub output_cost_per_token: Option<f64>,
33    /// USD per input token served from the prompt cache.
34    #[serde(default)]
35    pub cache_read_input_token_cost: Option<f64>,
36    /// USD per input token written to the prompt cache.
37    #[serde(default)]
38    pub cache_creation_input_token_cost: Option<f64>,
39    /// USD per reasoning token, where billed separately.
40    #[serde(default)]
41    pub output_cost_per_reasoning_token: Option<f64>,
42    /// USD per audio input token.
43    #[serde(default)]
44    pub input_cost_per_audio_token: Option<f64>,
45    /// Date the provider retires the model, as `YYYY-MM-DD`.
46    #[serde(default)]
47    pub deprecation_date: Option<String>,
48    // supports_* flags
49    /// Whether the model can call tools.
50    #[serde(default)]
51    pub supports_function_calling: Option<bool>,
52    /// Whether the model accepts images.
53    #[serde(default)]
54    pub supports_vision: Option<bool>,
55    /// Whether the model accepts audio.
56    #[serde(default)]
57    pub supports_audio_input: Option<bool>,
58    /// Whether the model can emit audio.
59    #[serde(default)]
60    pub supports_audio_output: Option<bool>,
61    /// Whether the model accepts video.
62    #[serde(default)]
63    pub supports_video_input: Option<bool>,
64    /// Whether the provider offers built-in web search.
65    #[serde(default)]
66    pub supports_web_search: Option<bool>,
67    /// Whether the model accepts PDFs directly.
68    #[serde(default)]
69    pub supports_pdf_input: Option<bool>,
70    /// Whether the model exposes reasoning or thinking output.
71    #[serde(default)]
72    pub supports_reasoning: Option<bool>,
73    /// Whether the model supports computer-use tooling.
74    #[serde(default)]
75    pub supports_computer_use: Option<bool>,
76    /// Whether the provider caches prompt prefixes.
77    #[serde(default)]
78    pub supports_prompt_caching: Option<bool>,
79    /// Whether an assistant turn may be prefilled.
80    #[serde(default)]
81    pub supports_assistant_prefill: Option<bool>,
82}
83
84/// Known provider prefix → canonical provider mapping.
85const PROVIDER_PREFIXES: &[(&str, &str)] = &[
86    ("openai/", "openai"),
87    ("anthropic/", "anthropic"),
88    ("gemini/", "gemini"),
89    ("deepseek/", "deepseek"),
90    ("xai/", "xai"),
91    ("mistral/", "mistral"),
92    ("meta-llama/", "meta"),
93    ("bedrock/", "bedrock"),
94    ("azure/", "azure"),
95    ("azure_ai/", "azure"),
96    ("vertex_ai/", "vertex_ai"),
97    ("groq/", "groq"),
98    ("together_ai/", "together_ai"),
99    ("fireworks_ai/", "fireworks_ai"),
100    ("replicate/", "replicate"),
101    ("cohere/", "cohere"),
102    ("cohere_chat/", "cohere"),
103    ("perplexity/", "perplexity"),
104    ("github_copilot/", "github_copilot"),
105    ("ollama/", "ollama"),
106    ("nvidia_nim/", "nvidia_nim"),
107    ("sambanova/", "sambanova"),
108    ("cerebras/", "cerebras"),
109    ("dashscope/", "alibaba"),
110];
111
112/// Strip known provider prefixes to get the canonical model ID.
113#[must_use]
114pub fn strip_provider_prefix(litellm_key: &str) -> &str {
115    for &(prefix, _) in PROVIDER_PREFIXES {
116        if let Some(rest) = litellm_key.strip_prefix(prefix) {
117            return rest;
118        }
119    }
120    litellm_key
121}
122
123/// Map a `LiteLLM` `litellm_provider` string to a provider tag.
124#[must_use]
125pub fn parse_provider(provider_str: &str) -> String {
126    let mapped = match provider_str {
127        "openai" | "text-completion-openai" | "chatgpt" => "openai",
128        "anthropic" => "anthropic",
129        "gemini" | "palm" => "gemini",
130        "deepseek" => "deepseek",
131        "xai" => "xai",
132        "mistral" | "codestral" | "text-completion-codestral" => "mistral",
133        "meta_llama" => "meta",
134        s if s.starts_with("bedrock") => "bedrock",
135        "azure" | "azure_ai" | "azure_text" => "azure",
136        s if s.starts_with("vertex_ai") => "vertex_ai",
137        "groq" => "groq",
138        "together_ai" => "together_ai",
139        s if s.starts_with("fireworks_ai") => "fireworks_ai",
140        "replicate" => "replicate",
141        "cohere" | "cohere_chat" => "cohere",
142        "perplexity" => "perplexity",
143        "github_copilot" => "github_copilot",
144        "dashscope" => "alibaba",
145        other => return other.to_string(),
146    };
147    mapped.to_string()
148}
149
150/// Map a `LiteLLM` `mode` string to a mode tag.
151#[must_use]
152pub fn parse_mode(mode_str: &str) -> Option<&'static str> {
153    match mode_str {
154        "chat" | "responses" => Some("chat"),
155        "embedding" => Some("embedding"),
156        "image_generation" | "image_edit" => Some("image_generation"),
157        "audio_transcription" => Some("audio_transcription"),
158        "audio_speech" => Some("audio_speech"),
159        "rerank" => Some("rerank"),
160        "moderation" => Some("moderation"),
161        "completion" => Some("completion"),
162        "video_generation" => Some("video_generation"),
163        "search" | "vector_store" => Some("search"),
164        "ocr" => Some("ocr"),
165        _ => None,
166    }
167}
168
169/// Intermediate representation used by both build.rs and runtime.
170#[derive(Debug, Clone)]
171pub struct ConvertedEntry {
172    /// Model id exactly as `LiteLLM` keys it, provider prefix included.
173    pub litellm_id: String,
174    /// Model id with the provider prefix stripped.
175    pub canonical_id: String,
176    /// Canonical provider slug.
177    pub provider: String,
178    /// What the model does, as a canonical mode string.
179    pub mode: String,
180    /// Largest prompt the model accepts, in tokens.
181    pub max_input_tokens: Option<u32>,
182    /// Largest completion the model produces, in tokens.
183    pub max_output_tokens: Option<u32>,
184    /// USD per input token.
185    pub input_cost_per_token: f64,
186    /// USD per output token.
187    pub output_cost_per_token: f64,
188    /// USD per cached input token read.
189    pub cache_read_per_token: Option<f64>,
190    /// USD per input token written to the cache.
191    pub cache_write_per_token: Option<f64>,
192    /// USD per reasoning token, where billed separately.
193    pub reasoning_per_token: Option<f64>,
194    /// USD per image token.
195    pub image_per_token: Option<f64>,
196    /// Capability names this model advertises.
197    pub abilities: Vec<&'static str>,
198    /// Retirement date, as `YYYY-MM-DD`.
199    pub deprecation_date: Option<String>,
200}
201
202/// Convert a single `LiteLLM` JSON entry into our intermediate representation.
203#[must_use]
204pub fn convert_litellm_entry(key: &str, raw: &RawLiteLLMEntry) -> Option<ConvertedEntry> {
205    let provider_str = raw.litellm_provider.as_deref()?;
206    let mode_str = raw.mode.as_deref()?;
207    let mode = parse_mode(mode_str)?;
208    let provider = parse_provider(provider_str);
209    let canonical_id = strip_provider_prefix(key).to_string();
210
211    let max_input = raw.max_input_tokens;
212    let max_output = raw.max_output_tokens.or(raw.max_tokens);
213
214    let mut abilities = Vec::new();
215    if raw.supports_function_calling == Some(true) {
216        abilities.push("ToolUse");
217    }
218    if raw.supports_vision == Some(true) {
219        abilities.push("Vision");
220    }
221    if raw.supports_audio_input == Some(true) {
222        abilities.push("Audio");
223    }
224    if raw.supports_audio_output == Some(true) {
225        abilities.push("AudioOutput");
226    }
227    if raw.supports_video_input == Some(true) {
228        abilities.push("Video");
229    }
230    if raw.supports_web_search == Some(true) {
231        abilities.push("WebSearch");
232    }
233    if raw.supports_pdf_input == Some(true) {
234        abilities.push("Pdf");
235    }
236    if raw.supports_reasoning == Some(true) {
237        abilities.push("Reasoning");
238    }
239    if raw.supports_computer_use == Some(true) {
240        abilities.push("ComputerUse");
241    }
242    if raw.supports_prompt_caching == Some(true) {
243        abilities.push("PromptCaching");
244    }
245    if raw.supports_assistant_prefill == Some(true) {
246        abilities.push("AssistantPrefill");
247    }
248    if mode == "image_generation" {
249        abilities.push("ImageGeneration");
250    }
251
252    Some(ConvertedEntry {
253        litellm_id: key.to_string(),
254        canonical_id,
255        provider,
256        mode: mode.to_string(),
257        max_input_tokens: max_input,
258        max_output_tokens: max_output,
259        input_cost_per_token: raw.input_cost_per_token.unwrap_or(0.0),
260        output_cost_per_token: raw.output_cost_per_token.unwrap_or(0.0),
261        cache_read_per_token: raw.cache_read_input_token_cost,
262        cache_write_per_token: raw.cache_creation_input_token_cost,
263        reasoning_per_token: raw.output_cost_per_reasoning_token,
264        image_per_token: raw.input_cost_per_audio_token,
265        abilities,
266        deprecation_date: raw.deprecation_date.clone(),
267    })
268}
269
270/// Parse the full `LiteLLM` JSON blob and return converted entries.
271///
272/// # Panics
273///
274/// Panics if `json_bytes` is not a JSON object. The only caller supplies either
275/// the committed snapshot or a payload already validated by the caller.
276#[must_use]
277pub fn parse_litellm_json(json_bytes: &[u8]) -> Vec<ConvertedEntry> {
278    let map: serde_json::Map<String, serde_json::Value> =
279        serde_json::from_slice(json_bytes).expect("invalid LiteLLM JSON");
280
281    let mut entries = Vec::new();
282    for (key, value) in &map {
283        if key == "sample_spec" {
284            continue;
285        }
286        let Ok(raw) = serde_json::from_value::<RawLiteLLMEntry>(value.clone()) else {
287            continue;
288        };
289        if let Some(entry) = convert_litellm_entry(key, &raw) {
290            entries.push(entry);
291        }
292    }
293    entries
294}