feat: model dir auto-scan, remote API models, model plaza with HF/ModelScope downloads
This commit is contained in:
Xianren Studio
@@ -77,5 +77,24 @@ fn open_db(path: &Path) -> Result<Connection> {
|
||||
conn.pragma_update(None, "journal_mode", "WAL")?;
|
||||
conn.pragma_update(None, "foreign_keys", "ON")?;
|
||||
conn.execute_batch(include_str!("schema.sql"))?;
|
||||
// 旧库迁移：为已存在的 models 表补充新增列
|
||||
ensure_column(&conn, "models", "kind", "TEXT NOT NULL DEFAULT 'local'")?;
|
||||
ensure_column(&conn, "models", "base_url", "TEXT")?;
|
||||
ensure_column(&conn, "models", "api_key", "TEXT")?;
|
||||
ensure_column(&conn, "models", "api_model", "TEXT")?;
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
fn ensure_column(conn: &Connection, table: &str, column: &str, decl: &str) -> rusqlite::Result<()> {
|
||||
let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
|
||||
let columns: Vec<String> = stmt
|
||||
.query_map([], |row| row.get::<_, String>(1))?
|
||||
.collect::<std::result::Result<_, _>>()?;
|
||||
if !columns.iter().any(|c| c == column) {
|
||||
conn.execute(
|
||||
&format!("ALTER TABLE {table} ADD COLUMN {column} {decl}"),
|
||||
[],
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
+138
-13
@@ -1,12 +1,15 @@
|
||||
use crate::error::Result;
|
||||
use rusqlite::{params, Connection};
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ModelInfo {
|
||||
pub id: String,
|
||||
pub repo_id: String,
|
||||
pub source: String,
|
||||
pub kind: String,
|
||||
pub file_name: String,
|
||||
pub file_path: String,
|
||||
pub file_size: i64,
|
||||
@@ -14,36 +17,54 @@ pub struct ModelInfo {
|
||||
pub family: Option<String>,
|
||||
pub status: String,
|
||||
pub sha256: Option<String>,
|
||||
pub base_url: Option<String>,
|
||||
pub api_key: Option<String>,
|
||||
pub api_model: Option<String>,
|
||||
pub meta: serde_json::Value,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Serialize)]
|
||||
pub struct ScanResult {
|
||||
pub added: usize,
|
||||
pub updated: usize,
|
||||
pub missing: usize,
|
||||
}
|
||||
|
||||
pub fn insert(
|
||||
db: &Connection,
|
||||
repo_id: &str,
|
||||
source: &str,
|
||||
kind: &str,
|
||||
file_name: &str,
|
||||
file_path: &str,
|
||||
file_size: i64,
|
||||
quant: Option<&str>,
|
||||
family: Option<&str>,
|
||||
sha256: Option<&str>,
|
||||
base_url: Option<&str>,
|
||||
api_key: Option<&str>,
|
||||
api_model: Option<&str>,
|
||||
meta: serde_json::Value,
|
||||
) -> Result<String> {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
db.execute(
|
||||
"INSERT INTO models (id, repo_id, source, file_name, file_path, file_size, quant, family, status, sha256, meta_json)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, 'ready', ?9, ?10)",
|
||||
"INSERT INTO models (id, repo_id, source, kind, file_name, file_path, file_size, quant, family, status, sha256, base_url, api_key, api_model, meta_json)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, 'ready', ?10, ?11, ?12, ?13, ?14)",
|
||||
params![
|
||||
id,
|
||||
repo_id,
|
||||
source,
|
||||
kind,
|
||||
file_name,
|
||||
file_path,
|
||||
file_size,
|
||||
quant,
|
||||
family,
|
||||
sha256,
|
||||
base_url,
|
||||
api_key,
|
||||
api_model,
|
||||
meta.to_string()
|
||||
],
|
||||
)?;
|
||||
@@ -52,7 +73,7 @@ pub fn insert(
|
||||
|
||||
pub fn list(db: &Connection) -> Result<Vec<ModelInfo>> {
|
||||
let mut stmt = db.prepare(
|
||||
"SELECT id, repo_id, source, file_name, file_path, file_size, quant, family, status, sha256, meta_json, created_at
|
||||
"SELECT id, repo_id, source, kind, file_name, file_path, file_size, quant, family, status, sha256, base_url, api_key, api_model, meta_json, created_at
|
||||
FROM models ORDER BY created_at DESC",
|
||||
)?;
|
||||
let rows = stmt.query_map([], row_to_model)?;
|
||||
@@ -65,7 +86,7 @@ pub fn list(db: &Connection) -> Result<Vec<ModelInfo>> {
|
||||
|
||||
pub fn get(db: &Connection, id: &str) -> Result<Option<ModelInfo>> {
|
||||
let mut stmt = db.prepare(
|
||||
"SELECT id, repo_id, source, file_name, file_path, file_size, quant, family, status, sha256, meta_json, created_at
|
||||
"SELECT id, repo_id, source, kind, file_name, file_path, file_size, quant, family, status, sha256, base_url, api_key, api_model, meta_json, created_at
|
||||
FROM models WHERE id = ?1",
|
||||
)?;
|
||||
let mut rows = stmt.query_map(params![id], row_to_model)?;
|
||||
@@ -80,21 +101,125 @@ pub fn remove(db: &Connection, id: &str) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn update_file_info(db: &Connection, id: &str, size: i64, status: &str) -> Result<()> {
|
||||
db.execute(
|
||||
"UPDATE models SET file_size = ?1, status = ?2 WHERE id = ?3",
|
||||
params![size, status, id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 扫描目录（含子目录）中的 .gguf 文件，与注册表同步。
|
||||
pub fn scan_directory(db: &Connection, dir: &Path) -> Result<ScanResult> {
|
||||
if !dir.exists() {
|
||||
return Ok(ScanResult::default());
|
||||
}
|
||||
let mut found = Vec::new();
|
||||
collect_gguf_files(dir, &mut found)?;
|
||||
|
||||
let existing = list(db)?;
|
||||
let mut by_path: HashMap<String, ModelInfo> = HashMap::new();
|
||||
for m in existing {
|
||||
by_path.insert(m.file_path.to_lowercase(), m);
|
||||
}
|
||||
|
||||
let mut result = ScanResult::default();
|
||||
for path in found {
|
||||
let meta = std::fs::metadata(&path)?;
|
||||
let name = path
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
let path_str = path.to_string_lossy().to_string();
|
||||
let key = path_str.to_lowercase();
|
||||
let size = meta.len() as i64;
|
||||
|
||||
if let Some(m) = by_path.remove(&key) {
|
||||
if m.file_size != size || m.status != "ready" {
|
||||
update_file_info(db, &m.id, size, "ready")?;
|
||||
result.updated += 1;
|
||||
}
|
||||
} else {
|
||||
insert(
|
||||
db,
|
||||
&name,
|
||||
"scan",
|
||||
"local",
|
||||
&name,
|
||||
&path_str,
|
||||
size,
|
||||
guess_quant(&name).as_deref(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
serde_json::json!({}),
|
||||
)?;
|
||||
result.added += 1;
|
||||
}
|
||||
}
|
||||
|
||||
for m in by_path.values() {
|
||||
if m.kind == "local" && m.status != "missing" {
|
||||
update_file_info(db, &m.id, m.file_size, "missing")?;
|
||||
result.missing += 1;
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn collect_gguf_files(dir: &Path, out: &mut Vec<PathBuf>) -> Result<()> {
|
||||
let entries = std::fs::read_dir(dir)?;
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
let file_type = entry.file_type()?;
|
||||
if file_type.is_dir() {
|
||||
collect_gguf_files(&path, out)?;
|
||||
} else if file_type.is_file() {
|
||||
let is_gguf = path
|
||||
.extension()
|
||||
.map(|e| e.eq_ignore_ascii_case("gguf"))
|
||||
.unwrap_or(false);
|
||||
let is_part = path.to_string_lossy().contains(".part");
|
||||
if is_gguf && !is_part {
|
||||
out.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn row_to_model(row: &rusqlite::Row<'_>) -> rusqlite::Result<ModelInfo> {
|
||||
let meta_json: String = row.get(10)?;
|
||||
let meta_json: String = row.get(14)?;
|
||||
Ok(ModelInfo {
|
||||
id: row.get(0)?,
|
||||
repo_id: row.get(1)?,
|
||||
source: row.get(2)?,
|
||||
file_name: row.get(3)?,
|
||||
file_path: row.get(4)?,
|
||||
file_size: row.get(5)?,
|
||||
quant: row.get(6)?,
|
||||
family: row.get(7)?,
|
||||
status: row.get(8)?,
|
||||
sha256: row.get(9)?,
|
||||
kind: row.get(3)?,
|
||||
file_name: row.get(4)?,
|
||||
file_path: row.get(5)?,
|
||||
file_size: row.get(6)?,
|
||||
quant: row.get(7)?,
|
||||
family: row.get(8)?,
|
||||
status: row.get(9)?,
|
||||
sha256: row.get(10)?,
|
||||
base_url: row.get(11)?,
|
||||
api_key: row.get(12)?,
|
||||
api_model: row.get(13)?,
|
||||
meta: serde_json::from_str(&meta_json).unwrap_or_else(|_| serde_json::json!({})),
|
||||
created_at: row.get(11)?,
|
||||
created_at: row.get(15)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn guess_quant(file_name: &str) -> Option<String> {
|
||||
const QUANTS: &[&str] = &[
|
||||
"q4_k_m", "q5_k_m", "q6_k", "q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "q2_k", "q3_k_m",
|
||||
"f16", "f32",
|
||||
];
|
||||
let lower = file_name.to_lowercase();
|
||||
QUANTS
|
||||
.iter()
|
||||
.find(|q| lower.contains(**q))
|
||||
.map(|s| s.to_string())
|
||||
}
|
||||
@@ -2,6 +2,7 @@ CREATE TABLE IF NOT EXISTS models (
|
||||
id TEXT PRIMARY KEY,
|
||||
repo_id TEXT NOT NULL,
|
||||
source TEXT NOT NULL DEFAULT 'local',
|
||||
kind TEXT NOT NULL DEFAULT 'local',
|
||||
file_name TEXT NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
file_size INTEGER NOT NULL DEFAULT 0,
|
||||
@@ -9,6 +10,9 @@ CREATE TABLE IF NOT EXISTS models (
|
||||
family TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'ready',
|
||||
sha256 TEXT,
|
||||
base_url TEXT,
|
||||
api_key TEXT,
|
||||
api_model TEXT,
|
||||
meta_json TEXT NOT NULL DEFAULT '{}',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
@@ -36,4 +40,3 @@ CREATE TABLE IF NOT EXISTS messages (
|
||||
tokens_out INTEGER,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user