ravenrs_ai/
models.rs

1//! The models: which fit, and their files in a folder.
2
3use std::fmt;
4use std::fs::{self, File};
5use std::io::{Read, Write};
6use std::path::{Path, PathBuf};
7use std::sync::atomic::{AtomicBool, Ordering};
8
9use sha2::{Digest, Sha256};
10
11use crate::Machine;
12
13/// The memory left for the system and the other apps while a model runs.
14const MEMORY_LEFT_FREE: u64 = 2_000_000_000;
15/// The disk left free after a download.
16const DISK_LEFT_FREE: u64 = 2_000_000_000;
17const CHUNK: usize = 1 << 16;
18/// What marks a file still being downloaded.
19const PART: &str = "part";
20
21/// A file of a model, pinned to one published version.
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub struct ModelFile {
24    /// Its name in the folder.
25    pub name: &'static str,
26    /// Where it comes from, pinned to a commit.
27    pub url: &'static str,
28    /// Its size in bytes.
29    pub bytes: u64,
30    /// Its SHA-256, in lowercase hexadecimal.
31    pub sha256: &'static str,
32}
33
34/// The kind of model, which decides how it runs and how to talk to it.
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36#[non_exhaustive]
37pub enum ModelFamily {
38    /// Qwen3.
39    Qwen3,
40    /// Qwen 2.5.
41    Qwen2,
42    /// Gemma 3.
43    Gemma3,
44    /// Llama 3.
45    Llama3,
46    /// Phi-3.
47    Phi3,
48    /// Mistral.
49    Mistral,
50    /// LFM3, from Liquid AI.
51    Lfm3,
52    /// Whisper, from OpenAI, which hears rather than writes.
53    Whisper,
54    /// Kokoro, from hexgrad, which only speaks.
55    Kokoro,
56}
57
58/// A model an app knows how to use.
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub struct LocalModel {
61    /// Its name in the folder, which never changes.
62    pub id: &'static str,
63    /// The name people read.
64    pub name: &'static str,
65    /// The license of its weights, which is not the app's.
66    pub license: &'static str,
67    /// How it runs.
68    pub family: ModelFamily,
69    /// The most memory it takes while answering, in bytes.
70    pub memory: u64,
71    /// How many of every hundred requests it understood when measured, if it was.
72    pub understood: Option<u8>,
73    /// Its files.
74    pub files: &'static [ModelFile],
75}
76
77impl LocalModel {
78    /// All of its files, in bytes.
79    pub fn bytes(&self) -> u64 {
80        self.files.iter().map(|file| file.bytes).sum()
81    }
82
83    /// Whether it can run here: enough memory for it, and room on the disk unless
84    /// it is already there.
85    pub fn fits(&self, machine: &Machine, installed: bool) -> bool {
86        let memory = machine.memory >= self.memory.saturating_add(MEMORY_LEFT_FREE);
87        let disk = installed || machine.free_disk >= self.bytes().saturating_add(DISK_LEFT_FREE);
88        memory && disk
89    }
90}
91
92/// Where a model stands in a folder.
93#[derive(Clone, Debug, PartialEq, Eq)]
94#[non_exhaustive]
95pub enum ModelState {
96    /// Not there.
97    Absent,
98    /// All its files are there, with the right sizes.
99    Installed,
100}
101
102/// Why a model did not arrive. Nothing half done stays looking installed.
103#[derive(Clone, Debug, PartialEq, Eq)]
104#[non_exhaustive]
105pub enum ModelError {
106    /// The connection failed or dropped.
107    Network(String),
108    /// Reading or writing the folder failed.
109    Disk(String),
110    /// What arrived is not the published file.
111    WrongFile,
112    /// Someone cancelled it.
113    Cancelled,
114}
115
116impl fmt::Display for ModelError {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        match self {
119            Self::Network(why) => write!(f, "the download failed: {why}"),
120            Self::Disk(why) => write!(f, "the model could not be written: {why}"),
121            Self::WrongFile => f.write_str("the file is not the published model"),
122            Self::Cancelled => f.write_str("the download was cancelled"),
123        }
124    }
125}
126
127impl std::error::Error for ModelError {}
128
129fn disk(error: std::io::Error) -> ModelError {
130    ModelError::Disk(error.to_string())
131}
132
133/// A folder that holds models, each in a folder of its own.
134///
135/// ```no_run
136/// use std::sync::atomic::AtomicBool;
137/// use ravenrs_ai::{LocalModel, ModelFamily, ModelFile, ModelFolder};
138///
139/// const MODEL: LocalModel = LocalModel {
140///     id: "tiny",
141///     name: "Tiny",
142///     license: "MIT",
143///     family: ModelFamily::Qwen3,
144///     memory: 2_000_000_000,
145///     understood: None,
146///     files: &[ModelFile {
147///         name: "tiny.gguf",
148///         url: "https://example.invalid/tiny.gguf",
149///         bytes: 1024,
150///         sha256: "0000000000000000000000000000000000000000000000000000000000000000",
151///     }],
152/// };
153/// let folder = ModelFolder::new("models");
154/// folder.download(&MODEL, &|received, total| println!("{received} of {total}"), &AtomicBool::new(false))?;
155/// folder.remove(&MODEL)?;
156/// # Ok::<(), ravenrs_ai::ModelError>(())
157/// ```
158#[derive(Clone, Debug, PartialEq, Eq)]
159pub struct ModelFolder {
160    root: PathBuf,
161}
162
163impl ModelFolder {
164    /// A folder at `root`. It is created when something is downloaded.
165    pub fn new(root: impl Into<PathBuf>) -> Self {
166        Self { root: root.into() }
167    }
168
169    /// Where it is.
170    pub fn root(&self) -> &Path {
171        &self.root
172    }
173
174    fn folder_of(&self, model: &LocalModel) -> PathBuf {
175        self.root.join(model.id)
176    }
177
178    /// Where a file of `model` is, or would be.
179    pub fn path_of(&self, model: &LocalModel, file: &ModelFile) -> PathBuf {
180        self.folder_of(model).join(file.name)
181    }
182
183    /// Where `model` stands. Quick: it compares sizes, as the hash was checked when
184    /// it arrived.
185    pub fn state(&self, model: &LocalModel) -> ModelState {
186        let whole = model
187            .files
188            .iter()
189            .all(|file| fs::metadata(self.path_of(model, file)).is_ok_and(|found| found.len() == file.bytes));
190        if whole { ModelState::Installed } else { ModelState::Absent }
191    }
192
193    /// Downloads what is missing of `model` and installs each file only if its hash
194    /// is right. It blocks until the end, so call it away from the UI thread.
195    ///
196    /// `progress` gets the bytes received and the total, and `cancelled` is looked
197    /// at between chunks.
198    ///
199    /// # Errors
200    ///
201    /// Why it did not arrive. Files already checked stay, so trying again goes on
202    /// from there, and the one being downloaded is erased.
203    pub fn download(
204        &self,
205        model: &LocalModel,
206        progress: &dyn Fn(u64, u64),
207        cancelled: &AtomicBool,
208    ) -> Result<(), ModelError> {
209        fs::create_dir_all(self.folder_of(model)).map_err(disk)?;
210        let total = model.bytes();
211        let mut done = 0;
212        for file in model.files {
213            let path = self.path_of(model, file);
214            if fs::metadata(&path).is_ok_and(|found| found.len() == file.bytes) {
215                done += file.bytes;
216                progress(done, total);
217                continue;
218            }
219            let part = path.with_extension(PART);
220            let result = fetch(file, &part, done, total, progress, cancelled);
221            if result.is_err() {
222                // Whatever it was, there is nothing to keep.
223                let _ = fs::remove_file(&part);
224            }
225            result?;
226            fs::rename(&part, &path).map_err(disk)?;
227            done += file.bytes;
228        }
229        Ok(())
230    }
231
232    /// Removes everything of `model`, including a download left halfway.
233    ///
234    /// # Errors
235    ///
236    /// When the disk refuses. A model that is not there is not an error.
237    pub fn remove(&self, model: &LocalModel) -> Result<(), ModelError> {
238        match fs::remove_dir_all(self.folder_of(model)) {
239            Err(error) if error.kind() != std::io::ErrorKind::NotFound => Err(disk(error)),
240            _ => Ok(()),
241        }
242    }
243}
244
245/// Downloads one file to `part`, checking its hash.
246fn fetch(
247    file: &ModelFile,
248    part: &Path,
249    before: u64,
250    total: u64,
251    progress: &dyn Fn(u64, u64),
252    cancelled: &AtomicBool,
253) -> Result<(), ModelError> {
254    let response = ureq::get(file.url).call().map_err(|error| ModelError::Network(error.to_string()))?;
255    let mut reader = response.into_body().into_reader();
256    let mut out = File::create(part).map_err(disk)?;
257    let mut hash = Sha256::new();
258    let mut buffer = vec![0_u8; CHUNK];
259    let mut received = 0_u64;
260    loop {
261        if cancelled.load(Ordering::Relaxed) {
262            return Err(ModelError::Cancelled);
263        }
264        let read = reader.read(&mut buffer).map_err(|error| ModelError::Network(error.to_string()))?;
265        if read == 0 {
266            break;
267        }
268        let chunk = buffer.get(..read).unwrap_or_default();
269        hash.update(chunk);
270        out.write_all(chunk).map_err(disk)?;
271        received += read as u64;
272        progress(before + received, total);
273    }
274    out.flush().map_err(disk)?;
275    let found: String = hash.finalize().iter().map(|byte| format!("{byte:02x}")).collect();
276    if found == file.sha256 && received == file.bytes { Ok(()) } else { Err(ModelError::WrongFile) }
277}
ravenrs_ai/models.rs

ravenrs_ai/
models.rs
