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: memory,
166 sha256,
167 digest_verified: None,
168 projectors,
169 })
170}
171pub fn parse_candidates(raw: &str) -> Result<Vec<RawModel>, ValidationError> {
172 let candidates: Vec<RawModel> = parse_document(raw)?;
173 require(
174 candidates.len() <= 10000,
175 "registry candidate limit exceeded",
176 )?;
177 for candidate in &candidates {
178 candidate.validate_shape()?;
179 }
180 Ok(candidates)
181}
182#[derive(Default, Debug, Serialize)]
183pub struct EnrichDiff {
184 pub added: Vec<String>,
185 pub updated: Vec<String>,
186 pub removed: Vec<String>,
187 pub skipped: Vec<String>,
188 pub capped: Vec<String>,
189}
190#[derive(Serialize)]
191pub struct EnrichResult {
192 pub catalog: Catalog,
193 pub diff: EnrichDiff,
194}
195pub fn enrich(
196 existing: &Catalog,
197 candidates: &[RawModel],
198 mode: Mode,
199 now: &str,
200 maximum: Option<usize>,
201) -> Result<EnrichResult, ValidationError> {
202 require(
203 (2..=crate::catalog::SCHEMA_VERSION).contains(&existing.schema_version)
204 && maximum != Some(0)
205 && candidates.len() <= 10000,
206 "invalid enrichment options",
207 )?;
208 require(
209 existing.generated_at.ends_with('Z')
210 && OffsetDateTime::parse(&existing.generated_at, &Rfc3339).is_ok(),
211 "invalid prior catalog timestamp",
212 )?;
213 let mut ids = std::collections::HashSet::new();
214 for model in &existing.models {
215 model.validate()?;
216 require(ids.insert(model.id.as_str()), "duplicate prior catalog id")?;
217 }
218 let now = OffsetDateTime::parse(now, &Rfc3339)
219 .map_err(|_| ValidationError("invalid enrichment clock".into()))?
220 .to_offset(time::UtcOffset::UTC);
221 let stamp = now
222 .format(time::macros::format_description!(
223 "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
224 ))
225 .map_err(|_| ValidationError("invalid enrichment clock".into()))?;
226 let today = &stamp[..10];
227 let newest = existing
228 .models
229 .iter()
230 .filter_map(|model| model.recency().map(|(day, _)| day))
231 .max()
232 .unwrap_or("");
233 let mut result: BTreeMap<String, CatalogModel> = existing
234 .models
235 .iter()
236 .map(|model| (model.id.clone(), model.clone()))
237 .collect();
238 let mut positions = BTreeMap::new();
239 let mut unique: Vec<&RawModel> = Vec::new();
240 for candidate in candidates {
241 candidate.validate_shape()?;
242 if let Some(index) = positions.get(&candidate.id) {
243 unique[*index] = candidate;
244 } else {
245 positions.insert(candidate.id.clone(), unique.len());
246 unique.push(candidate);
247 }
248 }
249 let mut diff = EnrichDiff::default();
250 for raw in unique {
251 let prior = result.get(&raw.id);
252 if !raw.open_weight || !LICENSES.contains(&raw.license.as_str()) {
253 if result.remove(&raw.id).is_some() {
254 diff.removed.push(raw.id.clone());
255 } else if mode == Mode::Backfill {
256 diff.skipped.push(raw.id.clone());
257 }
258 continue;
259 }
260 if mode == Mode::Incremental && prior.is_some() {
261 continue;
262 }
263 if raw.release_date.as_str() > today {
264 if prior.is_none() {
265 diff.skipped.push(raw.id.clone());
266 }
267 continue;
268 }
269 if mode == Mode::Incremental && raw.release_date.as_str() <= newest {
270 continue;
271 }
272 match raw.build(prior) {
273 Ok(built) => {
274 if let Some(prior) = prior {
275 let encoded = serde_json::to_value(&built)
276 .map_err(|error| ValidationError(error.to_string()))?;
277 let previous = serde_json::to_value(prior)
278 .map_err(|error| ValidationError(error.to_string()))?;
279 if encoded == previous {
280 continue;
281 }
282 diff.updated.push(raw.id.clone());
283 } else {
284 diff.added.push(raw.id.clone());
285 }
286 result.insert(raw.id.clone(), built);
287 }
288 Err(_) => diff.skipped.push(raw.id.clone()),
289 }
290 }
291 let mut models: Vec<_> = result.into_values().collect();
292 models.sort_by(|left, right| {
293 let day = |model: &CatalogModel| model.recency().map(|(day, _)| day.to_string());
294 day(right)
295 .cmp(&day(left))
296 .then_with(|| left.id.cmp(&right.id))
297 });
298 if let Some(maximum) = maximum
299 && models.len() > maximum
300 {
301 diff.capped = models.drain(maximum..).map(|model| model.id).collect();
302 }
303 let changed = !diff.added.is_empty()
304 || !diff.updated.is_empty()
305 || !diff.removed.is_empty()
306 || !diff.capped.is_empty();
307 diff.added.retain(|id| !diff.capped.contains(id));
308 diff.updated.retain(|id| !diff.capped.contains(id));
309 Ok(EnrichResult {
310 catalog: Catalog {
311 schema_version: crate::catalog::SCHEMA_VERSION,
312 generated_at: if changed {
313 stamp
314 } else {
315 existing.generated_at.clone()
316 },
317 models,
318 },
319 diff,
320 })
321}
322pub fn format_diff(diff: &EnrichDiff) -> String {
323 let mut text = format!(
324 "Refresh (dry-run):\n added: {}\n updated: {}\n removed: {}\n skipped: {}\n capped: {}\n",
325 diff.added.len(),
326 diff.updated.len(),
327 diff.removed.len(),
328 diff.skipped.len(),
329 diff.capped.len()
330 );
331 for (label, ids) in [
332 ("added ids", &diff.added),
333 ("updated ids", &diff.updated),
334 ("removed ids", &diff.removed),
335 ] {
336 if !ids.is_empty() {
337 let mut ids = ids.clone();
338 ids.sort();
339 text.push_str(&format!(" {label}: {}\n", ids.join(", ")));
340 }
341 }
342 text.push_str(" No catalog file was written.\n\n");
343 text
344}