467 lines
14 KiB
Rust
467 lines
14 KiB
Rust
use std::fmt;
|
|
use std::io;
|
|
use std::path::{Path, PathBuf};
|
|
use std::str::FromStr;
|
|
use std::sync::{LazyLock, Mutex};
|
|
|
|
use futures_util::StreamExt;
|
|
use serde::{Deserialize, Serialize};
|
|
use sha2::{Digest, Sha256};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
#[allow(non_camel_case_types)]
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum LlmModelId {
|
|
#[serde(rename = "qwen3_1_7b")]
|
|
Qwen3_1_7B_Q4,
|
|
#[serde(rename = "qwen3_4b_instruct_2507")]
|
|
Qwen3_4BInstruct2507Q4,
|
|
#[serde(rename = "qwen3_14b")]
|
|
Qwen3_14BQ5,
|
|
}
|
|
|
|
impl LlmModelId {
|
|
pub fn default_tier() -> Self {
|
|
Self::Qwen3_4BInstruct2507Q4
|
|
}
|
|
|
|
pub fn as_str(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => "qwen3_1_7b",
|
|
Self::Qwen3_4BInstruct2507Q4 => "qwen3_4b_instruct_2507",
|
|
Self::Qwen3_14BQ5 => "qwen3_14b",
|
|
}
|
|
}
|
|
|
|
pub fn display_name(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => "Qwen3 1.7B",
|
|
Self::Qwen3_4BInstruct2507Q4 => "Qwen3 4B Instruct 2507",
|
|
Self::Qwen3_14BQ5 => "Qwen3 14B",
|
|
}
|
|
}
|
|
|
|
pub fn file_name(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => "Qwen3-1.7B-Q4_K_M.gguf",
|
|
Self::Qwen3_4BInstruct2507Q4 => "Qwen3-4B-Instruct-2507-Q4_K_M.gguf",
|
|
Self::Qwen3_14BQ5 => "Qwen3-14B-Q5_K_M.gguf",
|
|
}
|
|
}
|
|
|
|
pub fn size_bytes(&self) -> u64 {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => 1_107_409_472,
|
|
Self::Qwen3_4BInstruct2507Q4 => 2_497_281_120,
|
|
Self::Qwen3_14BQ5 => 10_514_570_624,
|
|
}
|
|
}
|
|
|
|
pub fn minimum_ram_bytes(&self) -> u64 {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => 8 * 1024_u64.pow(3),
|
|
Self::Qwen3_4BInstruct2507Q4 => 16 * 1024_u64.pow(3),
|
|
Self::Qwen3_14BQ5 => 32 * 1024_u64.pow(3),
|
|
}
|
|
}
|
|
|
|
pub fn recommended_vram_bytes(&self) -> Option<u64> {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => None,
|
|
Self::Qwen3_4BInstruct2507Q4 => Some(8 * 1024_u64.pow(3)),
|
|
Self::Qwen3_14BQ5 => Some(16 * 1024_u64.pow(3)),
|
|
}
|
|
}
|
|
|
|
pub fn description(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => "Low tier for 8 GB RAM and CPU-heavy machines.",
|
|
Self::Qwen3_4BInstruct2507Q4 => {
|
|
"Default tier for cleanup and task extraction on 16 GB systems."
|
|
}
|
|
Self::Qwen3_14BQ5 => "High tier for 32 GB+ RAM and larger GPUs.",
|
|
}
|
|
}
|
|
|
|
pub fn hf_url(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => {
|
|
"https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/d7f544eead698dbd1f15126ef60b45a1e1933222/Qwen3-1.7B-Q4_K_M.gguf"
|
|
}
|
|
Self::Qwen3_4BInstruct2507Q4 => {
|
|
"https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/resolve/a06e946bb6b655725eafa393f4a9745d460374c9/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"
|
|
}
|
|
Self::Qwen3_14BQ5 => {
|
|
"https://huggingface.co/unsloth/Qwen3-14B-GGUF/resolve/a04a82c4739b3ef5fa6da7d10261db2c67dd1985/Qwen3-14B-Q5_K_M.gguf"
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn sha256(&self) -> &'static str {
|
|
match self {
|
|
Self::Qwen3_1_7B_Q4 => {
|
|
"de942b0819216caa3bfe487180dd1bb37398fa1c98cb42bb0bbac7ab7d6e8a12"
|
|
}
|
|
Self::Qwen3_4BInstruct2507Q4 => {
|
|
"bf52d44a54b81d44219833556849529ee96f09da673a38783dddc2e2eaf17881"
|
|
}
|
|
Self::Qwen3_14BQ5 => "6f87abc471bd509ad46aca4284b3cfa926d8114bc491bb0a7a3a7f74c16ef95b",
|
|
}
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for LlmModelId {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.write_str(self.as_str())
|
|
}
|
|
}
|
|
|
|
impl FromStr for LlmModelId {
|
|
type Err = String;
|
|
|
|
fn from_str(value: &str) -> Result<Self, Self::Err> {
|
|
match value {
|
|
"qwen3_1_7b" => Ok(Self::Qwen3_1_7B_Q4),
|
|
"qwen3_4b_instruct_2507" => Ok(Self::Qwen3_4BInstruct2507Q4),
|
|
"qwen3_14b" => Ok(Self::Qwen3_14BQ5),
|
|
other => Err(format!("Unknown LLM model id: {other}")),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct LlmModelInfo {
|
|
pub id: String,
|
|
pub display_name: &'static str,
|
|
pub file_name: &'static str,
|
|
pub size_bytes: u64,
|
|
pub description: &'static str,
|
|
pub minimum_ram_bytes: u64,
|
|
pub recommended_vram_bytes: Option<u64>,
|
|
}
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum DownloadError {
|
|
#[error("http error: {0}")]
|
|
Http(String),
|
|
#[error("io error: {0}")]
|
|
Io(#[from] io::Error),
|
|
#[error("sha256 mismatch: expected {expected}, got {actual}")]
|
|
ShaMismatch { expected: String, actual: String },
|
|
#[error("resume failed: server does not support range requests")]
|
|
ResumeUnsupported,
|
|
}
|
|
|
|
const ALL_MODELS: &[LlmModelId] = &[
|
|
LlmModelId::Qwen3_1_7B_Q4,
|
|
LlmModelId::Qwen3_4BInstruct2507Q4,
|
|
LlmModelId::Qwen3_14BQ5,
|
|
];
|
|
|
|
static ACTIVE_DOWNLOADS: LazyLock<Mutex<std::collections::HashSet<LlmModelId>>> =
|
|
LazyLock::new(|| Mutex::new(std::collections::HashSet::new()));
|
|
|
|
struct DownloadReservation {
|
|
id: LlmModelId,
|
|
}
|
|
|
|
impl DownloadReservation {
|
|
fn acquire(id: LlmModelId) -> Result<Self, DownloadError> {
|
|
let mut active = ACTIVE_DOWNLOADS
|
|
.lock()
|
|
.map_err(|_| DownloadError::Http("download lock poisoned".into()))?;
|
|
if !active.insert(id) {
|
|
return Err(DownloadError::Http(format!(
|
|
"download already in progress for {}",
|
|
id.as_str()
|
|
)));
|
|
}
|
|
Ok(Self { id })
|
|
}
|
|
}
|
|
|
|
impl Drop for DownloadReservation {
|
|
fn drop(&mut self) {
|
|
if let Ok(mut active) = ACTIVE_DOWNLOADS.lock() {
|
|
active.remove(&self.id);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn all_models() -> &'static [LlmModelId] {
|
|
ALL_MODELS
|
|
}
|
|
|
|
pub fn model_info(id: LlmModelId) -> LlmModelInfo {
|
|
LlmModelInfo {
|
|
id: id.as_str().to_string(),
|
|
display_name: id.display_name(),
|
|
file_name: id.file_name(),
|
|
size_bytes: id.size_bytes(),
|
|
description: id.description(),
|
|
minimum_ram_bytes: id.minimum_ram_bytes(),
|
|
recommended_vram_bytes: id.recommended_vram_bytes(),
|
|
}
|
|
}
|
|
|
|
pub fn recommend_tier(total_ram_bytes: u64, total_vram_bytes: Option<u64>) -> LlmModelId {
|
|
if total_vram_bytes.unwrap_or(0) >= 16 * 1024_u64.pow(3)
|
|
&& total_ram_bytes >= 32 * 1024_u64.pow(3)
|
|
{
|
|
LlmModelId::Qwen3_14BQ5
|
|
} else if total_vram_bytes.unwrap_or(0) >= 8 * 1024_u64.pow(3)
|
|
|| total_ram_bytes >= 16 * 1024_u64.pow(3)
|
|
{
|
|
LlmModelId::Qwen3_4BInstruct2507Q4
|
|
} else {
|
|
LlmModelId::Qwen3_1_7B_Q4
|
|
}
|
|
}
|
|
|
|
pub fn model_dir() -> PathBuf {
|
|
kon_core::paths::app_paths().llm_models_dir()
|
|
}
|
|
|
|
pub fn model_path(id: LlmModelId) -> PathBuf {
|
|
model_dir().join(id.file_name())
|
|
}
|
|
|
|
pub fn partial_download_path(id: LlmModelId) -> PathBuf {
|
|
model_path(id).with_extension("gguf.part")
|
|
}
|
|
|
|
pub fn is_downloaded(id: LlmModelId) -> bool {
|
|
model_path(id).exists()
|
|
}
|
|
|
|
pub fn delete_model(id: LlmModelId) -> io::Result<()> {
|
|
let final_path = model_path(id);
|
|
let partial_path = partial_download_path(id);
|
|
|
|
if final_path.exists() {
|
|
std::fs::remove_file(final_path)?;
|
|
}
|
|
if partial_path.exists() {
|
|
std::fs::remove_file(partial_path)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn download_model<F>(id: LlmModelId, on_progress: F) -> Result<(), DownloadError>
|
|
where
|
|
F: FnMut(u64, u64) + Send + 'static,
|
|
{
|
|
let _reservation = DownloadReservation::acquire(id)?;
|
|
let dest = model_path(id);
|
|
tokio::fs::create_dir_all(model_dir()).await?;
|
|
|
|
if dest.exists() {
|
|
let actual = sha256_file(&dest).await?;
|
|
if actual == id.sha256() {
|
|
return Ok(());
|
|
}
|
|
tokio::fs::remove_file(&dest).await?;
|
|
}
|
|
|
|
download_impl(id.hf_url(), id.sha256(), &dest, on_progress).await
|
|
}
|
|
|
|
async fn sha256_file(path: &Path) -> Result<String, io::Error> {
|
|
let mut hasher = Sha256::new();
|
|
let mut file = tokio::fs::File::open(path).await?;
|
|
let mut buffer = [0u8; 8192];
|
|
|
|
loop {
|
|
let count = file.read(&mut buffer).await?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
hasher.update(&buffer[..count]);
|
|
}
|
|
|
|
Ok(format!("{:x}", hasher.finalize()))
|
|
}
|
|
|
|
async fn download_impl<F>(
|
|
url: &str,
|
|
expected_sha: &str,
|
|
dest: &Path,
|
|
mut on_progress: F,
|
|
) -> Result<(), DownloadError>
|
|
where
|
|
F: FnMut(u64, u64) + Send + 'static,
|
|
{
|
|
let tmp = dest.with_extension("gguf.part");
|
|
let resume_from = tokio::fs::metadata(&tmp)
|
|
.await
|
|
.ok()
|
|
.map(|m| m.len())
|
|
.unwrap_or(0);
|
|
|
|
let client = reqwest::Client::builder()
|
|
.user_agent("kon/0.1.0")
|
|
.connect_timeout(std::time::Duration::from_secs(30))
|
|
.build()
|
|
.map_err(|e| DownloadError::Http(e.to_string()))?;
|
|
|
|
let mut request = client.get(url);
|
|
if resume_from > 0 {
|
|
request = request.header(reqwest::header::RANGE, format!("bytes={resume_from}-"));
|
|
}
|
|
|
|
let response = request
|
|
.send()
|
|
.await
|
|
.map_err(|e| DownloadError::Http(e.to_string()))?;
|
|
if resume_from > 0 && response.status() != reqwest::StatusCode::PARTIAL_CONTENT {
|
|
return Err(DownloadError::ResumeUnsupported);
|
|
}
|
|
if !response.status().is_success() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT
|
|
{
|
|
return Err(DownloadError::Http(format!("status {}", response.status())));
|
|
}
|
|
|
|
let total = if resume_from > 0 {
|
|
response
|
|
.headers()
|
|
.get(reqwest::header::CONTENT_RANGE)
|
|
.and_then(|value| value.to_str().ok())
|
|
.and_then(|value| value.rsplit('/').next())
|
|
.and_then(|value| value.parse::<u64>().ok())
|
|
.unwrap_or_else(|| response.content_length().unwrap_or(0) + resume_from)
|
|
} else {
|
|
response.content_length().unwrap_or(0)
|
|
};
|
|
|
|
let mut hasher = Sha256::new();
|
|
if resume_from > 0 {
|
|
let mut partial = tokio::fs::File::open(&tmp).await?;
|
|
let mut buffer = [0u8; 8192];
|
|
loop {
|
|
let count = partial.read(&mut buffer).await?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
hasher.update(&buffer[..count]);
|
|
}
|
|
}
|
|
|
|
let mut output = tokio::fs::OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.open(&tmp)
|
|
.await?;
|
|
|
|
let mut downloaded = resume_from;
|
|
let mut stream = response.bytes_stream();
|
|
while let Some(chunk) = stream.next().await {
|
|
let chunk = chunk.map_err(|e| DownloadError::Http(e.to_string()))?;
|
|
output.write_all(&chunk).await?;
|
|
hasher.update(&chunk);
|
|
downloaded += chunk.len() as u64;
|
|
on_progress(downloaded, total);
|
|
}
|
|
output.flush().await?;
|
|
drop(output);
|
|
|
|
let actual = format!("{:x}", hasher.finalize());
|
|
if actual != expected_sha {
|
|
tokio::fs::remove_file(&tmp).await.ok();
|
|
return Err(DownloadError::ShaMismatch {
|
|
expected: expected_sha.to_string(),
|
|
actual,
|
|
});
|
|
}
|
|
|
|
tokio::fs::rename(&tmp, dest).await?;
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::sync::{Arc, Mutex};
|
|
use tempfile::tempdir;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::TcpListener;
|
|
|
|
#[test]
|
|
fn model_path_contains_model_dir_and_filename() {
|
|
let path = model_path(LlmModelId::Qwen3_1_7B_Q4);
|
|
assert!(path.to_string_lossy().ends_with("Qwen3-1.7B-Q4_K_M.gguf"));
|
|
assert!(path.starts_with(model_dir()));
|
|
}
|
|
|
|
#[test]
|
|
fn recommend_tier_prefers_mid_by_default() {
|
|
let tier = recommend_tier(16 * 1024_u64.pow(3), None);
|
|
assert_eq!(tier, LlmModelId::Qwen3_4BInstruct2507Q4);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn download_impl_supports_resume_and_sha_verification() {
|
|
let fixture = b"hello resumed download".to_vec();
|
|
let expected_sha = format!("{:x}", Sha256::digest(&fixture));
|
|
let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = server.local_addr().unwrap();
|
|
let content = fixture.clone();
|
|
|
|
let server_task = tokio::spawn(async move {
|
|
let (mut socket, _) = server.accept().await.unwrap();
|
|
let mut request = vec![0u8; 2048];
|
|
let size = socket.read(&mut request).await.unwrap();
|
|
let request = String::from_utf8_lossy(&request[..size]).to_lowercase();
|
|
let range_start = request
|
|
.lines()
|
|
.find_map(|line| line.strip_prefix("range: bytes="))
|
|
.and_then(|line| line.strip_suffix('-'))
|
|
.and_then(|line| line.trim().parse::<usize>().ok());
|
|
|
|
if let Some(start) = range_start {
|
|
let body = &content[start..];
|
|
let response = format!(
|
|
"HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {}-{}/{}\r\nAccept-Ranges: bytes\r\n\r\n",
|
|
body.len(),
|
|
start,
|
|
content.len() - 1,
|
|
content.len()
|
|
);
|
|
socket.write_all(response.as_bytes()).await.unwrap();
|
|
socket.write_all(body).await.unwrap();
|
|
} else {
|
|
let response = format!(
|
|
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\n\r\n",
|
|
content.len()
|
|
);
|
|
socket.write_all(response.as_bytes()).await.unwrap();
|
|
socket.write_all(&content).await.unwrap();
|
|
}
|
|
});
|
|
|
|
let dir = tempdir().unwrap();
|
|
let dest = dir.path().join("fixture.gguf");
|
|
let part = dest.with_extension("gguf.part");
|
|
tokio::fs::write(&part, &fixture[..10]).await.unwrap();
|
|
|
|
let progress = Arc::new(Mutex::new(Vec::new()));
|
|
let progress_clone = progress.clone();
|
|
download_impl(
|
|
&format!("http://{addr}/fixture.gguf"),
|
|
&expected_sha,
|
|
&dest,
|
|
move |done, total| progress_clone.lock().unwrap().push((done, total)),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let saved = tokio::fs::read(&dest).await.unwrap();
|
|
assert_eq!(saved, fixture);
|
|
assert!(!part.exists());
|
|
assert!(!progress.lock().unwrap().is_empty());
|
|
|
|
server_task.await.unwrap();
|
|
}
|
|
}
|