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 .map(|v| v as u32);
52 let max_output = obj
53 .get("max_output_tokens")
54 .and_then(|v| v.as_u64())
55 .map(|v| v as u32);
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 supported.contains(&model.to_ascii_lowercase().replace('.', "-"))
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122
123 #[test]
124 fn loads_registry() {
125 let registry = Registry::load_embedded().unwrap();
126 assert!(!registry.models.is_empty());
127 }
128
129 #[test]
130 fn parses_registry_from_json_string() {
131 let registry = Registry::from_json_str(
132 r#"{
133 "sample_spec": {},
134 "test/model": {
135 "input_cost_per_token": 0.000001,
136 "output_cost_per_token": 0.000002,
137 "max_input_tokens": 4096,
138 "max_output_tokens": 1024,
139 "mode": "chat",
140 "litellm_provider": "test"
141 }
142 }"#,
143 )
144 .unwrap();
145
146 let model = registry.get("test/model").unwrap();
147 assert_eq!(model.input_cost_per_1k, Some(0.001));
148 assert_eq!(model.output_cost_per_1k, Some(0.002));
149 assert_eq!(model.max_input_tokens, Some(4096));
150 assert_eq!(model.max_output_tokens, Some(1024));
151 }
152
153 #[test]
154 fn embedded_registry_retains_metadata_coverage() {
155 let registry = Registry::load_embedded().unwrap();
156 let priced_models = registry
157 .models
158 .values()
159 .filter(|model| model.output_cost_per_1k.is_some())
160 .count();
161 let models_with_context = registry
162 .models
163 .values()
164 .filter(|model| model.max_input_tokens.is_some())
165 .count();
166
167 assert!(
168 priced_models > 1000,
169 "expected broad completion-pricing coverage, found {priced_models} models"
170 );
171 assert!(
172 models_with_context > 1000,
173 "expected broad context-window coverage, found {models_with_context} models"
174 );
175 }
176
177 #[test]
178 fn embedded_capabilities_drive_anthropic_output_config_support() {
179 assert!(embedded_model_supports_output_config("claude-sonnet-5"));
180 assert!(embedded_model_supports_output_config("claude-opus-4.8"));
181 assert!(!embedded_model_supports_output_config(
182 "claude-3-7-sonnet-20250219"
183 ));
184 assert!(!embedded_model_supports_output_config(
185 "claude-3-5-sonnet-20241022"
186 ));
187 }
188}