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 quantizations.push(sized_quantization(
94 count,
95 &self.architecture,
96 &strip_control(&quant.name),
97 quant.disk_bytes,
98 quant.sha256.as_ref().map(|sha| strip_control(sha)),
99 quant.projectors.clone(),
100 )?);
101 }
102 let model = CatalogModel {
103 id: strip_control(&self.id),
104 family: strip_control(&self.family),
105 params: strip_control(&self.params),
106 architecture: self.architecture.clone(),
107 active_params: self
108 .active_params
109 .as_ref()
110 .map(|value| strip_control(value)),
111 license: self.license.clone(),
112 open_weight: self.open_weight,
113 context_length: self.context_length,
114 capabilities: self.capabilities.clone(),
115 release_date: Some(self.release_date.clone()),
116 added_at: None,
117 provenance: crate::catalog::EntryProvenance::Curated,
118 source: Source {
119 ollama: self
120 .source
121 .ollama
122 .as_ref()
123 .map(|value| strip_control(value)),
124 hf: self.source.hf.as_ref().map(|value| strip_control(value)),
125 ..Default::default()
126 },
127 quantizations,
128 kv_bytes_per_token: prior.and_then(|model| model.kv_bytes_per_token),
129 benchmark_proxy: prior.and_then(|model| model.benchmark_proxy),
130 };
131 model.validate()?;
132 Ok(model)
133 }
134}
135pub fn sized_quantization(
137 params: f64,
138 architecture: &Architecture,
139 name: &str,
140 disk_bytes: f64,
141 sha256: Option<String>,
142 projectors: Vec<crate::sizing::ProjectorArtifact>,
143) -> Result<Quantization, ValidationError> {
144 let bits = quant_bits(name);
145 require(
146 bits.is_some() || !matches!(architecture, Architecture::Moe),
147 "unknown MoE quantization",
148 )?;
149 let projector_bytes = projectors
150 .iter()
151 .map(|projector| projector.bytes as f64)
152 .sum::<f64>();
153 require(
154 projector_bytes < disk_bytes,
155 "projector size exceeds aggregate weights",
156 )?;
157 let resident = (disk_bytes - projector_bytes)
158 .max(bits.map(|bits| (params * bits / 8.0).ceil()).unwrap_or(0.0))
159 + projector_bytes;
160 let memory = resident + (resident * 0.15).ceil();
161 Ok(Quantization {
162 name: name.into(),
163 disk_bytes,
164 min_ram_bytes: memory,
165 min_vram_bytes: if name.eq_ignore_ascii_case("I2_S") {
166 0.0
167 } else {
168 memory
169 },
170 sha256,
171 digest_verified: None,
172 projectors,
173 })
174}
175pub fn parse_candidates(raw: &str) -> Result<Vec<RawModel>, ValidationError> {
176 let candidates: Vec<RawModel> = parse_document(raw)?;
177 require(
178 candidates.len() <= 10000,
179 "registry candidate limit exceeded",
180 )?;
181 for candidate in &candidates {
182 candidate.validate_shape()?;
183 }
184 Ok(candidates)
185}
186#[derive(Default, Debug, Serialize)]
187pub struct EnrichDiff {
188 pub added: Vec<String>,
189 pub updated: Vec<String>,
190 pub removed: Vec<String>,
191 pub skipped: Vec<String>,
192 pub capped: Vec<String>,
193}
194#[derive(Serialize)]
195pub struct EnrichResult {
196 pub catalog: Catalog,
197 pub diff: EnrichDiff,
198}
199pub fn enrich(
200 existing: &Catalog,
201 candidates: &[RawModel],
202 mode: Mode,
203 now: &str,
204 maximum: Option<usize>,
205) -> Result<EnrichResult, ValidationError> {
206 require(
207 (2..=crate::catalog::SCHEMA_VERSION).contains(&existing.schema_version)
208 && maximum != Some(0)
209 && candidates.len() <= 10000,
210 "invalid enrichment options",
211 )?;
212 require(
213 existing.generated_at.ends_with('Z')
214 && OffsetDateTime::parse(&existing.generated_at, &Rfc3339).is_ok(),
215 "invalid prior catalog timestamp",
216 )?;
217 let mut ids = std::collections::HashSet::new();
218 for model in &existing.models {
219 model.validate()?;
220 require(ids.insert(model.id.as_str()), "duplicate prior catalog id")?;
221 }
222 let now = OffsetDateTime::parse(now, &Rfc3339)
223 .map_err(|_| ValidationError("invalid enrichment clock".into()))?
224 .to_offset(time::UtcOffset::UTC);
225 let stamp = now
226 .format(time::macros::format_description!(
227 "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
228 ))
229 .map_err(|_| ValidationError("invalid enrichment clock".into()))?;
230 let today = &stamp[..10];
231 let newest = existing
232 .models
233 .iter()
234 .filter_map(|model| model.recency().map(|(day, _)| day))
235 .max()
236 .unwrap_or("");
237 let mut result: BTreeMap<String, CatalogModel> = existing
238 .models
239 .iter()
240 .map(|model| (model.id.clone(), model.clone()))
241 .collect();
242 let mut positions = BTreeMap::new();
243 let mut unique: Vec<&RawModel> = Vec::new();
244 for candidate in candidates {
245 candidate.validate_shape()?;
246 if let Some(index) = positions.get(&candidate.id) {
247 unique[*index] = candidate;
248 } else {
249 positions.insert(candidate.id.clone(), unique.len());
250 unique.push(candidate);
251 }
252 }
253 let mut diff = EnrichDiff::default();
254 for raw in unique {
255 let prior = result.get(&raw.id);
256 if !raw.open_weight || !LICENSES.contains(&raw.license.as_str()) {
257 if result.remove(&raw.id).is_some() {
258 diff.removed.push(raw.id.clone());
259 } else if mode == Mode::Backfill {
260 diff.skipped.push(raw.id.clone());
261 }
262 continue;
263 }
264 if mode == Mode::Incremental && prior.is_some() {
265 continue;
266 }
267 if raw.release_date.as_str() > today {
268 if prior.is_none() {
269 diff.skipped.push(raw.id.clone());
270 }
271 continue;
272 }
273 if mode == Mode::Incremental && raw.release_date.as_str() <= newest {
274 continue;
275 }
276 match raw.build(prior) {
277 Ok(built) => {
278 if let Some(prior) = prior {
279 let encoded = serde_json::to_value(&built)
280 .map_err(|error| ValidationError(error.to_string()))?;
281 let previous = serde_json::to_value(prior)
282 .map_err(|error| ValidationError(error.to_string()))?;
283 if encoded == previous {
284 continue;
285 }
286 diff.updated.push(raw.id.clone());
287 } else {
288 diff.added.push(raw.id.clone());
289 }
290 result.insert(raw.id.clone(), built);
291 }
292 Err(_) => diff.skipped.push(raw.id.clone()),
293 }
294 }
295 let mut models: Vec<_> = result.into_values().collect();
296 models.sort_by(|left, right| {
297 let day = |model: &CatalogModel| model.recency().map(|(day, _)| day.to_string());
298 day(right)
299 .cmp(&day(left))
300 .then_with(|| left.id.cmp(&right.id))
301 });
302 if let Some(maximum) = maximum
303 && models.len() > maximum
304 {
305 diff.capped = models.drain(maximum..).map(|model| model.id).collect();
306 }
307 let changed = !diff.added.is_empty()
308 || !diff.updated.is_empty()
309 || !diff.removed.is_empty()
310 || !diff.capped.is_empty();
311 diff.added.retain(|id| !diff.capped.contains(id));
312 diff.updated.retain(|id| !diff.capped.contains(id));
313 Ok(EnrichResult {
314 catalog: Catalog {
315 schema_version: crate::catalog::SCHEMA_VERSION,
316 generated_at: if changed {
317 stamp
318 } else {
319 existing.generated_at.clone()
320 },
321 models,
322 },
323 diff,
324 })
325}
326pub fn format_diff(diff: &EnrichDiff) -> String {
327 let mut text = format!(
328 "Refresh (dry-run):\n added: {}\n updated: {}\n removed: {}\n skipped: {}\n capped: {}\n",
329 diff.added.len(),
330 diff.updated.len(),
331 diff.removed.len(),
332 diff.skipped.len(),
333 diff.capped.len()
334 );
335 for (label, ids) in [
336 ("added ids", &diff.added),
337 ("updated ids", &diff.updated),
338 ("removed ids", &diff.removed),
339 ] {
340 if !ids.is_empty() {
341 let mut ids = ids.clone();
342 ids.sort();
343 text.push_str(&format!(" {label}: {}\n", ids.join(", ")));
344 }
345 }
346 text.push_str(" No catalog file was written.\n\n");
347 text
348}