1use serde_json::Value;
2use std::collections::{HashMap, HashSet};
3use std::sync::OnceLock;
4
5use crate::error::{LiteLLMError, Result};
6
7#[derive(Debug, Clone)]
8pub struct ModelPricing {
9 pub input_cost_per_1k: Option<f64>,
10 pub output_cost_per_1k: Option<f64>,
11 pub max_input_tokens: Option<u32>,
12 pub max_output_tokens: Option<u32>,
13 pub mode: Option<String>,
14 pub provider: Option<String>,
15}
16
17#[derive(Debug, Clone)]
18pub struct Registry {
19 pub models: HashMap<String, ModelPricing>,
20}
21
22impl Registry {
23 pub fn load_embedded() -> Result<Self> {
24 Self::from_json_str(include_str!("../data/model_prices_and_context_window.json"))
25 }
26
27 pub fn from_json_str(raw: &str) -> Result<Self> {
28 let json: Value = serde_json::from_str(raw)
29 .map_err(|e| LiteLLMError::Parse(format!("model registry: {e}")))?;
30 let mut models = HashMap::new();
31 let map = json
32 .as_object()
33 .ok_or_else(|| LiteLLMError::Parse("model registry not an object".into()))?;
34 for (name, entry) in map {
35 if name == "sample_spec" {
36 continue;
37 }
38 if let Some(obj) = entry.as_object() {
39 let input = obj
40 .get("input_cost_per_token")
41 .and_then(|v| v.as_f64())
42 .map(|v| v * 1000.0);
43 let output = obj
44 .get("output_cost_per_token")
45 .and_then(|v| v.as_f64())
46 .map(|v| v * 1000.0);
47 let max_input = obj
48 .get("max_input_tokens")
49 .or_else(|| obj.get("max_tokens"))
50 .and_then(|v| v.as_u64())
51 .and_then(|v| u32::try_from(v).ok());
52 let max_output = obj
53 .get("max_output_tokens")
54 .and_then(|v| v.as_u64())
55 .and_then(|v| u32::try_from(v).ok());
56 let mode = obj
57 .get("mode")
58 .and_then(|v| v.as_str())
59 .map(|s| s.to_string());
60 let provider = obj
61 .get("litellm_provider")
62 .and_then(|v| v.as_str())
63 .map(|s| s.to_string());
64 models.insert(
65 name.to_string(),
66 ModelPricing {
67 input_cost_per_1k: input,
68 output_cost_per_1k: output,
69 max_input_tokens: max_input,
70 max_output_tokens: max_output,
71 mode,
72 provider,
73 },
74 );
75 }
76 }
77 Ok(Self { models })
78 }
79
80 pub fn get(&self, model: &str) -> Option<&ModelPricing> {
81 self.models.get(model)
82 }
83
84 pub fn estimate_cost(&self, model: &str, input_tokens: u32, output_tokens: u32) -> Option<f64> {
85 let pricing = self.models.get(model)?;
86 let input = pricing
87 .input_cost_per_1k
88 .map(|v| v * input_tokens as f64 / 1000.0)?;
89 let output = pricing
90 .output_cost_per_1k
91 .map(|v| v * output_tokens as f64 / 1000.0)?;
92 Some(input + output)
93 }
94}
95
96pub(crate) fn embedded_model_supports_output_config(model: &str) -> bool {
97 static SUPPORTED_MODELS: OnceLock<HashSet<String>> = OnceLock::new();
98
99 let supported = SUPPORTED_MODELS.get_or_init(|| {
100 let raw = include_str!("../data/model_prices_and_context_window.json");
101 let Ok(Value::Object(models)) = serde_json::from_str::<Value>(raw) else {
102 return HashSet::new();
103 };
104 models
105 .into_iter()
106 .filter_map(|(name, metadata)| {
107 metadata
108 .get("supports_output_config")
109 .and_then(Value::as_bool)
110 .unwrap_or(false)
111 .then(|| name.to_ascii_lowercase().replace('.', "-"))
112 })
113 .collect()
114 });
115
116 let normalized = model.to_ascii_lowercase().replace('.', "-");
117 if supported.contains(&normalized) {
118 return true;
119 }
120
121 let Some((base, suffix)) = normalized.rsplit_once('-') else {
122 return false;
123 };
124 suffix.len() == 8
125 && suffix.bytes().all(|byte| byte.is_ascii_digit())
126 && supported.contains(base)
127}
128
129#[cfg(test)]
130mod tests {
131 use super::*;
132
133 #[test]
134 fn loads_registry() {
135 let registry = Registry::load_embedded().unwrap();
136 assert!(!registry.models.is_empty());
137 }
138
139 #[test]
140 fn parses_registry_from_json_string() {
141 let registry = Registry::from_json_str(
142 r#"{
143 "sample_spec": {},
144 "test/model": {
145 "input_cost_per_token": 0.000001,
146 "output_cost_per_token": 0.000002,
147 "max_input_tokens": 4096,
148 "max_output_tokens": 1024,
149 "mode": "chat",
150 "litellm_provider": "test"
151 }
152 }"#,
153 )
154 .unwrap();
155
156 let model = registry.get("test/model").unwrap();
157 assert_eq!(model.input_cost_per_1k, Some(0.001));
158 assert_eq!(model.output_cost_per_1k, Some(0.002));
159 assert_eq!(model.max_input_tokens, Some(4096));
160 assert_eq!(model.max_output_tokens, Some(1024));
161 }
162
163 #[test]
164 fn token_counts_beyond_u32_are_dropped_not_truncated() {
165 let registry = Registry::from_json_str(
166 r#"{
167 "sample_spec": {},
168 "test/overflow": {
169 "max_input_tokens": 5000000000,
170 "max_output_tokens": 4294967296,
171 "mode": "chat"
172 }
173 }"#,
174 )
175 .unwrap();
176
177 let model = registry.get("test/overflow").unwrap();
178 assert_eq!(model.max_input_tokens, None);
179 assert_eq!(model.max_output_tokens, None);
180 }
181
182 #[test]
183 fn embedded_registry_retains_metadata_coverage() {
184 let registry = Registry::load_embedded().unwrap();
185 let priced_models = registry
186 .models
187 .values()
188 .filter(|model| model.output_cost_per_1k.is_some())
189 .count();
190 let models_with_context = registry
191 .models
192 .values()
193 .filter(|model| model.max_input_tokens.is_some())
194 .count();
195
196 assert!(
197 priced_models > 1000,
198 "expected broad completion-pricing coverage, found {priced_models} models"
199 );
200 assert!(
201 models_with_context > 1000,
202 "expected broad context-window coverage, found {models_with_context} models"
203 );
204 }
205
206 #[test]
207 fn embedded_capabilities_drive_anthropic_output_config_support() {
208 assert!(embedded_model_supports_output_config("claude-sonnet-5"));
209 assert!(embedded_model_supports_output_config("claude-opus-4.8"));
210 assert!(!embedded_model_supports_output_config(
211 "claude-3-7-sonnet-20250219"
212 ));
213 assert!(!embedded_model_supports_output_config(
214 "claude-3-5-sonnet-20241022"
215 ));
216 }
217
218 #[test]
219 fn embedded_capabilities_retry_without_one_trailing_date_token() {
220 assert!(embedded_model_supports_output_config(
221 "claude-opus-4.8-20260315"
222 ));
223 }
224
225 #[test]
226 fn embedded_capabilities_do_not_promote_unflagged_dated_models() {
227 assert!(!embedded_model_supports_output_config(
228 "claude-haiku-4-5-20260711"
229 ));
230 }
231
232 #[test]
233 fn embedded_capabilities_leave_non_dated_misses_unsupported() {
234 assert!(!embedded_model_supports_output_config(
235 "claude-sonnet-5-latest"
236 ));
237 }
238
239 #[test]
240 fn embedded_capabilities_strip_at_most_one_trailing_date_token() {
241 assert!(!embedded_model_supports_output_config(
242 "claude-sonnet-5-20260710-20260711"
243 ));
244 }
245}