use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ModelSource {
HuggingFace {
repo_id: String,
},
LocalFile {
path: PathBuf,
},
Custom {
url: String,
},
}
impl ModelSource {
pub fn huggingface(repo_id: impl Into<String>) -> Self {
Self::HuggingFace { repo_id: repo_id.into() }
}
pub fn local(path: impl Into<PathBuf>) -> Self {
Self::LocalFile { path: path.into() }
}
pub fn custom(url: impl Into<String>) -> Self {
Self::Custom { url: url.into() }
}
pub fn display_string(&self) -> String {
match self {
Self::HuggingFace { repo_id } => format!("hf://{repo_id}"),
Self::LocalFile { path } => format!("file://{}", path.display()),
Self::Custom { url } => url.clone(),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelEntry {
pub name: String,
pub version: String,
pub sha256: String,
pub size_bytes: u64,
pub source: ModelSource,
pub local_path: Option<PathBuf>,
pub format: Option<String>,
pub metadata: HashMap<String, String>,
}
impl ModelEntry {
pub fn new(
name: impl Into<String>,
version: impl Into<String>,
sha256: impl Into<String>,
size_bytes: u64,
source: ModelSource,
) -> Self {
Self {
name: name.into(),
version: version.into(),
sha256: sha256.into(),
size_bytes,
source,
local_path: None,
format: None,
metadata: HashMap::new(),
}
}
pub fn with_local_path(mut self, path: impl Into<PathBuf>) -> Self {
self.local_path = Some(path.into());
self
}
pub fn with_format(mut self, format: impl Into<String>) -> Self {
self.format = Some(format.into());
self
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
pub fn is_local(&self) -> bool {
self.local_path.as_ref().is_some_and(|p| p.exists())
}
pub fn size_mb(&self) -> f64 {
self.size_bytes as f64 / (1024.0 * 1024.0)
}
pub fn size_gb(&self) -> f64 {
self.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
}
}