Skip to main content

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}