hf_fetch_model/discover.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Model family discovery and search via the `HuggingFace` Hub API.
4//!
5//! Queries the HF Hub for popular models, extracts `model_type` metadata,
6//! compares against locally cached families, and fetches model card metadata.
7
8use std::collections::{BTreeMap, HashMap};
9use std::future::Future;
10use std::hash::BuildHasher;
11use std::sync::Arc;
12
13use serde::Deserialize;
14
15use crate::error::FetchError;
16
17/// A model found by searching the `HuggingFace` Hub.
18#[derive(Debug, Clone)]
19pub struct SearchResult {
20 /// The repository identifier (e.g., `"RWKV/RWKV7-Goose-World3-1.5B-HF"`).
21 pub model_id: String,
22 /// Total download count.
23 pub downloads: u64,
24 /// Library framework (e.g., `"transformers"`, `"peft"`, `"diffusers"`), if reported.
25 pub library_name: Option<String>,
26 /// Pipeline task tag (e.g., `"text-generation"`), if reported.
27 pub pipeline_tag: Option<String>,
28 /// Tags from the model's metadata (e.g., `["gguf", "conversational"]`).
29 pub tags: Vec<String>,
30}
31
32/// A model family discovered from the `HuggingFace` Hub.
33#[derive(Debug, Clone)]
34pub struct DiscoveredFamily {
35 /// The `model_type` identifier (e.g., `"gpt_neox"`, `"llama"`).
36 pub model_type: String,
37 /// The most-downloaded representative model for this family.
38 pub top_model: String,
39 /// Download count of the representative model.
40 pub downloads: u64,
41}
42
43/// JSON response structure for an individual model from the HF API.
44#[derive(Debug, Deserialize)]
45struct ApiModelEntry {
46 #[serde(rename = "modelId")]
47 model_id: String,
48 #[serde(default)]
49 downloads: u64,
50 #[serde(default)]
51 config: Option<ApiConfig>,
52 #[serde(default)]
53 library_name: Option<String>,
54 #[serde(default)]
55 pipeline_tag: Option<String>,
56 #[serde(default)]
57 tags: Vec<String>,
58}
59
60/// The `config` object embedded in a model API response.
61#[derive(Debug, Deserialize)]
62struct ApiConfig {
63 model_type: Option<String>,
64}
65
66/// Access control status of a model on the `HuggingFace` Hub.
67///
68/// Some models require users to accept license terms before downloading.
69/// The gating mode determines whether approval is automatic or manual.
70#[derive(Debug, Clone, PartialEq, Eq)]
71#[non_exhaustive]
72pub enum GateStatus {
73 /// No gate — anyone can download without restrictions.
74 Open,
75 /// Automatic approval after the user accepts terms on the Hub.
76 Auto,
77 /// Manual approval by the model author after the user requests access.
78 Manual,
79}
80
81impl GateStatus {
82 /// Returns `true` if the model requires accepting terms before download.
83 #[must_use]
84 pub const fn is_gated(&self) -> bool {
85 matches!(self, Self::Auto | Self::Manual)
86 }
87}
88
89impl std::fmt::Display for GateStatus {
90 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91 match self {
92 Self::Open => write!(f, "open"),
93 Self::Auto => write!(f, "auto"),
94 Self::Manual => write!(f, "manual"),
95 }
96 }
97}
98
99/// Metadata from a `HuggingFace` model card.
100///
101/// Extracted from the single-model API endpoint
102/// (`GET /api/models/{owner}/{model}`). All fields are optional
103/// because model cards may omit any of them.
104#[derive(Debug, Clone)]
105pub struct ModelCardMetadata {
106 /// SPDX license identifier (e.g., `"apache-2.0"`).
107 pub license: Option<String>,
108 /// Pipeline tag (e.g., `"text-generation"`).
109 pub pipeline_tag: Option<String>,
110 /// Tags associated with the model (e.g., `["pytorch", "safetensors"]`).
111 pub tags: Vec<String>,
112 /// Library name (e.g., `"transformers"`, `"vllm"`).
113 pub library_name: Option<String>,
114 /// Languages the model supports (e.g., `["en", "fr"]`).
115 pub languages: Vec<String>,
116 /// Access control status (open, auto-gated, or manually gated).
117 pub gated: GateStatus,
118}
119
120/// JSON response for a single model from `GET /api/models/{model_id}`.
121#[derive(Debug, Deserialize)]
122struct ApiModelDetail {
123 #[serde(default)]
124 pipeline_tag: Option<String>,
125 #[serde(default)]
126 tags: Vec<String>,
127 #[serde(default)]
128 library_name: Option<String>,
129 #[serde(default)]
130 gated: ApiGated,
131 #[serde(default, rename = "cardData")]
132 card_data: Option<ApiCardData>,
133}
134
135/// The `cardData` sub-object (parsed YAML front matter from the model README).
136#[derive(Debug, Deserialize)]
137struct ApiCardData {
138 #[serde(default)]
139 license: Option<String>,
140 #[serde(default)]
141 language: Option<ApiLanguage>,
142}
143
144/// Languages in `cardData` can be a single string or a list of strings.
145#[derive(Debug, Deserialize)]
146#[serde(untagged)]
147enum ApiLanguage {
148 Single(String),
149 Multiple(Vec<String>),
150}
151
152/// The `gated` field can be `false` (boolean) or a string like `"auto"` / `"manual"`.
153#[derive(Debug, Deserialize)]
154#[serde(untagged)]
155enum ApiGated {
156 Bool(bool),
157 Mode(String),
158}
159
160impl Default for ApiGated {
161 fn default() -> Self {
162 Self::Bool(false)
163 }
164}
165
166const PAGE_SIZE: usize = 100;
167const HF_API_BASE: &str = "https://huggingface.co/api/models";
168
169/// Queries the `HuggingFace` Hub API for top models by downloads
170/// and returns families not present in the local cache.
171///
172/// # Arguments
173///
174/// * `local_families` — Set of `model_type` values already cached locally.
175/// * `max_models` — Maximum number of models to scan (paginated in batches of 100).
176/// * `tag` — Optional tag filter (e.g., `"gguf"`, `"bitsandbytes"`). When set,
177/// only models carrying this tag contribute to family discovery.
178///
179/// # Errors
180///
181/// Returns [`FetchError::Http`] if any API request fails.
182pub async fn discover_new_families<S: BuildHasher>(
183 local_families: &std::collections::HashSet<String, S>,
184 max_models: usize,
185 tag: Option<&str>,
186) -> Result<Vec<DiscoveredFamily>, FetchError> {
187 let client = reqwest::Client::new();
188 let mut remote_families: BTreeMap<String, (String, u64)> = BTreeMap::new();
189 let mut offset: usize = 0;
190
191 while offset < max_models {
192 let page_limit = PAGE_SIZE.min(max_models.saturating_sub(offset));
193 let page_limit_str = page_limit.to_string();
194 let offset_str = offset.to_string();
195
196 // BORROW: explicit .as_str() instead of Deref coercion
197 let mut query_params: Vec<(&str, &str)> = vec![
198 ("config", "true"),
199 ("sort", "downloads"),
200 ("direction", "-1"),
201 ("limit", page_limit_str.as_str()),
202 ("offset", offset_str.as_str()),
203 ];
204 if let Some(t) = tag {
205 query_params.push(("filter", t));
206 }
207
208 let response = client
209 .get(HF_API_BASE)
210 .query(&query_params)
211 .send()
212 .await
213 .map_err(|e| FetchError::Http(e.to_string()))?;
214
215 if !response.status().is_success() {
216 return Err(FetchError::Http(format!(
217 "HF API returned status {}",
218 response.status()
219 )));
220 }
221
222 let models: Vec<ApiModelEntry> = response
223 .json()
224 .await
225 .map_err(|e| FetchError::Http(e.to_string()))?;
226
227 if models.is_empty() {
228 break;
229 }
230
231 for model in &models {
232 // Client-side tag filter: the HF API may ignore the `filter` query
233 // parameter when combined with other params, so verify the tag is
234 // actually present on each returned model.
235 if let Some(t) = tag
236 && !model.tags.iter().any(|model_tag| {
237 // BORROW: explicit .as_str() instead of Deref coercion
238 model_tag.as_str().eq_ignore_ascii_case(t)
239 })
240 {
241 continue;
242 }
243
244 // BORROW: explicit .as_ref() and .as_deref() for Option<String>
245 let model_type = model.config.as_ref().and_then(|c| c.model_type.as_deref());
246
247 if let Some(mt) = model_type {
248 remote_families
249 .entry(mt.to_owned())
250 .or_insert_with(|| (model.model_id.clone(), model.downloads));
251 }
252 }
253
254 offset = offset.saturating_add(models.len());
255 }
256
257 // Filter to families not already cached locally.
258 // BORROW: explicit .as_str() instead of Deref coercion
259 let discovered: Vec<DiscoveredFamily> = remote_families
260 .into_iter()
261 .filter(|(mt, _)| !local_families.contains(mt.as_str()))
262 .map(|(model_type, (top_model, downloads))| DiscoveredFamily {
263 model_type,
264 top_model,
265 downloads,
266 })
267 .collect();
268
269 Ok(discovered)
270}
271
272/// Normalizes common quantization synonyms in a search query so that
273/// variant spellings (e.g., `"8bit"`, `"8-bit"`, `"int8"`) produce
274/// consistent results.
275#[must_use]
276fn normalize_quantization_terms(query: &str) -> String {
277 /// Synonym groups: all variants map to the first (canonical) form.
278 const SYNONYMS: &[(&[&str], &str)] = &[
279 (&["8bit", "8-bit", "int8"], "8-bit"),
280 (&["4bit", "4-bit", "int4"], "4-bit"),
281 (&["fp8", "float8"], "fp8"),
282 ];
283
284 query
285 .split_whitespace()
286 .map(|token| {
287 // BORROW: explicit .to_lowercase() for case-insensitive comparison
288 let lower = token.to_lowercase();
289 for &(variants, canonical) in SYNONYMS {
290 // BORROW: explicit .as_str() instead of Deref coercion
291 if variants.contains(&lower.as_str()) {
292 // BORROW: explicit .to_owned() for &str → owned String
293 return (*canonical).to_owned();
294 }
295 }
296 // BORROW: explicit .to_owned() for &str → owned String
297 token.to_owned()
298 })
299 .collect::<Vec<_>>()
300 .join(" ")
301}
302
303/// Searches the `HuggingFace` Hub for models matching a query string.
304///
305/// Optionally filters by `library` framework (e.g., `"transformers"`, `"peft"`),
306/// `pipeline` task tag (e.g., `"text-generation"`), and/or `tag` (e.g., `"gguf"`).
307/// Library and pipeline filters are sent as query parameters; tag is sent via the
308/// `filter` parameter. All three are also applied client-side for correctness.
309///
310/// Common quantization synonyms (`"8bit"` / `"8-bit"` / `"int8"`,
311/// `"4bit"` / `"4-bit"` / `"int4"`, `"fp8"` / `"float8"`) are normalized
312/// before querying the API so that variant spellings return consistent results.
313///
314/// Results are sorted by download count (most popular first).
315///
316/// # Arguments
317///
318/// * `query` — Free-text search string (e.g., `"RWKV-7"`, `"llama 3"`).
319/// * `limit` — Maximum number of results to return.
320/// * `library` — Optional library filter (e.g., `"peft"`, `"transformers"`).
321/// * `pipeline` — Optional pipeline tag filter (e.g., `"text-generation"`).
322/// * `tag` — Optional tag filter (e.g., `"gguf"`, `"conversational"`).
323/// * `token` — Authentication token (or `None` for anonymous requests).
324/// Gated repos are typically still visible without one (gating restricts
325/// content downloads, not search visibility or metadata), but a private
326/// repo the token has access to is only ever found with it.
327///
328/// # Errors
329///
330/// Returns [`FetchError::Http`] if the API request fails.
331pub async fn search_models(
332 query: &str,
333 limit: usize,
334 library: Option<&str>,
335 pipeline: Option<&str>,
336 tag: Option<&str>,
337 token: Option<&str>,
338) -> Result<Vec<SearchResult>, FetchError> {
339 let normalized = normalize_quantization_terms(query);
340 let client = crate::chunked::build_client(token)?;
341
342 // BORROW: explicit .as_str() instead of Deref coercion
343 let mut query_params: Vec<(&str, &str)> = vec![
344 ("search", normalized.as_str()),
345 ("sort", "downloads"),
346 ("direction", "-1"),
347 ];
348 if let Some(lib) = library {
349 query_params.push(("library", lib));
350 }
351 if let Some(pipe) = pipeline {
352 query_params.push(("pipeline_tag", pipe));
353 }
354 if let Some(t) = tag {
355 query_params.push(("filter", t));
356 }
357
358 let response = client
359 .get(HF_API_BASE)
360 .query(&query_params)
361 .query(&[("limit", limit)])
362 .send()
363 .await
364 .map_err(|e| FetchError::Http(e.to_string()))?;
365
366 if !response.status().is_success() {
367 return Err(FetchError::Http(format!(
368 "HF API returned status {}",
369 response.status()
370 )));
371 }
372
373 let models: Vec<ApiModelEntry> = response
374 .json()
375 .await
376 .map_err(|e| FetchError::Http(e.to_string()))?;
377
378 // Client-side filtering: the HF search API may ignore library/pipeline_tag/filter
379 // query parameters when combined with the `search` parameter, so we filter
380 // the results ourselves to guarantee correctness.
381 let results = models
382 .into_iter()
383 .filter(|m| {
384 if let Some(lib) = library {
385 match m.library_name {
386 // BORROW: explicit .as_str() instead of Deref coercion
387 Some(ref name) if name.as_str().eq_ignore_ascii_case(lib) => {}
388 _ => return false,
389 }
390 }
391 if let Some(pipe) = pipeline {
392 match m.pipeline_tag {
393 // BORROW: explicit .as_str() instead of Deref coercion
394 Some(ref t) if t.as_str().eq_ignore_ascii_case(pipe) => {}
395 _ => return false,
396 }
397 }
398 if let Some(t) = tag
399 && !m.tags.iter().any(|model_tag| {
400 // BORROW: explicit .as_str() instead of Deref coercion
401 model_tag.as_str().eq_ignore_ascii_case(t)
402 })
403 {
404 return false;
405 }
406 true
407 })
408 .map(|m| SearchResult {
409 model_id: m.model_id,
410 downloads: m.downloads,
411 library_name: m.library_name,
412 pipeline_tag: m.pipeline_tag,
413 tags: m.tags,
414 })
415 .collect();
416
417 Ok(results)
418}
419
420/// How a repo's (or a filtered listing's) `.gguf` files relate to each other.
421#[derive(Debug, Clone, PartialEq, Eq)]
422#[non_exhaustive]
423pub enum GgufFileSetKind {
424 /// A single logical file split via `llama.cpp`'s official
425 /// `*-NNNNN-of-MMMMM.gguf` convention — summing sizes across the set is
426 /// correct (you need every shard). `first_shard` is the filename of the
427 /// split whose index is `1` — under `llama.cpp`'s own convention, the
428 /// *only* shard carrying the `GGUF` metadata `KV` table (and therefore
429 /// any backlink key); later shards' headers hold tensor info alone.
430 /// Carried here, rather than re-derived by `pick_backlink_representative`
431 /// via its own separate scan, so classification and "which file is
432 /// shard 1" can never drift apart.
433 Sharded {
434 /// The filename of the shard whose parsed index is `1`.
435 first_shard: String,
436 },
437 /// Two or more mutually-exclusive quantization alternatives — summing
438 /// sizes across the set is a category error, since nobody downloads more
439 /// than one.
440 QuantAlternatives,
441 /// Zero or one `.gguf` file present — nothing to disambiguate; today's
442 /// summing behavior is already correct.
443 NotApplicable,
444}
445
446/// Returns `true` when `filename` has a `.gguf` extension, case-insensitively.
447#[must_use]
448pub fn is_gguf_filename(filename: &str) -> bool {
449 filename.to_ascii_lowercase().ends_with(".gguf")
450}
451
452/// Parses `filename` against `llama.cpp`'s split-file convention
453/// (`<prefix>-<index>-of-<total>.gguf`, `index`/`total` equal-width,
454/// zero-padded digit strings). Returns `(prefix, index, total)` on a match.
455fn parse_gguf_split_name(filename: &str) -> Option<(&str, u32, u32)> {
456 if !is_gguf_filename(filename) {
457 return None;
458 }
459 // `.gguf` is 5 ASCII bytes, so this length is always a valid boundary.
460 let stem = filename.get(..filename.len().checked_sub(5)?)?;
461 let (before, total_str) = stem.split_once("-of-")?;
462 let (prefix, index_str) = before.rsplit_once('-')?;
463 if index_str.is_empty()
464 || total_str.is_empty()
465 || index_str.len() != total_str.len()
466 || !index_str.bytes().all(|b| b.is_ascii_digit())
467 || !total_str.bytes().all(|b| b.is_ascii_digit())
468 {
469 return None;
470 }
471 let index: u32 = index_str.parse().ok()?;
472 let total: u32 = total_str.parse().ok()?;
473 Some((prefix, index, total))
474}
475
476/// Classifies a repo's `.gguf` files as a sharded single file or a set of
477/// mutually-exclusive quant alternatives.
478///
479/// `Sharded` requires **every** `.gguf` file in `filenames` to match
480/// `parse_gguf_split_name` with the same prefix and total, the file count
481/// to equal that total, and the indices to cover `1..=total` exactly.
482/// Anything else with two or more `.gguf` files — a mismatched prefix, a
483/// file outside the split convention entirely, a missing or duplicate index
484/// — is `QuantAlternatives`.
485#[must_use]
486pub fn classify_gguf_files(filenames: &[&str]) -> GgufFileSetKind {
487 let gguf: Vec<&str> = filenames
488 .iter()
489 .copied()
490 .filter(|f| is_gguf_filename(f))
491 .collect();
492 if gguf.len() <= 1 {
493 return GgufFileSetKind::NotApplicable;
494 }
495
496 let Some(parsed): Option<Vec<(&str, u32, u32)>> =
497 gguf.iter().map(|f| parse_gguf_split_name(f)).collect()
498 else {
499 return GgufFileSetKind::QuantAlternatives;
500 };
501 // INDEX: gguf.len() > 1 checked above, so parsed (same length) is non-empty
502 let Some(&(first_prefix, _, first_total)) = parsed.first() else {
503 return GgufFileSetKind::QuantAlternatives;
504 };
505 let same_group = parsed
506 .iter()
507 .all(|&(prefix, _, total)| prefix == first_prefix && total == first_total);
508 if !same_group {
509 return GgufFileSetKind::QuantAlternatives;
510 }
511 // CAST: u32 -> usize, total is a small shard count from a zero-padded filename token
512 #[allow(clippy::as_conversions)]
513 let total_usize = first_total as usize;
514 if gguf.len() != total_usize {
515 return GgufFileSetKind::QuantAlternatives;
516 }
517
518 let mut indices: Vec<u32> = parsed.iter().map(|&(_, index, _)| index).collect();
519 indices.sort_unstable();
520 let complete = indices.iter().enumerate().all(|(i, &index)| {
521 // CAST: usize -> u32, i is bounded by total_usize which came from a u32
522 #[allow(clippy::as_conversions, clippy::cast_possible_truncation)]
523 let expected = i as u32 + 1;
524 index == expected
525 });
526
527 if !complete {
528 return GgufFileSetKind::QuantAlternatives;
529 }
530
531 // INDEX: `complete` just verified the indices cover `1..=total_usize`
532 // exactly, so exactly one entry in `parsed` (same order/length as
533 // `gguf`, both filtered from `filenames` together) has index `1`.
534 let Some((&first_shard, _)) = gguf
535 .iter()
536 .zip(parsed.iter())
537 .find(|&(_, &(_, index, _))| index == 1)
538 else {
539 // EXPLICIT: unreachable given `complete` above; degrade to
540 // QuantAlternatives instead of panicking if this invariant is ever
541 // violated by a future edit.
542 return GgufFileSetKind::QuantAlternatives;
543 };
544 GgufFileSetKind::Sharded {
545 // BORROW: explicit .to_owned() for &str → owned String field
546 first_shard: first_shard.to_owned(),
547 }
548}
549
550/// Computes the min/max size across `.gguf` files in `filenames_with_size`,
551/// when they are mutually-exclusive quant alternatives (see
552/// [`classify_gguf_files`]). Returns `None` when the set is `Sharded` or
553/// `NotApplicable` (summing is already correct there) or when none of the
554/// `.gguf` files carry a known size.
555///
556/// Takes `(filename, size)` pairs rather than a concrete file type so both
557/// the remote listing (`repo::RepoFile`, whose `size` is `Option<u64>`) and
558/// the local cache listing (`cache::CacheFileUsage`, whose `size` is always
559/// known) can share one implementation. `size` is `Option<u64>` — and every
560/// filename must be passed, not just the ones with a known size —
561/// because classification (`Sharded` vs `QuantAlternatives`) depends on
562/// seeing the *complete* file set: a genuinely sharded repo where the Hub
563/// API happens not to report one shard's size would otherwise lose that
564/// shard's filename before `classify_gguf_files` ever saw it, undercounting
565/// the shard total and misclassifying the whole set as `QuantAlternatives`.
566/// Only the size-known subset feeds the min/max once classification itself
567/// has already run on every filename.
568#[must_use]
569pub fn gguf_size_range<'a, I>(filenames_with_size: I) -> Option<(u64, u64)>
570where
571 I: IntoIterator<Item = (&'a str, Option<u64>)>,
572{
573 let pairs: Vec<(&str, Option<u64>)> = filenames_with_size.into_iter().collect();
574 let filenames: Vec<&str> = pairs.iter().map(|&(name, _)| name).collect();
575 if !matches!(
576 classify_gguf_files(&filenames),
577 GgufFileSetKind::QuantAlternatives
578 ) {
579 return None;
580 }
581
582 let sizes: Vec<u64> = pairs
583 .iter()
584 .filter(|&&(name, _)| is_gguf_filename(name))
585 .filter_map(|&(_, size)| size)
586 .collect();
587 let min = sizes.iter().copied().min()?;
588 let max = sizes.iter().copied().max()?;
589 Some((min, max))
590}
591
592/// A repo's aggregate size, aware of the `.gguf` quant-alternatives case.
593///
594/// `total` is always the raw sum of every listed file's size — well-defined,
595/// if not always the most useful number. When `quant_alternatives` is `true`,
596/// `size_min`/`size_max` give the more honest range: the smallest and
597/// largest `.gguf` file, since the repo's files are mutually-exclusive
598/// choices rather than parts of a whole.
599#[derive(Debug, Clone, Copy)]
600#[non_exhaustive]
601pub struct RepoSizeSummary {
602 /// Sum of every listed file's size.
603 pub total: u64,
604 /// Whether this repo's `.gguf` files are quant alternatives (see [`GgufFileSetKind`]).
605 pub quant_alternatives: bool,
606 /// Smallest `.gguf` file's size, when `quant_alternatives` is `true`.
607 pub size_min: Option<u64>,
608 /// Largest `.gguf` file's size, when `quant_alternatives` is `true`.
609 pub size_max: Option<u64>,
610}
611
612/// Fetches a repo's file listing and summarizes its size, detecting the
613/// `.gguf` quant-alternatives case via [`classify_gguf_files`].
614///
615/// Additive alongside [`fetch_repo_total_size`] (kept unchanged for existing
616/// callers/downstream consumers) rather than replacing it — this is the
617/// quant-aware variant `hf-fm search --show size` uses.
618///
619/// # Errors
620///
621/// Returns [`FetchError::Http`] if the API request fails or returns a
622/// non-success status. Returns [`FetchError::RepoNotFound`] if the
623/// repository does not exist on the Hub.
624pub async fn fetch_repo_size_summary(
625 repo_id: &str,
626 client: &reqwest::Client,
627) -> Result<RepoSizeSummary, FetchError> {
628 let files = crate::repo::list_repo_files_with_metadata(repo_id, None, None, client).await?;
629 let total: u64 = files.iter().filter_map(|f| f.size).sum();
630
631 // BORROW: explicit .as_str() instead of Deref coercion
632 let sized: Vec<(&str, Option<u64>)> = files
633 .iter()
634 .map(|f| (f.filename.as_str(), f.size))
635 .collect();
636 let range = gguf_size_range(sized);
637
638 Ok(RepoSizeSummary {
639 total,
640 quant_alternatives: range.is_some(),
641 size_min: range.map(|(min, _)| min),
642 size_max: range.map(|(_, max)| max),
643 })
644}
645
646/// Returns the total size in bytes of all files in the given repository's
647/// `main` revision, summed across `siblings[].size` from the
648/// `/api/models/{repo_id}?blobs=true` endpoint.
649///
650/// Used by `hf-fm search --show size` to enrich result rows with a total-repo
651/// size column. Failures are surfaced to the caller so individual repo
652/// lookups can be skipped (the search itself is not aborted on a single
653/// 404 / network blip).
654///
655/// # Arguments
656///
657/// * `repo_id` — The full model identifier (e.g., `"org/model"`).
658/// * `client` — Shared `reqwest::Client` (callers fan out N lookups concurrently).
659///
660/// # Errors
661///
662/// Returns [`FetchError::Http`] if the API request
663/// fails or returns a non-success status.
664/// Returns [`FetchError::RepoNotFound`] if
665/// the repository does not exist on the Hub.
666pub async fn fetch_repo_total_size(
667 repo_id: &str,
668 client: &reqwest::Client,
669) -> Result<u64, FetchError> {
670 let files = crate::repo::list_repo_files_with_metadata(repo_id, None, None, client).await?;
671 Ok(files.iter().filter_map(|f| f.size).sum())
672}
673
674/// Backlink verification outcome for a [`QuantCandidate`].
675#[derive(Debug, Clone)]
676#[non_exhaustive]
677pub enum QuantVerification {
678 /// A `.gguf` file's `general.source.url` / `general.base_model.*.repo_url`
679 /// metadata confirms this repo targets the base model.
680 Verified,
681 /// No verification could run: the candidate has no `.gguf` file, or its
682 /// GGUF metadata carries no recognized backlink key. The naming match
683 /// stands alone.
684 Unverified,
685 /// The backlink check itself failed (network error, timeout, a gated
686 /// repo) rather than returning a definite answer. The naming match
687 /// stands alone; the reason is kept for display.
688 CheckFailed(String),
689}
690
691/// A quant-sibling repo candidate discovered for a base model, plus its file
692/// listing (reused by callers building a size table, so [`discover_quant_siblings`]
693/// is the only Hub round-trip needed per candidate beyond the initial search).
694#[derive(Debug, Clone)]
695#[non_exhaustive]
696pub struct QuantCandidate {
697 /// The candidate repository identifier.
698 pub repo_id: String,
699 /// Backlink verification outcome (see [`QuantVerification`]).
700 pub verification: QuantVerification,
701 /// The candidate repo's file listing (sizes and SHA256 where known).
702 pub files: Vec<crate::repo::RepoFile>,
703}
704
705impl QuantCandidate {
706 /// Builds a [`QuantCandidate`] directly.
707 ///
708 /// `#[non_exhaustive]` blocks struct-literal construction from outside
709 /// this crate, so this is the canonical way to build one from the `hf-fm`
710 /// binary crate — used by its test suite; production code only ever
711 /// receives `QuantCandidate` values from [`discover_quant_siblings`].
712 #[must_use]
713 pub fn new(
714 repo_id: String,
715 verification: QuantVerification,
716 files: Vec<crate::repo::RepoFile>,
717 ) -> Self {
718 Self {
719 repo_id,
720 verification,
721 files,
722 }
723 }
724}
725
726/// Known `GGUF` metadata keys that point back at a source `HuggingFace` repo,
727/// per `llama.cpp`'s `gguf-py` metadata writer conventions.
728fn gguf_source_backlinks(metadata: &HashMap<String, String>) -> Vec<&str> {
729 metadata
730 .iter()
731 .filter(|(key, _)| {
732 key.as_str() == "general.source.url"
733 || key.as_str() == "general.source.huggingface.repository"
734 || (key.starts_with("general.base_model.") && key.ends_with(".repo_url"))
735 })
736 // BORROW: explicit .as_str() instead of Deref coercion
737 .map(|(_, value)| value.as_str())
738 .collect()
739}
740
741/// Picks the `.gguf` file whose header `build_quant_candidate` should
742/// inspect for a `general.source.url`-style backlink.
743///
744/// A genuinely sharded `GGUF` file (see [`classify_gguf_files`]) carries its
745/// metadata `KV` table in the *first* split only — `llama.cpp`'s own
746/// convention; later shards' headers hold tensor info alone, no metadata.
747/// Picking the smallest shard by size (very often the *last* one, not the
748/// first) would inspect a header with no backlink key at all and misreport
749/// [`QuantVerification::Unverified`] even when the repo genuinely names its
750/// source. For a mutually-exclusive quant-alternatives set (or a lone
751/// `.gguf` file), smallest-by-size stays the right choice: it is the
752/// cheapest header to fetch, and every alternative carries the same
753/// backlink (they are re-quantizations of the same source checkpoint).
754///
755/// Reads which file is shard 1 straight off [`GgufFileSetKind::Sharded`]'s
756/// own `first_shard` field rather than re-deriving it here, so this can
757/// never disagree with `classify_gguf_files`'s own notion of "sharded".
758fn pick_backlink_representative(
759 gguf_files: &[crate::repo::RepoFile],
760) -> Option<&crate::repo::RepoFile> {
761 // BORROW: explicit .as_str() instead of Deref coercion
762 let filenames: Vec<&str> = gguf_files.iter().map(|f| f.filename.as_str()).collect();
763 if let GgufFileSetKind::Sharded { first_shard } = classify_gguf_files(&filenames)
764 && let Some(file) = gguf_files.iter().find(|f| f.filename == first_shard)
765 {
766 return Some(file);
767 }
768 gguf_files.iter().min_by_key(|f| f.size.unwrap_or(u64::MAX))
769}
770
771/// Lists `candidate_repo_id`'s files and, if it holds at least one `.gguf`
772/// file, cross-checks a representative one's metadata backlink (see
773/// [`pick_backlink_representative`]) against `base_repo_id`.
774///
775/// Returns `None` only when a backlink is present and explicitly names a
776/// *different* repo — a naming collision, not a real sibling. Every other
777/// outcome (no `.gguf` file, no backlink key, or the header fetch itself
778/// failing) is returned as a kept candidate with the reason recorded in
779/// [`QuantVerification`], per [`discover_quant_siblings`]'s contract that a
780/// transient failure must not hide a real candidate. A file-listing failure
781/// (the repo cannot be enumerated at all) drops the candidate silently —
782/// there would be nothing to render for it regardless of verification status.
783async fn build_quant_candidate(
784 candidate_repo_id: String,
785 base_repo_id: &str,
786 token: Option<&str>,
787 client: &reqwest::Client,
788) -> Option<QuantCandidate> {
789 let files = crate::repo::list_repo_files_with_metadata(&candidate_repo_id, token, None, client)
790 .await
791 .ok()?;
792
793 let gguf_files: Vec<_> = files
794 .iter()
795 .filter(|f| is_gguf_filename(&f.filename))
796 .cloned()
797 .collect();
798
799 let Some(representative) = pick_backlink_representative(&gguf_files) else {
800 return Some(QuantCandidate {
801 repo_id: candidate_repo_id,
802 verification: QuantVerification::Unverified,
803 files,
804 });
805 };
806
807 match crate::inspect::inspect_gguf(&candidate_repo_id, &representative.filename, token, None)
808 .await
809 {
810 Err(e) => Some(QuantCandidate {
811 repo_id: candidate_repo_id,
812 verification: QuantVerification::CheckFailed(e.to_string()),
813 files,
814 }),
815 Ok((info, _source, _stats)) => {
816 let backlinks = info
817 .metadata
818 .as_ref()
819 .map(|m| gguf_source_backlinks(m))
820 .unwrap_or_default();
821 if backlinks.is_empty() {
822 return Some(QuantCandidate {
823 repo_id: candidate_repo_id,
824 verification: QuantVerification::Unverified,
825 files,
826 });
827 }
828 // BORROW: explicit .to_lowercase() for case-insensitive substring match
829 let base_lower = base_repo_id.to_lowercase();
830 let matched = backlinks
831 .iter()
832 .any(|b| b.to_lowercase().contains(&base_lower));
833 if matched {
834 Some(QuantCandidate {
835 repo_id: candidate_repo_id,
836 verification: QuantVerification::Verified,
837 files,
838 })
839 } else {
840 None // EXPLICIT: backlink present but points elsewhere — not a real sibling
841 }
842 }
843 }
844}
845
846/// Discovers quant-sibling repos for a base model.
847///
848/// Searches the Hub for the base model's short name (the part of
849/// `base_repo_id` after `/`), keeps results whose repo ID contains that name
850/// as a case-insensitive substring (excluding `base_repo_id` itself), then
851/// fans out `build_quant_candidate` across the survivors through a bounded
852/// `tokio::sync::Semaphore` (8 permits, mirroring [`fetch_repo_sizes_concurrent`]).
853///
854/// Sibling discovery has no dedicated Hub endpoint — quant repos
855/// overwhelmingly either name themselves `<base>-<SCHEME>` or carry a GGUF
856/// metadata backlink to the original, so this combines both signals: the
857/// naming match decides the candidate *pool*, and the backlink (when present
858/// and checkable) raises confidence without ever silently dropping a
859/// candidate over a transient network failure.
860///
861/// # Errors
862///
863/// Returns [`FetchError::Http`] if the initial Hub search fails. Per-candidate
864/// failures below that point are absorbed into [`QuantVerification::CheckFailed`]
865/// rather than aborting the whole discovery.
866pub async fn discover_quant_siblings(
867 base_repo_id: &str,
868 token: Option<&str>,
869 client: &reqwest::Client,
870) -> Result<Vec<QuantCandidate>, FetchError> {
871 let short_name = base_repo_id.rsplit('/').next().unwrap_or(base_repo_id);
872 let results = search_models(short_name, 50, None, None, None, token).await?;
873
874 // BORROW: explicit .to_lowercase() for case-insensitive substring match
875 let short_name_lower = short_name.to_lowercase();
876 let semaphore = Arc::new(tokio::sync::Semaphore::new(8));
877 let mut set: tokio::task::JoinSet<Option<QuantCandidate>> = tokio::task::JoinSet::new();
878
879 for result in results {
880 if result.model_id.eq_ignore_ascii_case(base_repo_id) {
881 continue; // EXPLICIT: the base repo itself, not a sibling
882 }
883 if !result.model_id.to_lowercase().contains(&short_name_lower) {
884 continue;
885 }
886
887 let limiter = Arc::clone(&semaphore);
888 let client = client.clone();
889 // BORROW: explicit .to_owned() — the spawned task must be 'static
890 let base_repo_owned = base_repo_id.to_owned();
891 let token_owned = token.map(str::to_owned);
892 set.spawn(async move {
893 let _permit = limiter.acquire_owned().await.ok()?;
894 // BORROW: explicit .as_deref() for Option<String> → Option<&str>
895 build_quant_candidate(
896 result.model_id,
897 &base_repo_owned,
898 token_owned.as_deref(),
899 &client,
900 )
901 .await
902 });
903 }
904
905 let mut candidates: Vec<QuantCandidate> = Vec::new();
906 while let Some(joined) = set.join_next().await {
907 if let Ok(Some(candidate)) = joined {
908 candidates.push(candidate);
909 }
910 }
911 candidates.sort_by(|a, b| a.repo_id.cmp(&b.repo_id));
912
913 Ok(candidates)
914}
915
916/// Fans out an async per-item operation through a bounded
917/// `tokio::sync::Semaphore`, running at most `concurrency` futures at once.
918/// Returns one `Option<R>` per input item, in the same order as `items` —
919/// regardless of which task finishes first. `None` marks a per-item outcome
920/// the caller should treat as absent: `f` returning `None`, or (defensively)
921/// a task whose permit was never acquired or that panicked. Neither case
922/// aborts the fan-out; every other item still runs to completion.
923///
924/// The single bounded-concurrency scaffold behind every per-repo network
925/// fan-out in this crate — [`fetch_repo_sizes_concurrent`],
926/// [`fetch_repo_size_summaries_concurrent`], [`fetch_tags_concurrent`], and
927/// (through the `hf-fm` CLI, a separate crate that can only reach a `pub`
928/// item) `quants --fits`'s bounded offload-plan inspection — so a future
929/// change to permit-acquisition, panic-handling, or concurrency semantics
930/// only needs to happen once. Order-preserving (rather than the map-keyed
931/// shape an earlier version of this helper had) because `--fits` needs
932/// every row's verdict, in row order, even the ones that produced no
933/// result — a `HashMap` has no way to represent "present but absent",
934/// only "absent". Callers that want a `repo_id`-keyed map instead (every
935/// caller above except `--fits`) zip `items` back onto the result, which is
936/// why `T` does not need `Eq + Hash` here even though every current caller's
937/// `T` happens to satisfy it.
938///
939/// `f` is `Fn`, not `FnOnce`, since it is called once per item — shared
940/// state it needs (e.g. an HTTP client) should be cloned inside the closure
941/// body on each call, not moved out of the closure's own captured environment.
942pub async fn fan_out_bounded<T, R, F, Fut>(
943 items: Vec<T>,
944 concurrency: usize,
945 f: F,
946) -> Vec<Option<R>>
947where
948 T: Send + 'static,
949 R: Send + 'static,
950 F: Fn(T) -> Fut + Send + Sync + 'static,
951 Fut: Future<Output = Option<R>> + Send + 'static,
952{
953 let len = items.len();
954 let semaphore = Arc::new(tokio::sync::Semaphore::new(concurrency));
955 let f = Arc::new(f);
956 let mut set: tokio::task::JoinSet<(usize, Option<R>)> = tokio::task::JoinSet::new();
957
958 for (index, item) in items.into_iter().enumerate() {
959 let limiter = Arc::clone(&semaphore);
960 let f = Arc::clone(&f);
961 set.spawn(async move {
962 let Ok(_permit) = limiter.acquire_owned().await else {
963 return (index, None);
964 };
965 (index, f(item).await)
966 });
967 }
968
969 let mut slots: Vec<Option<R>> = (0..len).map(|_| None).collect();
970 while let Some(joined) = set.join_next().await {
971 if let Ok((index, result)) = joined
972 && let Some(slot) = slots.get_mut(index)
973 {
974 *slot = result;
975 }
976 }
977 slots
978}
979
980/// Zips `keys` back onto `fan_out_bounded`'s order-aligned `Vec<Option<R>>`,
981/// dropping absent slots — the map-building step every `repo_id`-keyed
982/// fan-out caller in this module needs after switching to the order-
983/// preserving [`fan_out_bounded`].
984fn zip_into_map<T: Eq + std::hash::Hash, R>(
985 keys: Vec<T>,
986 results: Vec<Option<R>>,
987) -> HashMap<T, R> {
988 keys.into_iter()
989 .zip(results)
990 .filter_map(|(key, result)| result.map(|r| (key, r)))
991 .collect()
992}
993
994/// Fans out [`fetch_repo_total_size`] across the given repository IDs through
995/// a bounded `tokio::sync::Semaphore` (8 permits) to stay friendly to the HF
996/// Hub on `--limit 100`-style invocations.
997///
998/// Per-repo failures (network errors, 404s, missing `size` fields) are
999/// silently dropped from the returned map; callers should render rows whose
1000/// `repo_id` is absent from the map with a placeholder (`—`). The search
1001/// itself is not aborted on a single failure.
1002///
1003/// # Arguments
1004///
1005/// * `repo_ids` — Owned list of model identifiers. Ownership is moved into
1006/// the spawned tasks so each future is `'static`.
1007#[must_use]
1008pub async fn fetch_repo_sizes_concurrent(repo_ids: Vec<String>) -> HashMap<String, u64> {
1009 let client = reqwest::Client::new();
1010 // BORROW: explicit .clone() — fan_out_bounded consumes `repo_ids`, but
1011 // the original ids are still needed afterward to key the returned map.
1012 let keys = repo_ids.clone();
1013 let results = fan_out_bounded(repo_ids, 8, move |repo_id| {
1014 let client = client.clone();
1015 // EXPLICIT: per-repo failure intentionally swallowed — the caller
1016 // renders the row with "—" rather than aborting the search.
1017 async move { fetch_repo_total_size(&repo_id, &client).await.ok() }
1018 })
1019 .await;
1020 zip_into_map(keys, results)
1021}
1022
1023/// Fans out [`fetch_repo_size_summary`] across the given repository IDs
1024/// through a bounded `tokio::sync::Semaphore` (8 permits) — the quant-aware
1025/// counterpart `hf-fm search --show size` uses.
1026///
1027/// Per-repo failures are silently dropped from the returned map; callers
1028/// should render rows whose `repo_id` is absent with a placeholder (`—`).
1029/// This is deliberately distinct from a client-build failure (see Errors
1030/// below): one bad repo is an expected, per-item outcome across a fan-out
1031/// this wide, but a token too malformed to become an HTTP header is a
1032/// caller configuration error every row would otherwise fail identically
1033/// and silently for — worth surfacing loudly instead, the same way
1034/// [`search_models`] already does for the same failure mode.
1035///
1036/// # Arguments
1037///
1038/// * `repo_ids` — Owned list of model identifiers. Ownership is moved into
1039/// the spawned tasks so each future is `'static`.
1040/// * `token` — Authentication token (or `None` for anonymous requests),
1041/// applied to every fanned-out request via [`crate::chunked::build_client`]
1042/// — without this, a private or gated repo silently fails to size (the
1043/// same per-repo failure path as a 404, indistinguishable to the caller).
1044///
1045/// # Errors
1046///
1047/// Returns [`FetchError::Http`] if `token` cannot be turned into a valid
1048/// HTTP header value (via [`crate::chunked::build_client`]) — before any
1049/// per-repo request is attempted.
1050pub async fn fetch_repo_size_summaries_concurrent(
1051 repo_ids: Vec<String>,
1052 token: Option<&str>,
1053) -> Result<HashMap<String, RepoSizeSummary>, FetchError> {
1054 let client = crate::chunked::build_client(token)?;
1055 // BORROW: explicit .clone() — fan_out_bounded consumes `repo_ids`, but
1056 // the original ids are still needed afterward to key the returned map.
1057 let keys = repo_ids.clone();
1058 let results = fan_out_bounded(repo_ids, 8, move |repo_id| {
1059 let client = client.clone();
1060 // EXPLICIT: per-repo failure intentionally swallowed — the caller
1061 // renders the row with "—" rather than aborting the search.
1062 async move { fetch_repo_size_summary(&repo_id, &client).await.ok() }
1063 })
1064 .await;
1065 Ok(zip_into_map(keys, results))
1066}
1067
1068/// Fans out [`fetch_model_card`] across the given repository IDs through a
1069/// bounded `tokio::sync::Semaphore` (8 permits) and returns a map from
1070/// `repo_id` to the model card's tag list.
1071///
1072/// Per-repo failures (network errors, 404s, missing models) are silently
1073/// dropped from the returned map. Callers that want strict semantics should
1074/// treat absence as "no tags known". Uses the same bounded fan-out
1075/// scaffold as [`fetch_repo_sizes_concurrent`], [`fan_out_bounded`].
1076///
1077/// # Arguments
1078///
1079/// * `repo_ids` — Owned list of model identifiers. Ownership is moved into
1080/// the spawned tasks so each future is `'static`.
1081#[must_use]
1082pub async fn fetch_tags_concurrent(repo_ids: Vec<String>) -> HashMap<String, Vec<String>> {
1083 // BORROW: explicit .clone() — fan_out_bounded consumes `repo_ids`, but
1084 // the original ids are still needed afterward to key the returned map.
1085 let keys = repo_ids.clone();
1086 let results = fan_out_bounded(repo_ids, 8, move |repo_id| async move {
1087 // EXPLICIT: per-repo failure intentionally swallowed — missing
1088 // tags mean the row simply doesn't match any --tag filter (the
1089 // user's listing is not aborted on a single 404 / network blip).
1090 fetch_model_card(&repo_id).await.ok().map(|card| card.tags)
1091 })
1092 .await;
1093 zip_into_map(keys, results)
1094}
1095
1096/// Fetches model card metadata for a specific model from the `HuggingFace` Hub.
1097///
1098/// Queries `GET https://huggingface.co/api/models/{model_id}` and extracts
1099/// license, pipeline tag, tags, library name, and languages from the response.
1100///
1101/// # Arguments
1102///
1103/// * `model_id` — The full model identifier (e.g., `"mistralai/Ministral-3-3B-Instruct-2512"`).
1104///
1105/// # Errors
1106///
1107/// Returns [`FetchError::Http`] if the API request fails or the model is not found.
1108pub async fn fetch_model_card(model_id: &str) -> Result<ModelCardMetadata, FetchError> {
1109 let client = reqwest::Client::new();
1110 let url = format!("{HF_API_BASE}/{model_id}");
1111
1112 let response = client
1113 .get(url.as_str()) // BORROW: explicit .as_str()
1114 .send()
1115 .await
1116 .map_err(|e| FetchError::Http(e.to_string()))?;
1117
1118 if !response.status().is_success() {
1119 return Err(FetchError::Http(format!(
1120 "HF API returned status {} for model {model_id}",
1121 response.status()
1122 )));
1123 }
1124
1125 let detail: ApiModelDetail = response
1126 .json()
1127 .await
1128 .map_err(|e| FetchError::Http(e.to_string()))?;
1129
1130 let (license, languages) = if let Some(card) = detail.card_data {
1131 let langs = match card.language {
1132 Some(ApiLanguage::Single(s)) => vec![s],
1133 Some(ApiLanguage::Multiple(v)) => v,
1134 None => Vec::new(),
1135 };
1136 (card.license, langs)
1137 } else {
1138 (None, Vec::new())
1139 };
1140
1141 let gated = match detail.gated {
1142 ApiGated::Bool(false) => GateStatus::Open,
1143 ApiGated::Mode(ref mode) if mode.eq_ignore_ascii_case("manual") => GateStatus::Manual,
1144 ApiGated::Bool(true) | ApiGated::Mode(_) => GateStatus::Auto,
1145 };
1146
1147 Ok(ModelCardMetadata {
1148 license,
1149 pipeline_tag: detail.pipeline_tag,
1150 tags: detail.tags,
1151 library_name: detail.library_name,
1152 languages,
1153 gated,
1154 })
1155}
1156
1157/// Fetches the raw README text for a `HuggingFace` model repository.
1158///
1159/// Downloads `README.md` from the repository at the given revision.
1160/// Returns `Ok(None)` if the file does not exist (HTTP 404).
1161///
1162/// # Arguments
1163///
1164/// * `model_id` — The full model identifier (e.g., `"mistralai/Ministral-3-3B-Instruct-2512"`).
1165/// * `revision` — Git revision to fetch (defaults to `"main"` when `None`).
1166/// * `token` — Optional authentication token.
1167///
1168/// # Errors
1169///
1170/// Returns [`FetchError::Http`] if the request fails (other than 404).
1171pub async fn fetch_readme(
1172 model_id: &str,
1173 revision: Option<&str>,
1174 token: Option<&str>,
1175) -> Result<Option<String>, FetchError> {
1176 let rev = revision.unwrap_or("main");
1177 let url = crate::chunked::build_download_url(model_id, rev, "README.md");
1178 let client = crate::chunked::build_client(token)?;
1179
1180 let response = client
1181 .get(url.as_str()) // BORROW: explicit .as_str() instead of Deref coercion
1182 .send()
1183 .await
1184 .map_err(|e| FetchError::Http(format!("failed to fetch README for {model_id}: {e}")))?;
1185
1186 if response.status() == reqwest::StatusCode::NOT_FOUND {
1187 return Ok(None);
1188 }
1189
1190 if !response.status().is_success() {
1191 return Err(FetchError::Http(format!(
1192 "README request for {model_id} returned status {}",
1193 response.status()
1194 )));
1195 }
1196
1197 let text = response
1198 .text()
1199 .await
1200 .map_err(|e| FetchError::Http(format!("failed to read README for {model_id}: {e}")))?;
1201
1202 Ok(Some(text))
1203}
1204
1205#[cfg(test)]
1206mod tests {
1207 use super::*;
1208
1209 #[test]
1210 fn normalize_8bit_variants() {
1211 assert_eq!(normalize_quantization_terms("AWQ 8bit"), "AWQ 8-bit");
1212 assert_eq!(normalize_quantization_terms("AWQ 8-bit"), "AWQ 8-bit");
1213 assert_eq!(normalize_quantization_terms("AWQ int8"), "AWQ 8-bit");
1214 assert_eq!(normalize_quantization_terms("AWQ INT8"), "AWQ 8-bit");
1215 }
1216
1217 #[test]
1218 fn normalize_4bit_variants() {
1219 assert_eq!(normalize_quantization_terms("GPTQ 4bit"), "GPTQ 4-bit");
1220 assert_eq!(normalize_quantization_terms("GPTQ INT4"), "GPTQ 4-bit");
1221 assert_eq!(normalize_quantization_terms("GPTQ 4-bit"), "GPTQ 4-bit");
1222 }
1223
1224 #[test]
1225 fn normalize_fp8_variants() {
1226 assert_eq!(normalize_quantization_terms("FP8"), "fp8");
1227 assert_eq!(normalize_quantization_terms("float8"), "fp8");
1228 assert_eq!(normalize_quantization_terms("fp8"), "fp8");
1229 }
1230
1231 #[test]
1232 fn normalize_passthrough() {
1233 assert_eq!(normalize_quantization_terms("llama 3"), "llama 3");
1234 assert_eq!(normalize_quantization_terms("RWKV-7"), "RWKV-7");
1235 }
1236
1237 // ---------- quants sibling discovery ----------
1238
1239 #[test]
1240 fn gguf_source_backlinks_finds_source_url() {
1241 let mut metadata = HashMap::new();
1242 metadata.insert(
1243 "general.source.url".to_owned(),
1244 "https://huggingface.co/poolside/Laguna-XS-2.1".to_owned(),
1245 );
1246 metadata.insert("general.architecture".to_owned(), "llama".to_owned());
1247 assert_eq!(
1248 gguf_source_backlinks(&metadata),
1249 vec!["https://huggingface.co/poolside/Laguna-XS-2.1"]
1250 );
1251 }
1252
1253 #[test]
1254 fn gguf_source_backlinks_finds_base_model_repo_url() {
1255 let mut metadata = HashMap::new();
1256 metadata.insert(
1257 "general.base_model.0.repo_url".to_owned(),
1258 "https://huggingface.co/poolside/Laguna-XS-2.1".to_owned(),
1259 );
1260 assert_eq!(gguf_source_backlinks(&metadata).len(), 1);
1261 }
1262
1263 #[test]
1264 fn gguf_source_backlinks_finds_huggingface_repository_key() {
1265 let mut metadata = HashMap::new();
1266 metadata.insert(
1267 "general.source.huggingface.repository".to_owned(),
1268 "poolside/Laguna-XS-2.1".to_owned(),
1269 );
1270 assert_eq!(gguf_source_backlinks(&metadata).len(), 1);
1271 }
1272
1273 #[test]
1274 fn gguf_source_backlinks_ignores_unrelated_keys() {
1275 let mut metadata = HashMap::new();
1276 metadata.insert("general.architecture".to_owned(), "llama".to_owned());
1277 metadata.insert("general.name".to_owned(), "Laguna-XS-2.1-GGUF".to_owned());
1278 assert!(gguf_source_backlinks(&metadata).is_empty());
1279 }
1280
1281 // ---------- classify_gguf_files ----------
1282
1283 #[test]
1284 fn classify_gguf_files_not_applicable_for_zero_or_one_file() {
1285 assert_eq!(classify_gguf_files(&[]), GgufFileSetKind::NotApplicable);
1286 assert_eq!(
1287 classify_gguf_files(&["model-Q4_K_M.gguf"]),
1288 GgufFileSetKind::NotApplicable
1289 );
1290 }
1291
1292 #[test]
1293 fn classify_gguf_files_recognizes_a_complete_shard_set() {
1294 let files = [
1295 "model-00001-of-00003.gguf",
1296 "model-00002-of-00003.gguf",
1297 "model-00003-of-00003.gguf",
1298 ];
1299 assert_eq!(
1300 classify_gguf_files(&files),
1301 GgufFileSetKind::Sharded {
1302 first_shard: "model-00001-of-00003.gguf".to_owned()
1303 }
1304 );
1305 }
1306
1307 #[test]
1308 fn classify_gguf_files_recognizes_a_shard_set_regardless_of_listing_order() {
1309 let files = [
1310 "model-00003-of-00003.gguf",
1311 "model-00001-of-00003.gguf",
1312 "model-00002-of-00003.gguf",
1313 ];
1314 // `first_shard` must always be the index-1 filename, not simply the
1315 // first one encountered in `files`.
1316 assert_eq!(
1317 classify_gguf_files(&files),
1318 GgufFileSetKind::Sharded {
1319 first_shard: "model-00001-of-00003.gguf".to_owned()
1320 }
1321 );
1322 }
1323
1324 #[test]
1325 fn classify_gguf_files_flags_quant_alternatives() {
1326 let files = ["model-Q4_K_M.gguf", "model-Q5_K_M.gguf", "model-Q8_0.gguf"];
1327 assert_eq!(
1328 classify_gguf_files(&files),
1329 GgufFileSetKind::QuantAlternatives
1330 );
1331 }
1332
1333 #[test]
1334 fn classify_gguf_files_flags_a_missing_shard_index() {
1335 // Claims 3-of-3 but only two files are present — incomplete.
1336 let files = ["model-00001-of-00003.gguf", "model-00003-of-00003.gguf"];
1337 assert_eq!(
1338 classify_gguf_files(&files),
1339 GgufFileSetKind::QuantAlternatives
1340 );
1341 }
1342
1343 #[test]
1344 fn classify_gguf_files_flags_a_duplicate_shard_index() {
1345 let files = [
1346 "model-00001-of-00003.gguf",
1347 "model-00001-of-00003.gguf",
1348 "model-00003-of-00003.gguf",
1349 ];
1350 assert_eq!(
1351 classify_gguf_files(&files),
1352 GgufFileSetKind::QuantAlternatives
1353 );
1354 }
1355
1356 #[test]
1357 fn classify_gguf_files_flags_mismatched_prefixes() {
1358 let files = ["model-a-00001-of-00002.gguf", "model-b-00002-of-00002.gguf"];
1359 assert_eq!(
1360 classify_gguf_files(&files),
1361 GgufFileSetKind::QuantAlternatives
1362 );
1363 }
1364
1365 #[test]
1366 fn classify_gguf_files_ignores_non_gguf_files() {
1367 let files = [
1368 "model-00001-of-00002.gguf",
1369 "model-00002-of-00002.gguf",
1370 "config.json",
1371 "README.md",
1372 ];
1373 assert_eq!(
1374 classify_gguf_files(&files),
1375 GgufFileSetKind::Sharded {
1376 first_shard: "model-00001-of-00002.gguf".to_owned()
1377 }
1378 );
1379 }
1380
1381 // ---------- gguf_size_range ----------
1382
1383 #[test]
1384 fn gguf_size_range_computes_min_max_for_quant_alternatives() {
1385 let files = vec![
1386 ("model-Q8_0.gguf", Some(20_000)),
1387 ("model-Q4_K_M.gguf", Some(10_000)),
1388 ("model-Q3_K_S.gguf", Some(14_000)),
1389 ];
1390 assert_eq!(gguf_size_range(files), Some((10_000, 20_000)));
1391 }
1392
1393 #[test]
1394 fn gguf_size_range_returns_none_for_a_sharded_set() {
1395 let files = vec![
1396 ("model-00001-of-00002.gguf", Some(10_000)),
1397 ("model-00002-of-00002.gguf", Some(10_000)),
1398 ];
1399 assert_eq!(gguf_size_range(files), None);
1400 }
1401
1402 #[test]
1403 fn gguf_size_range_still_classifies_sharded_when_one_shard_has_no_known_size() {
1404 // Regression test: a genuinely sharded 3-file set where the Hub API
1405 // happened not to report one shard's size must still classify as
1406 // Sharded (via the full filename list) and therefore return None —
1407 // dropping the unsized shard's filename before classification would
1408 // leave only 2 of 3 expected filenames, tripping the "file count
1409 // must equal the parsed total" check and misclassifying the whole
1410 // set as QuantAlternatives, which then reports a bogus min/max
1411 // range over only the 2 known-size shards.
1412 let files = vec![
1413 ("model-00001-of-00003.gguf", Some(5_000)),
1414 ("model-00002-of-00003.gguf", None),
1415 ("model-00003-of-00003.gguf", Some(5_000)),
1416 ];
1417 assert_eq!(gguf_size_range(files), None);
1418 }
1419
1420 #[test]
1421 fn gguf_size_range_excludes_unsized_files_from_min_max_but_keeps_classifying_on_all() {
1422 // A quant-alternatives set where one candidate's size is unknown:
1423 // classification still sees all 3 filenames (so 3 mutually-exclusive
1424 // `.gguf` files are correctly detected), but the unsized entry is
1425 // excluded from the min/max computation itself.
1426 let files = vec![
1427 ("model-Q8_0.gguf", Some(20_000)),
1428 ("model-Q4_K_M.gguf", None),
1429 ("model-Q3_K_S.gguf", Some(14_000)),
1430 ];
1431 assert_eq!(gguf_size_range(files), Some((14_000, 20_000)));
1432 }
1433
1434 // ---------- pick_backlink_representative ----------
1435
1436 fn repo_file(filename: &str, size: u64) -> crate::repo::RepoFile {
1437 crate::repo::RepoFile {
1438 // BORROW: explicit .to_owned() for &str → owned String field
1439 filename: filename.to_owned(),
1440 size: Some(size),
1441 sha256: None,
1442 }
1443 }
1444
1445 #[test]
1446 fn pick_backlink_representative_picks_first_shard_for_a_sharded_set() {
1447 // The metadata KV table (and any backlink key) lives only in the
1448 // first split, per llama.cpp's own convention — even though it is
1449 // not the smallest file here, it must still be the one chosen.
1450 let files = vec![
1451 repo_file("model-00001-of-00003.gguf", 5_000),
1452 repo_file("model-00002-of-00003.gguf", 5_000),
1453 repo_file("model-00003-of-00003.gguf", 1_000),
1454 ];
1455 assert_eq!(
1456 pick_backlink_representative(&files).map(|f| f.filename.as_str()),
1457 Some("model-00001-of-00003.gguf")
1458 );
1459 }
1460
1461 #[test]
1462 fn pick_backlink_representative_picks_smallest_for_quant_alternatives() {
1463 let files = vec![
1464 repo_file("model-Q8_0.gguf", 20_000),
1465 repo_file("model-Q4_K_M.gguf", 10_000),
1466 repo_file("model-Q3_K_S.gguf", 14_000),
1467 ];
1468 assert_eq!(
1469 pick_backlink_representative(&files).map(|f| f.filename.as_str()),
1470 Some("model-Q4_K_M.gguf")
1471 );
1472 }
1473
1474 #[test]
1475 fn pick_backlink_representative_picks_the_lone_file() {
1476 let files = vec![repo_file("model.gguf", 10_000)];
1477 assert_eq!(
1478 pick_backlink_representative(&files).map(|f| f.filename.as_str()),
1479 Some("model.gguf")
1480 );
1481 }
1482
1483 #[test]
1484 fn pick_backlink_representative_returns_none_for_no_gguf_files() {
1485 assert!(pick_backlink_representative(&[]).is_none());
1486 }
1487
1488 // ---------- fan_out_bounded ----------
1489
1490 #[tokio::test]
1491 async fn fan_out_bounded_preserves_item_order() {
1492 let items = vec!["a".to_owned(), "b".to_owned(), "c".to_owned()];
1493 let out =
1494 fan_out_bounded(items, 2, |item| async move { Some(format!("{item}-done")) }).await;
1495
1496 assert_eq!(
1497 out,
1498 vec![
1499 Some("a-done".to_owned()),
1500 Some("b-done".to_owned()),
1501 Some("c-done".to_owned()),
1502 ]
1503 );
1504 }
1505
1506 #[tokio::test]
1507 async fn fan_out_bounded_keeps_a_none_slot_for_per_item_failures() {
1508 let items = vec!["keep".to_owned(), "drop".to_owned(), "keep2".to_owned()];
1509 let out = fan_out_bounded(items, 2, |item| async move {
1510 if item == "drop" { None } else { Some(item) }
1511 })
1512 .await;
1513
1514 assert_eq!(
1515 out,
1516 vec![Some("keep".to_owned()), None, Some("keep2".to_owned()),]
1517 );
1518 }
1519
1520 #[tokio::test]
1521 async fn zip_into_map_drops_none_slots_and_keys_by_item() {
1522 let keys = vec!["a".to_owned(), "b".to_owned(), "c".to_owned()];
1523 let results = vec![Some(1), None, Some(3)];
1524
1525 let map = zip_into_map(keys, results);
1526
1527 assert_eq!(map.len(), 2);
1528 assert_eq!(map.get("a"), Some(&1));
1529 assert!(!map.contains_key("b"));
1530 assert_eq!(map.get("c"), Some(&3));
1531 }
1532
1533 #[tokio::test]
1534 async fn fan_out_bounded_respects_concurrency_limit() {
1535 // 6 items, concurrency 2: an atomic counter of in-flight tasks must
1536 // never exceed 2 at any point during the run.
1537 let in_flight = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1538 let max_seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1539 let items: Vec<u32> = (0..6).collect();
1540
1541 let in_flight_for_closure = Arc::clone(&in_flight);
1542 let max_seen_for_closure = Arc::clone(&max_seen);
1543 let out = fan_out_bounded(items, 2, move |item| {
1544 let in_flight = Arc::clone(&in_flight_for_closure);
1545 let max_seen = Arc::clone(&max_seen_for_closure);
1546 async move {
1547 let now = in_flight.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
1548 max_seen.fetch_max(now, std::sync::atomic::Ordering::SeqCst);
1549 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1550 in_flight.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
1551 Some(item)
1552 }
1553 })
1554 .await;
1555
1556 assert_eq!(out.len(), 6);
1557 assert!(
1558 max_seen.load(std::sync::atomic::Ordering::SeqCst) <= 2,
1559 "observed more than 2 tasks in flight at once"
1560 );
1561 }
1562
1563 // ---------- fetch_repo_size_summaries_concurrent ----------
1564
1565 #[tokio::test]
1566 async fn fetch_repo_size_summaries_concurrent_reports_a_malformed_token_loudly() {
1567 // A token with an embedded newline can never become a valid HTTP
1568 // header value — `build_client` fails before any per-repo request
1569 // is attempted. This must surface as an `Err`, the same way
1570 // `search_models` already fails loudly for the identical cause,
1571 // rather than silently degrading to an empty map indistinguishable
1572 // from "every repo failed to size".
1573 let result =
1574 fetch_repo_size_summaries_concurrent(vec!["org/model".to_owned()], Some("bad\ntoken"))
1575 .await;
1576 assert!(
1577 result.is_err(),
1578 "a malformed token must be reported as an error"
1579 );
1580 }
1581}