1use crate::{
2 catalog::{Catalog, CatalogModel, LICENSES, Source, parse_document, require},
3 reports::strip_control,
4 sizing::{Architecture, Quantization, ValidationError, parse_param_count, quant_bits},
5};
6use serde::{Deserialize, Serialize};
7use std::collections::BTreeMap;
8use time::{OffsetDateTime, format_description::well_known::Rfc3339};
9
10#[derive(Clone, Copy, PartialEq, Eq, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum Mode {
13 Backfill,
14 Incremental,
15}
16#[derive(Clone, Deserialize, Serialize)]
17#[serde(rename_all = "camelCase", deny_unknown_fields)]
18pub struct RawQuant {
19 pub name: String,
20 pub disk_bytes: f64,
21 #[serde(skip_serializing_if = "Option::is_none")]
22 pub sha256: Option<String>,
23 #[serde(default, skip_serializing_if = "Vec::is_empty")]
24 pub projectors: Vec<crate::sizing::ProjectorArtifact>,
25}
26#[derive(Clone, Deserialize, Serialize)]
27#[serde(deny_unknown_fields)]
28pub struct RawSource {
29 #[serde(skip_serializing_if = "Option::is_none")]
30 pub ollama: Option<String>,
31 #[serde(skip_serializing_if = "Option::is_none")]
32 pub hf: Option<String>,
33}
34#[derive(Clone, Deserialize, Serialize)]
35#[serde(rename_all = "camelCase", deny_unknown_fields)]
36pub struct RawModel {
37 pub id: String,
38 pub family: String,
39 pub params: String,
40 pub architecture: Architecture,
41 #[serde(skip_serializing_if = "Option::is_none")]
42 pub active_params: Option<String>,
43 pub license: String,
44 pub open_weight: bool,
45 pub context_length: f64,
46 pub capabilities: Vec<String>,
47 pub release_date: String,
48 pub source: RawSource,
49 pub quantizations: Vec<RawQuant>,
50}
51impl RawModel {
52 fn validate_shape(&self) -> Result<(), ValidationError> {
53 let integer = |value: f64| {
54 value.is_finite()
55 && value > 0.0
56 && value.fract() == 0.0
57 && value <= 9_007_199_254_740_991.0
58 };
59 require(
60 [
61 &self.id,
62 &self.family,
63 &self.params,
64 &self.license,
65 &self.release_date,
66 ]
67 .iter()
68 .all(|text| !text.is_empty())
69 && self
70 .active_params
71 .as_ref()
72 .is_none_or(|text| !text.is_empty())
73 && integer(self.context_length)
74 && !self.capabilities.is_empty()
75 && self.capabilities.iter().all(|text| !text.is_empty())
76 && [&self.source.ollama, &self.source.hf]
77 .into_iter()
78 .flatten()
79 .all(|text| !text.is_empty())
80 && !self.quantizations.is_empty()
81 && self.quantizations.iter().all(|quant| {
82 !quant.name.is_empty()
83 && integer(quant.disk_bytes)
84 && quant.sha256.as_ref().is_none_or(|text| !text.is_empty())
85 }),
86 "invalid registry payload",
87 )
88 }
89 fn build(&self, prior: Option<&CatalogModel>) -> Result<CatalogModel, ValidationError> {
90 let count = parse_param_count(&self.params)?;
91 let mut quantizations = Vec::new();
92 for quant in &self.quantizations {
93 let bits = quant_bits(&quant.name);
94 require(
95 bits.is_some() || !matches!(self.architecture, Architecture::Moe),
96 "unknown MoE quantization",
97 )?;
98 let projector_bytes = quant
99 .projectors
100 .iter()
101 .map(|projector| projector.bytes as f64)
102 .sum::<f64>();
103 require(
104 projector_bytes < quant.disk_bytes,
105 "projector size exceeds aggregate weights",
106 )?;
107 let resident = (quant.disk_bytes - projector_bytes)
108 .max(bits.map(|bits| (count * bits / 8.0).ceil()).unwrap_or(0.0))
109 + projector_bytes;
110 let memory = resident + (resident * 0.15).ceil();
111 quantizations.push(Quantization {
112 name: strip_control(&quant.name),
113 disk_bytes: quant.disk_bytes,
114 min_ram_bytes: memory,
115 min_vram_bytes: memory,
116 sha256: quant.sha256.as_ref().map(|sha| strip_control(sha)),
117 digest_verified: None,
118 projectors: quant.projectors.clone(),
119 });
120 }
121 let model = CatalogModel {
122 id: strip_control(&self.id),
123 family: strip_control(&self.family),
124 params: strip_control(&self.params),
125 architecture: self.architecture.clone(),
126 active_params: self
127 .active_params
128 .as_ref()
129 .map(|value| strip_control(value)),
130 license: self.license.clone(),
131 open_weight: self.open_weight,
132 context_length: self.context_length,
133 capabilities: self.capabilities.clone(),
134 release_date: self.release_date.clone(),
135 source: Source {
136 ollama: self
137 .source
138 .ollama
139 .as_ref()
140 .map(|value| strip_control(value)),
141 hf: self.source.hf.as_ref().map(|value| strip_control(value)),
142 ..Default::default()
143 },
144 quantizations,
145 kv_bytes_per_token: prior.and_then(|model| model.kv_bytes_per_token),
146 benchmark_proxy: prior.and_then(|model| model.benchmark_proxy),
147 };
148 model.validate()?;
149 Ok(model)
150 }
151}
152pub fn parse_candidates(raw: &str) -> Result<Vec<RawModel>, ValidationError> {
153 let candidates: Vec<RawModel> = parse_document(raw)?;
154 require(
155 candidates.len() <= 10000,
156 "registry candidate limit exceeded",
157 )?;
158 for candidate in &candidates {
159 candidate.validate_shape()?;
160 }
161 Ok(candidates)
162}
163#[derive(Default, Debug, Serialize)]
164pub struct EnrichDiff {
165 pub added: Vec<String>,
166 pub updated: Vec<String>,
167 pub removed: Vec<String>,
168 pub skipped: Vec<String>,
169 pub capped: Vec<String>,
170}
171#[derive(Serialize)]
172pub struct EnrichResult {
173 pub catalog: Catalog,
174 pub diff: EnrichDiff,
175}
176pub fn enrich(
177 existing: &Catalog,
178 candidates: &[RawModel],
179 mode: Mode,
180 now: &str,
181 maximum: Option<usize>,
182) -> Result<EnrichResult, ValidationError> {
183 require(
184 existing.schema_version == 2 && maximum != Some(0) && candidates.len() <= 10000,
185 "invalid enrichment options",
186 )?;
187 require(
188 existing.generated_at.ends_with('Z')
189 && OffsetDateTime::parse(&existing.generated_at, &Rfc3339).is_ok(),
190 "invalid prior catalog timestamp",
191 )?;
192 let mut ids = std::collections::HashSet::new();
193 for model in &existing.models {
194 model.validate()?;
195 require(ids.insert(model.id.as_str()), "duplicate prior catalog id")?;
196 }
197 let now = OffsetDateTime::parse(now, &Rfc3339)
198 .map_err(|_| ValidationError("invalid enrichment clock".into()))?
199 .to_offset(time::UtcOffset::UTC);
200 let stamp = now
201 .format(time::macros::format_description!(
202 "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
203 ))
204 .map_err(|_| ValidationError("invalid enrichment clock".into()))?;
205 let today = &stamp[..10];
206 let newest = existing
207 .models
208 .iter()
209 .map(|model| model.release_date.as_str())
210 .max()
211 .unwrap_or("");
212 let mut result: BTreeMap<String, CatalogModel> = existing
213 .models
214 .iter()
215 .map(|model| (model.id.clone(), model.clone()))
216 .collect();
217 let mut positions = BTreeMap::new();
218 let mut unique: Vec<&RawModel> = Vec::new();
219 for candidate in candidates {
220 candidate.validate_shape()?;
221 if let Some(index) = positions.get(&candidate.id) {
222 unique[*index] = candidate;
223 } else {
224 positions.insert(candidate.id.clone(), unique.len());
225 unique.push(candidate);
226 }
227 }
228 let mut diff = EnrichDiff::default();
229 for raw in unique {
230 let prior = result.get(&raw.id);
231 if !raw.open_weight || !LICENSES.contains(&raw.license.as_str()) {
232 if result.remove(&raw.id).is_some() {
233 diff.removed.push(raw.id.clone());
234 } else if mode == Mode::Backfill {
235 diff.skipped.push(raw.id.clone());
236 }
237 continue;
238 }
239 if mode == Mode::Incremental && prior.is_some() {
240 continue;
241 }
242 if raw.release_date.as_str() > today {
243 if prior.is_none() {
244 diff.skipped.push(raw.id.clone());
245 }
246 continue;
247 }
248 if mode == Mode::Incremental && raw.release_date.as_str() <= newest {
249 continue;
250 }
251 match raw.build(prior) {
252 Ok(built) => {
253 if let Some(prior) = prior {
254 let encoded = serde_json::to_value(&built)
255 .map_err(|error| ValidationError(error.to_string()))?;
256 let previous = serde_json::to_value(prior)
257 .map_err(|error| ValidationError(error.to_string()))?;
258 if encoded == previous {
259 continue;
260 }
261 diff.updated.push(raw.id.clone());
262 } else {
263 diff.added.push(raw.id.clone());
264 }
265 result.insert(raw.id.clone(), built);
266 }
267 Err(_) => diff.skipped.push(raw.id.clone()),
268 }
269 }
270 let mut models: Vec<_> = result.into_values().collect();
271 models.sort_by(|left, right| {
272 right
273 .release_date
274 .cmp(&left.release_date)
275 .then_with(|| left.id.cmp(&right.id))
276 });
277 if let Some(maximum) = maximum
278 && models.len() > maximum
279 {
280 diff.capped = models.drain(maximum..).map(|model| model.id).collect();
281 }
282 let changed = !diff.added.is_empty()
283 || !diff.updated.is_empty()
284 || !diff.removed.is_empty()
285 || !diff.capped.is_empty();
286 diff.added.retain(|id| !diff.capped.contains(id));
287 diff.updated.retain(|id| !diff.capped.contains(id));
288 Ok(EnrichResult {
289 catalog: Catalog {
290 schema_version: 2,
291 generated_at: if changed {
292 stamp
293 } else {
294 existing.generated_at.clone()
295 },
296 models,
297 },
298 diff,
299 })
300}
301pub fn format_diff(diff: &EnrichDiff) -> String {
302 let mut text = format!(
303 "Refresh (dry-run):\n added: {}\n updated: {}\n removed: {}\n skipped: {}\n capped: {}\n",
304 diff.added.len(),
305 diff.updated.len(),
306 diff.removed.len(),
307 diff.skipped.len(),
308 diff.capped.len()
309 );
310 for (label, ids) in [
311 ("added ids", &diff.added),
312 ("updated ids", &diff.updated),
313 ("removed ids", &diff.removed),
314 ] {
315 if !ids.is_empty() {
316 let mut ids = ids.clone();
317 ids.sort();
318 text.push_str(&format!(" {label}: {}\n", ids.join(", ")));
319 }
320 }
321 text.push_str(" No catalog file was written.\n\n");
322 text
323}