Skip to main content

yapi_ai/
llama.rs

1//! llama.cpp: a local `llama-server` in router mode as a provider. Its
2//! loaded models are chat models over the OpenAI-compatible API and
3//! classifiers over `llama-cpp-classify`.
4//!
5//! Port of `extensions/llama/client.ts` and `provider.ts` in pi-coding-agent
6//! `v1.0.0`, where llama.cpp is the built-in extension `builtin:llama.cpp`.
7
8use std::time::Duration;
9
10use serde::Deserialize;
11use serde_json::{Map, Value, json};
12use tokio_util::sync::CancellationToken;
13use yapi_types::model::{ClassifierModel, Model};
14
15use crate::model_catalog::ProviderModels;
16
17/// The provider id.
18pub const PROVIDER_ID: &str = "llama.cpp";
19/// Where the server runs unless configured.
20pub const DEFAULT_SERVER_URL: &str = "http://127.0.0.1:8080";
21/// The credential setting with the server URL.
22pub const BASE_URL_ENV: &str = "LLAMA_BASE_URL";
23/// pi gives each server request 15s.
24const REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
25
26/// A model's state on the server.
27#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
28pub struct Status {
29    /// `unloaded`, `loading`, `loaded`, `downloading` or `sleeping`.
30    pub value: String,
31    /// The server's arguments for the model.
32    #[serde(default)]
33    pub args: Option<Vec<String>>,
34    /// Whether its last load failed.
35    #[serde(default)]
36    pub failed: Option<bool>,
37    /// The exit code of a failed load.
38    #[serde(default)]
39    pub exit_code: Option<i64>,
40    /// Download progress by file: `{done, total}`.
41    #[serde(default)]
42    pub progress: Option<Value>,
43}
44
45/// What the model takes and produces.
46#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
47pub struct Architecture {
48    /// Input modalities, such as `image`.
49    #[serde(default)]
50    pub input_modalities: Option<Vec<String>>,
51}
52
53/// Facts about the model's weights and runtime.
54#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
55pub struct Meta {
56    /// The context size it runs with.
57    #[serde(default)]
58    pub n_ctx: Option<u64>,
59    /// The context size it was trained with.
60    #[serde(default)]
61    pub n_ctx_train: Option<u64>,
62}
63
64/// One entry of the server's `/models`.
65#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
66pub struct ModelInfo {
67    /// The model id requests name.
68    pub id: String,
69    /// Its state.
70    pub status: Status,
71    /// Its modalities.
72    #[serde(default)]
73    pub architecture: Option<Architecture>,
74    /// Where it comes from, such as `preset`.
75    #[serde(default)]
76    pub source: Option<String>,
77    /// Its weights and runtime.
78    #[serde(default)]
79    pub meta: Option<Meta>,
80}
81
82impl ModelInfo {
83    /// Loaded, or asleep and woken by the next request.
84    pub fn is_loaded(&self) -> bool {
85        matches!(self.status.value.as_str(), "loaded" | "sleeping")
86    }
87}
88
89/// What `/props` reports.
90#[derive(Clone, Debug, Default, PartialEq)]
91pub struct Props {
92    /// Whether the router loads unloaded presets on first use.
93    pub models_autoload: Option<bool>,
94    /// The model's chat template.
95    pub chat_template: Option<String>,
96}
97
98/// pi's `normalizeLlamaServerUrl`: an http(s) URL without query, fragment,
99/// trailing slash or `/v1`.
100pub fn normalize_server_url(value: &str) -> Result<String, String> {
101    let mut url = url::Url::parse(value.trim()).map_err(|_| "Invalid URL".to_owned())?;
102    if !matches!(url.scheme(), "http" | "https") {
103        return Err("Server URL must use http or https".into());
104    }
105    url.set_fragment(None);
106    url.set_query(None);
107    let path = url.path().trim_end_matches('/');
108    let path = path.strip_suffix("/v1").unwrap_or(path).to_owned();
109    url.set_path(if path.is_empty() { "/" } else { &path });
110    Ok(url.to_string().trim_end_matches('/').to_owned())
111}
112
113/// The OpenAI-compatible base URL of a server.
114pub fn inference_url(server_url: &str) -> String {
115    let normalized = normalize_server_url(server_url)
116        .unwrap_or_else(|_| server_url.trim_end_matches('/').to_owned());
117    format!("{normalized}/v1")
118}
119
120/// A client of one server.
121#[derive(Clone, Debug)]
122pub struct Client {
123    /// The normalized server URL.
124    pub server_url: String,
125    api_key: Option<String>,
126}
127
128fn error_message(payload: &Value, fallback: String) -> String {
129    payload["error"]["message"]
130        .as_str()
131        .filter(|message| !message.is_empty())
132        .map_or(fallback, str::to_owned)
133}
134
135impl Client {
136    /// A client of `server_url`, sending `api_key` when there is one.
137    pub fn new(server_url: &str, api_key: Option<String>) -> Result<Client, String> {
138        Ok(Client {
139            server_url: normalize_server_url(server_url)?,
140            api_key: api_key.filter(|key| !key.is_empty()),
141        })
142    }
143
144    async fn request(
145        &self,
146        method: reqwest::Method,
147        path: &str,
148        body: Option<Value>,
149        cancel: &CancellationToken,
150    ) -> Result<Value, String> {
151        let mut request = crate::http::client()
152            .request(method, format!("{}{path}", self.server_url))
153            .timeout(REQUEST_TIMEOUT);
154        if let Some(body) = body {
155            request = request
156                .header("Content-Type", "application/json")
157                .body(yapi_types::json::stringify(&body));
158        }
159        if let Some(key) = &self.api_key {
160            request = request.header("Authorization", format!("Bearer {key}"));
161        }
162        let Some(response) = cancel.run_until_cancelled(request.send()).await else {
163            return Err("This operation was aborted".into());
164        };
165        let response = response.map_err(|err| {
166            if err.is_timeout() {
167                "The operation was aborted due to timeout".to_owned()
168            } else {
169                "fetch failed".to_owned()
170            }
171        })?;
172        let status = response.status();
173        let payload: Value = response
174            .bytes()
175            .await
176            .ok()
177            .and_then(|bytes| serde_json::from_slice(&bytes).ok())
178            .unwrap_or(Value::Null);
179        if !status.is_success() {
180            return Err(error_message(
181                &payload,
182                format!("llama.cpp returned HTTP {}", status.as_u16()),
183            ));
184        }
185        Ok(payload)
186    }
187
188    /// The server's models. `reload` rescans its model directories.
189    pub async fn list(
190        &self,
191        reload: bool,
192        cancel: &CancellationToken,
193    ) -> Result<Vec<ModelInfo>, String> {
194        let path = if reload {
195            "/models?reload=1"
196        } else {
197            "/models"
198        };
199        let payload = self
200            .request(reqwest::Method::GET, path, None, cancel)
201            .await?;
202        let Some(data) = payload["data"].as_array() else {
203            return Err("llama.cpp returned an invalid model catalog".into());
204        };
205        let valid = |entry: &Value| entry["id"].is_string() && entry["status"]["value"].is_string();
206        if !data.iter().all(valid) {
207            return Err("Server is not running in llama.cpp router mode".into());
208        }
209        Ok(data
210            .iter()
211            .filter_map(|entry| serde_json::from_value(entry.clone()).ok())
212            .collect())
213    }
214
215    /// `/props`, for `model` without loading it, or for the server.
216    pub async fn props(
217        &self,
218        model: Option<&str>,
219        cancel: &CancellationToken,
220    ) -> Result<Props, String> {
221        let path = match model {
222            Some(model) => format!(
223                "/props?{}",
224                url::form_urlencoded::Serializer::new(String::new())
225                    .append_pair("model", model)
226                    .append_pair("autoload", "false")
227                    .finish()
228            ),
229            None => "/props".into(),
230        };
231        let payload = self
232            .request(reqwest::Method::GET, &path, None, cancel)
233            .await?;
234        Ok(Props {
235            models_autoload: payload["models_autoload"].as_bool(),
236            chat_template: payload["chat_template"].as_str().map(str::to_owned),
237        })
238    }
239
240    /// POSTs `{"model": model}` to `path`.
241    async fn post_model(
242        &self,
243        path: &str,
244        model: &str,
245        cancel: &CancellationToken,
246    ) -> Result<(), String> {
247        let body = Some(json!({ "model": model }));
248        self.request(reqwest::Method::POST, path, body, cancel)
249            .await
250            .map(drop)
251    }
252
253    /// Asks the router to load `model`.
254    pub async fn load(&self, model: &str, cancel: &CancellationToken) -> Result<(), String> {
255        self.post_model("/models/load", model, cancel).await
256    }
257
258    /// Asks the router to unload `model`, or to stop loading or downloading it.
259    pub async fn unload(&self, model: &str, cancel: &CancellationToken) -> Result<(), String> {
260        self.post_model("/models/unload", model, cancel).await
261    }
262
263    /// Unloads `model` and waits until it is unloaded.
264    pub async fn unload_and_wait(
265        &self,
266        model: &str,
267        cancel: &CancellationToken,
268    ) -> Result<(), String> {
269        self.unload(model, cancel).await?;
270        loop {
271            let models = self.list(false, cancel).await?;
272            match models.iter().find(|entry| entry.id == model) {
273                Some(entry) if entry.status.value != "unloaded" => {}
274                _ => return Ok(()),
275            }
276            sleep(Duration::from_millis(100), cancel).await?;
277        }
278    }
279
280    /// Asks the router to download `model`, a Hugging Face `repo[:quant]`.
281    pub async fn download(&self, model: &str, cancel: &CancellationToken) -> Result<(), String> {
282        self.post_model("/models", model, cancel).await
283    }
284
285    /// Loads `model` and waits until it is loaded, reporting progress.
286    pub async fn load_and_wait(
287        &self,
288        model: &str,
289        progress: &(dyn Fn(Progress) + Send + Sync),
290        cancel: &CancellationToken,
291    ) -> Result<ModelInfo, String> {
292        self.load(model, cancel).await?;
293        progress(Progress::message("Loading model"));
294        loop {
295            let models = self.list(false, cancel).await?;
296            match models.into_iter().find(|entry| entry.id == model) {
297                Some(entry) if entry.status.value == "loaded" => return Ok(entry),
298                Some(entry) if entry.status.failed == Some(true) => {
299                    return Err(match entry.status.exit_code {
300                        Some(code) => format!("Model exited with code {code}"),
301                        None => "Model failed to load".into(),
302                    });
303                }
304                _ => {}
305            }
306            sleep(Duration::from_millis(250), cancel).await?;
307        }
308    }
309
310    /// Downloads `model`, waits until the router lists it and returns the
311    /// catalog after a rescan.
312    pub async fn download_and_wait(
313        &self,
314        model: &str,
315        progress: &(dyn Fn(Progress) + Send + Sync),
316        cancel: &CancellationToken,
317    ) -> Result<Vec<ModelInfo>, String> {
318        self.download(model, cancel).await?;
319        progress(Progress::message("Downloading model"));
320        let mut downloading = false;
321        let mut polls = 0;
322        loop {
323            let models = self.list(false, cancel).await?;
324            polls += 1;
325            match models.iter().find(|entry| entry.id == model) {
326                Some(entry) if entry.status.value == "downloading" => {
327                    downloading = true;
328                    if let Some(report) = entry.status.progress.as_ref().and_then(download_progress)
329                    {
330                        progress(report);
331                    }
332                }
333                Some(_) if downloading || polls >= 2 => return self.list(true, cancel).await,
334                _ => {}
335            }
336            sleep(Duration::from_millis(500), cancel).await?;
337        }
338    }
339}
340
341async fn sleep(duration: Duration, cancel: &CancellationToken) -> Result<(), String> {
342    cancel
343        .run_until_cancelled(tokio::time::sleep(duration))
344        .await
345        .ok_or_else(|| "Cancelled".into())
346}
347
348/// What a long operation reports.
349#[derive(Clone, Debug, PartialEq)]
350pub struct Progress {
351    /// What happens now.
352    pub message: String,
353    /// How far along, from 0 to 1, when known.
354    pub ratio: Option<f64>,
355    /// Detail such as bytes downloaded.
356    pub detail: Option<String>,
357}
358
359impl Progress {
360    fn message(message: &str) -> Progress {
361        Progress {
362            message: message.into(),
363            ratio: None,
364            detail: None,
365        }
366    }
367}
368
369/// pi's `formatBytes`.
370pub fn format_bytes(bytes: f64) -> String {
371    if bytes < 1024.0 {
372        return format!("{bytes} B");
373    }
374    let units = ["KiB", "MiB", "GiB", "TiB"];
375    let mut value = bytes / 1024.0;
376    let mut unit = units[0];
377    for next in &units[1..] {
378        if value < 1024.0 {
379            break;
380        }
381        value /= 1024.0;
382        unit = next;
383    }
384    let digits = if value >= 10.0 { 1 } else { 2 };
385    format!("{} {unit}", yapi_types::js::to_fixed(value, digits))
386}
387
388/// pi's `parseDownloadProgress` over a model's per-file `{done, total}`.
389fn download_progress(data: &Value) -> Option<Progress> {
390    let files = data
391        .get("progress")
392        .filter(|nested| nested.is_object())
393        .unwrap_or(data);
394    let (mut done, mut total) = (0.0, 0.0);
395    for entry in files.as_object()?.values() {
396        if let (Some(d), Some(t)) = (entry["done"].as_f64(), entry["total"].as_f64()) {
397            done += d;
398            total += t;
399        }
400    }
401    (total > 0.0).then(|| Progress {
402        message: "Downloading model".into(),
403        ratio: Some(done / total),
404        detail: Some(format!("{} / {}", format_bytes(done), format_bytes(total))),
405    })
406}
407
408/// Whether a model can serve requests: loaded or asleep, or an unloaded
409/// preset the router loads on first use.
410pub fn is_selectable(model: &ModelInfo, router_autoload: bool) -> bool {
411    match model.status.value.as_str() {
412        "loaded" | "sleeping" => true,
413        "unloaded" => {
414            router_autoload
415                && model.status.failed != Some(true)
416                && model.source.as_deref() == Some("preset")
417        }
418        _ => false,
419    }
420}
421
422/// The `--ctx-size`, `-c` or `-ctx` argument the server runs the model with.
423fn configured_context(model: &ModelInfo) -> Option<u64> {
424    let args = model.status.args.as_deref().unwrap_or_default();
425    args.windows(2).find_map(|pair| {
426        matches!(pair[0].as_str(), "--ctx-size" | "-c" | "-ctx")
427            .then(|| pair[1].parse::<u64>().ok().filter(|value| *value > 0))
428            .flatten()
429    })
430}
431
432/// pi's `contextWindowOf`: the runtime context, the configured one, the
433/// last known one, the training context, or 128000.
434pub fn context_window(model: &ModelInfo, cached: Option<u64>) -> u64 {
435    let meta = model.meta.clone().unwrap_or_default();
436    meta.n_ctx
437        .filter(|value| *value > 0)
438        .or_else(|| configured_context(model))
439        .or(cached.filter(|value| *value > 0))
440        .or(meta.n_ctx_train.filter(|value| *value > 0))
441        .unwrap_or(128_000)
442}
443
444const ZERO_COST: &str = r#"{"input":0,"output":0,"cacheRead":0,"cacheWrite":0}"#;
445
446/// pi's `toPiModel`: the chat model of a server model.
447pub fn chat_model(
448    model: &ModelInfo,
449    server_url: &str,
450    props: Option<&Props>,
451    cached: Option<u64>,
452) -> Option<Model> {
453    let window = context_window(model, cached);
454    let reasoning = props
455        .and_then(|props| props.chat_template.as_deref())
456        .is_some_and(|template| template.contains("enable_thinking"));
457    let image = model
458        .architecture
459        .as_ref()
460        .and_then(|architecture| architecture.input_modalities.as_ref())
461        .is_some_and(|modalities| modalities.iter().any(|kind| kind == "image"));
462    let mut value = json!({
463        "id": model.id,
464        "name": model.id,
465        "api": "openai-completions",
466        "provider": PROVIDER_ID,
467        "baseUrl": inference_url(server_url),
468        "reasoning": reasoning,
469    });
470    if reasoning {
471        value["thinkingLevelMap"] = json!({
472            "off": "off", "minimal": null, "low": null, "medium": "medium", "high": null, "xhigh": null
473        });
474    }
475    value["input"] = if image {
476        json!(["text", "image"])
477    } else {
478        json!(["text"])
479    };
480    value["cost"] = serde_json::from_str(ZERO_COST).unwrap_or(Value::Null);
481    value["contextWindow"] = json!(window);
482    value["maxTokens"] = json!(window);
483    let mut compat = json!({
484        "supportsStore": false,
485        "supportsDeveloperRole": false,
486        "supportsReasoningEffort": false,
487        "supportsUsageInStreaming": true,
488        "supportsStrictMode": false,
489        "maxTokensField": "max_tokens",
490    });
491    if reasoning {
492        compat["thinkingFormat"] = json!("qwen-chat-template");
493    }
494    value["compat"] = compat;
495    serde_json::from_value(value).ok()
496}
497
498/// pi's `toPiClassifierModel`: the same model as a classifier.
499pub fn classifier_model(
500    model: &ModelInfo,
501    server_url: &str,
502    cached: Option<u64>,
503) -> Option<ClassifierModel> {
504    serde_json::from_value(json!({
505        "type": "classifier",
506        "id": model.id,
507        "name": model.id,
508        "api": "llama-cpp-classify",
509        "provider": PROVIDER_ID,
510        "baseUrl": server_url,
511        "input": ["text"],
512        "cost": serde_json::from_str::<Value>(ZERO_COST).unwrap_or(Value::Null),
513        "contextWindow": context_window(model, cached),
514    }))
515    .ok()
516}
517
518/// The models a server serves now: loaded ones with their templates, then
519/// the same models as classifiers. `cached` gives earlier context windows.
520pub async fn fetch_models(
521    client: &Client,
522    cached: &Map<String, Value>,
523    cancel: &CancellationToken,
524) -> Result<ProviderModels, String> {
525    let catalog = client.list(false, cancel).await?;
526    let autoload = catalog
527        .iter()
528        .any(|model| model.status.value == "unloaded" && model.source.as_deref() == Some("preset"))
529        && client
530            .props(None, cancel)
531            .await
532            .ok()
533            .and_then(|props| props.models_autoload)
534            == Some(true);
535    let selectable: Vec<&ModelInfo> = catalog
536        .iter()
537        .filter(|model| is_selectable(model, autoload))
538        .collect();
539    let window = |id: &str| cached.get(id).and_then(Value::as_u64);
540    let mut models = ProviderModels::default();
541    for model in &selectable {
542        // Only loaded models give their template without side effects.
543        let props = if model.status.value == "loaded" {
544            Some(client.props(Some(&model.id), cancel).await?)
545        } else {
546            None
547        };
548        models.chat.extend(chat_model(
549            model,
550            &client.server_url,
551            props.as_ref(),
552            window(&model.id),
553        ));
554    }
555    for model in &selectable {
556        models.classifiers.extend(classifier_model(
557            model,
558            &client.server_url,
559            window(&model.id),
560        ));
561    }
562    Ok(models)
563}
564
565/// Hugging Face, where the router downloads models from.
566pub mod huggingface {
567    use std::path::PathBuf;
568
569    use serde_json::Value;
570    use tokio_util::sync::CancellationToken;
571
572    const DEFAULT_URL: &str = "https://huggingface.co";
573
574    /// A search result.
575    #[derive(Clone, Debug, PartialEq)]
576    pub struct Model {
577        /// `owner/repository`.
578        pub id: String,
579        /// Downloads in the last month.
580        pub downloads: u64,
581    }
582
583    /// A GGUF quantization of a repository.
584    #[derive(Clone, Debug, PartialEq)]
585    pub struct Quantization {
586        /// Such as `Q4_K_M`.
587        pub name: String,
588        /// Bytes over all shards, when every shard reports its size.
589        pub size: Option<u64>,
590    }
591
592    /// A repository's access and quantizations.
593    #[derive(Clone, Debug, PartialEq)]
594    pub struct Details {
595        /// The repository id.
596        pub id: String,
597        /// `auto` or `manual` approval, when gated.
598        pub gated: Option<String>,
599        /// Quantizations, Q4_K_M first, then by size.
600        pub quantizations: Vec<Quantization>,
601    }
602
603    /// pi's `findHuggingFaceToken`: `HF_TOKEN`, then the token files of the
604    /// Hugging Face CLI. `env` reads the environment.
605    pub fn find_token(
606        env: &dyn Fn(&str) -> Option<String>,
607        home: Option<PathBuf>,
608    ) -> Option<String> {
609        if let Some(token) = env("HF_TOKEN")
610            .map(|token| token.trim().to_owned())
611            .filter(|token| !token.is_empty())
612        {
613            return Some(token);
614        }
615        let mut paths: Vec<PathBuf> = Vec::new();
616        paths.extend(env("HF_TOKEN_PATH").map(PathBuf::from));
617        paths.extend(env("HF_HOME").map(|home| PathBuf::from(home).join("token")));
618        paths.extend(
619            env("XDG_CACHE_HOME")
620                .map(|cache| PathBuf::from(cache).join("huggingface").join("token")),
621        );
622        paths.extend(home.map(|home| home.join(".cache").join("huggingface").join("token")));
623        let mut seen = Vec::new();
624        for path in paths {
625            if seen.contains(&path) {
626                continue;
627            }
628            if let Ok(token) = std::fs::read_to_string(&path) {
629                let token = token.trim();
630                if !token.is_empty() {
631                    return Some(token.to_owned());
632                }
633            }
634            seen.push(path);
635        }
636        None
637    }
638
639    /// A Hugging Face API client.
640    #[derive(Clone, Debug)]
641    pub struct Client {
642        token: Option<String>,
643        base_url: String,
644    }
645
646    impl Client {
647        /// A client of `base_url`, or huggingface.co, sending `token`.
648        pub fn new(token: Option<String>, base_url: Option<&str>) -> Client {
649            Client {
650                token,
651                base_url: base_url
652                    .unwrap_or(DEFAULT_URL)
653                    .trim_end_matches('/')
654                    .to_owned(),
655            }
656        }
657
658        async fn request(&self, path: &str, cancel: &CancellationToken) -> Result<Value, String> {
659            let mut request = crate::http::client()
660                .get(format!("{}{path}", self.base_url))
661                .timeout(super::REQUEST_TIMEOUT);
662            if let Some(token) = &self.token {
663                request = request.header("Authorization", format!("Bearer {token}"));
664            }
665            let Some(response) = cancel.run_until_cancelled(request.send()).await else {
666                return Err("This operation was aborted".into());
667            };
668            let response = response.map_err(|_| "fetch failed".to_owned())?;
669            let status = response.status().as_u16();
670            let header = |name: &str| {
671                response
672                    .headers()
673                    .get(name)
674                    .and_then(|value| value.to_str().ok())
675                    .map(str::to_owned)
676            };
677            let (retry_after, rate_limit) = (header("retry-after"), header("ratelimit"));
678            let payload: Value = response
679                .bytes()
680                .await
681                .ok()
682                .and_then(|bytes| serde_json::from_slice(&bytes).ok())
683                .unwrap_or(Value::Null);
684            if (200..300).contains(&status) {
685                return Ok(payload);
686            }
687            if status == 429 {
688                let delay = retry_after
689                    .and_then(|value| value.trim().parse::<f64>().ok())
690                    .filter(|delay| *delay != 0.0)
691                    .map(|delay| delay.to_string())
692                    .or_else(|| {
693                        rate_limit.and_then(|value| {
694                            value
695                                .split(';')
696                                .find_map(|part| part.trim().strip_prefix("t="))
697                                .filter(|digits| {
698                                    !digits.is_empty() && digits.chars().all(|c| c.is_ascii_digit())
699                                })
700                                .map(str::to_owned)
701                        })
702                    });
703                return Err(match delay {
704                    Some(delay) => format!("Hugging Face rate limit reached; retry in {delay}s"),
705                    None => "Hugging Face rate limit reached".into(),
706                });
707            }
708            Err(payload["error"]
709                .as_str()
710                .filter(|error| !error.is_empty())
711                .map_or_else(
712                    || format!("Hugging Face returned HTTP {status}"),
713                    str::to_owned,
714                ))
715        }
716
717        /// GGUF repositories matching `query`, most downloaded first.
718        pub async fn search(
719            &self,
720            query: &str,
721            cancel: &CancellationToken,
722        ) -> Result<Vec<Model>, String> {
723            let query = url::form_urlencoded::Serializer::new(String::new())
724                .append_pair("search", query)
725                .append_pair("filter", "gguf")
726                .append_pair("sort", "downloads")
727                .append_pair("direction", "-1")
728                .append_pair("limit", "20")
729                .finish();
730            let payload = self
731                .request(&format!("/api/models?{query}"), cancel)
732                .await?;
733            let Some(entries) = payload.as_array() else {
734                return Err("Hugging Face returned invalid search results".into());
735            };
736            Ok(entries
737                .iter()
738                .filter_map(|entry| {
739                    Some(Model {
740                        id: entry["id"].as_str()?.to_owned(),
741                        downloads: entry["downloads"].as_f64().map_or(0, |count| count as u64),
742                    })
743                })
744                .collect())
745        }
746
747        /// A repository's access and quantizations.
748        pub async fn details(
749            &self,
750            id: &str,
751            cancel: &CancellationToken,
752        ) -> Result<Details, String> {
753            let encoded: Vec<String> = id
754                .split('/')
755                .map(|part| url::form_urlencoded::byte_serialize(part.as_bytes()).collect())
756                .collect();
757            let payload = self
758                .request(
759                    &format!("/api/models/{}?blobs=true", encoded.join("/")),
760                    cancel,
761                )
762                .await?;
763            if !payload.is_object() {
764                return Err("Hugging Face returned invalid model details".into());
765            }
766            let mut sizes: Vec<(String, u64, bool)> = Vec::new();
767            for file in payload["siblings"].as_array().into_iter().flatten() {
768                let Some(path) = file["rfilename"].as_str() else {
769                    continue;
770                };
771                if !path.to_lowercase().ends_with(".gguf") {
772                    continue;
773                }
774                let name = path.rsplit('/').next().unwrap_or(path);
775                if name.to_lowercase().starts_with("mmproj") {
776                    continue;
777                }
778                let Some(quantization) = quantization(&name[..name.len() - 5]) else {
779                    continue;
780                };
781                let index = match sizes.iter().position(|(name, _, _)| *name == quantization) {
782                    Some(index) => index,
783                    None => {
784                        sizes.push((quantization, 0, true));
785                        sizes.len() - 1
786                    }
787                };
788                match file["size"].as_f64() {
789                    Some(size) => sizes[index].1 += size as u64,
790                    None => sizes[index].2 = false,
791                }
792            }
793            let mut quantizations: Vec<Quantization> = sizes
794                .into_iter()
795                .map(|(name, total, complete)| Quantization {
796                    name,
797                    size: complete.then_some(total),
798                })
799                .collect();
800            quantizations.sort_by(|left, right| {
801                if left.name == "Q4_K_M" {
802                    return std::cmp::Ordering::Less;
803                }
804                if right.name == "Q4_K_M" {
805                    return std::cmp::Ordering::Greater;
806                }
807                left.size
808                    .unwrap_or(u64::MAX)
809                    .cmp(&right.size.unwrap_or(u64::MAX))
810                    .then_with(|| yapi_types::collate::locale_compare(&left.name, &right.name))
811            });
812            let gated = payload["gated"]
813                .as_str()
814                .filter(|gated| matches!(*gated, "auto" | "manual"))
815                .map(str::to_owned);
816            Ok(Details {
817                id: payload["id"].as_str().unwrap_or(id).to_owned(),
818                gated,
819                quantizations,
820            })
821        }
822    }
823
824    /// The quantization a GGUF file stem names, upper-cased, as pi's
825    /// `QUANTIZATION_PATTERN` matches it once the shard suffix is removed.
826    pub fn quantization(stem: &str) -> Option<String> {
827        let stem = match stem.rfind("-of-") {
828            Some(of)
829                if stem.len() == of + 9
830                    && stem[of + 4..].chars().all(|c| c.is_ascii_digit())
831                    && of >= 6
832                    && stem.as_bytes()[of - 6] == b'-'
833                    && stem[of - 5..of].chars().all(|c| c.is_ascii_digit()) =>
834            {
835                &stem[..of - 6]
836            }
837            _ => stem,
838        };
839        static PATTERN: std::sync::LazyLock<regex_lite::Regex> = std::sync::LazyLock::new(|| {
840            regex_lite::Regex::new(
841                r"(?i)(?:^|[-_.])((?:UD-)?(?:IQ\d(?:_[A-Z0-9]+)+|Q\d(?:_[A-Z0-9]+)+|BF16|F16|F32|MXFP\d(?:_[A-Z0-9]+)*))$",
842            )
843            .expect("constant pattern; `normalizes_server_urls`")
844        });
845        PATTERN
846            .captures(stem)
847            .and_then(|captures| captures.get(1))
848            .map(|found| found.as_str().to_uppercase())
849    }
850}
851
852#[cfg(test)]
853mod tests {
854    use super::*;
855
856    #[test]
857    fn normalizes_server_urls() {
858        assert_eq!(
859            normalize_server_url(" http://127.0.0.1:8080/v1/?x=1#y ").unwrap(),
860            "http://127.0.0.1:8080"
861        );
862        assert_eq!(
863            normalize_server_url("https://host/llama/").unwrap(),
864            "https://host/llama"
865        );
866        assert_eq!(
867            normalize_server_url("ftp://host").unwrap_err(),
868            "Server URL must use http or https"
869        );
870        assert_eq!(inference_url("http://h:1"), "http://h:1/v1");
871        assert_eq!(format_bytes(1536.0), "1.50 KiB");
872        assert_eq!(format_bytes(12.0 * 1024.0 * 1024.0), "12.0 MiB");
873        assert_eq!(
874            huggingface::quantization("Qwen3-8B-Q4_K_M-00001-of-00002").as_deref(),
875            Some("Q4_K_M")
876        );
877        assert_eq!(
878            huggingface::quantization("model.ud-iq2_xxs").as_deref(),
879            Some("UD-IQ2_XXS")
880        );
881        assert_eq!(huggingface::quantization("readme"), None);
882    }
883
884    #[test]
885    fn builds_models_like_pi() {
886        let info: ModelInfo = serde_json::from_value(json!({
887            "id": "qwen", "status": {"value": "loaded", "args": ["-c", "8192"]},
888            "architecture": {"input_modalities": ["text", "image"]}
889        }))
890        .unwrap();
891        let props = Props {
892            models_autoload: None,
893            chat_template: Some("{% if enable_thinking %}".into()),
894        };
895        let model = chat_model(&info, "http://h:8080", Some(&props), None).unwrap();
896        assert_eq!(model.context_window, 8192);
897        assert!(model.reasoning);
898        assert_eq!(model.base_url, "http://h:8080/v1");
899        assert!(model.accepts_images());
900        let classifier = classifier_model(&info, "http://h:8080", Some(4096)).unwrap();
901        assert_eq!(classifier.api, "llama-cpp-classify");
902        assert_eq!(classifier.base_url, "http://h:8080");
903        let preset: ModelInfo = serde_json::from_value(json!({
904            "id": "p", "status": {"value": "unloaded"}, "source": "preset"
905        }))
906        .unwrap();
907        assert!(is_selectable(&preset, true));
908        assert!(!is_selectable(&preset, false));
909        assert_eq!(context_window(&preset, None), 128_000);
910    }
911}