Skip to content

Commit 3b0b8ad

Browse files
kerthcetkerthcet
committed
optimize the modelInfo
Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent ff7a96d commit 3b0b8ad

8 files changed

Lines changed: 49 additions & 49 deletions

File tree

‎src/api/tests.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use tower::util::ServiceExt; // for `oneshot` and `ready`
1414

1515
use super::routes::create_router;
1616
use crate::backend::mock::MockEngine;
17-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry};
17+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
1818

1919
/// Helper to create test app with a pre-registered test model
2020
/// Returns the router and the temp directory (which must be kept alive)
@@ -27,15 +27,15 @@ fn create_test_app() -> (axum::Router, TempDir) {
2727
let test_model = ModelInfo {
2828
uuid: "test-uuid".to_string(),
2929
name: "test-model".to_string(),
30+
provider: "test".to_string(),
3031
author: Some("test-author".to_string()),
3132
task: Some("text-generation".to_string()),
3233
model_series: Some("test-series".to_string()),
33-
provider: "test".to_string(),
3434
license: Some("MIT".to_string()),
3535
created_at: chrono::Utc::now().to_rfc3339(),
3636
updated_at: chrono::Utc::now().to_rfc3339(),
3737
metadata: ModelMetadata {
38-
artifact: ArtifactInfo {
38+
cache: CacheInfo {
3939
revision: "test-rev".to_string(),
4040
size: 1000,
4141
path: "/tmp/test-model".to_string(),

‎src/cli/commands.rs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -141,12 +141,12 @@ pub async fn run(cli: Cli) {
141141
"MODEL", "TASK", "PROVIDER", "REVISION", "SIZE", "CREATED"
142142
]);
143143
for model in models {
144-
let size_str = format_size_decimal(model.metadata.artifact.size);
144+
let size_str = format_size_decimal(model.metadata.cache.size);
145145

146-
let revision_short = if model.metadata.artifact.revision.len() > 8 {
147-
&model.metadata.artifact.revision[..8]
146+
let revision_short = if model.metadata.cache.revision.len() > 8 {
147+
&model.metadata.cache.revision[..8]
148148
} else {
149-
&model.metadata.artifact.revision
149+
&model.metadata.cache.revision
150150
};
151151

152152
let created_str = format_time_ago(&model.created_at);
@@ -230,7 +230,7 @@ pub async fn run(cli: Cli) {
230230
#[cfg(test)]
231231
mod tests {
232232
use super::*;
233-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
233+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
234234
use tempfile::TempDir;
235235

236236
// Helper to create a test model
@@ -245,15 +245,15 @@ mod tests {
245245
ModelInfo {
246246
uuid: revision.to_string(),
247247
name: name.to_string(),
248+
provider: "huggingface".to_string(),
248249
author: Some("test-author".to_string()),
249250
task: Some("text-generation".to_string()),
250251
model_series: Some("gpt2".to_string()),
251-
provider: "huggingface".to_string(),
252252
license: Some("mit".to_string()),
253253
created_at: "2025-01-01T00:00:00Z".to_string(),
254254
updated_at: "2025-01-01T00:00:00Z".to_string(),
255255
metadata: ModelMetadata {
256-
artifact: ArtifactInfo {
256+
cache: CacheInfo {
257257
revision: revision.to_string(),
258258
size: 1000,
259259
path: "/tmp/test".to_string(),
@@ -375,7 +375,7 @@ mod tests {
375375

376376
// Update the model
377377
let mut updated_model = create_test_model("test/updated-model", "v2");
378-
updated_model.metadata.artifact.size = 2000;
378+
updated_model.metadata.cache.size = 2000;
379379
updated_model.created_at = "2025-01-05T00:00:00Z".to_string();
380380
updated_model.updated_at = "2025-01-05T00:00:00Z".to_string();
381381

@@ -387,7 +387,7 @@ mod tests {
387387
// updated_at should be new
388388
assert_eq!(result.updated_at, "2025-01-05T00:00:00Z");
389389
// Other fields should be updated
390-
assert_eq!(result.metadata.artifact.revision, "v2");
391-
assert_eq!(result.metadata.artifact.size, 2000);
390+
assert_eq!(result.metadata.cache.revision, "v2");
391+
assert_eq!(result.metadata.cache.size, 2000);
392392
}
393393
}

‎src/cli/inspect.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,12 @@ pub fn display(model: &ModelInfo) {
6868
// Artifact section
6969
println!(" artifact:");
7070
println!(" provider: {}", model.provider);
71-
println!(" revision: {}", model.metadata.artifact.revision);
71+
println!(" revision: {}", model.metadata.cache.revision);
7272
println!(
7373
" size: {}",
74-
format_size_decimal(model.metadata.artifact.size)
74+
format_size_decimal(model.metadata.cache.size)
7575
);
76-
println!(" cache_path: {}", model.metadata.artifact.path);
76+
println!(" cache_path: {}", model.metadata.cache.path);
7777
println!("status:");
7878
println!(" created: {}", format_time_ago(&model.created_at));
7979
println!(" updated: {}", format_time_ago(&model.updated_at));
@@ -82,7 +82,7 @@ pub fn display(model: &ModelInfo) {
8282
#[cfg(test)]
8383
mod tests {
8484
use super::*;
85-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
85+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
8686
use tempfile::TempDir;
8787

8888
fn create_test_model(name: &str, uuid: &str) -> ModelInfo {
@@ -94,15 +94,15 @@ mod tests {
9494
ModelInfo {
9595
uuid: uuid.to_string(),
9696
name: name.to_string(),
97+
provider: "huggingface".to_string(),
9798
author: Some("test-author".to_string()),
9899
task: Some("text-generation".to_string()),
99100
model_series: Some("gpt2".to_string()),
100-
provider: "huggingface".to_string(),
101101
license: Some("mit".to_string()),
102102
created_at: "2025-01-01T00:00:00Z".to_string(),
103103
updated_at: "2025-01-01T00:00:00Z".to_string(),
104104
metadata: ModelMetadata {
105-
artifact: ArtifactInfo {
105+
cache: CacheInfo {
106106
revision: uuid.to_string(),
107107
size: 1000,
108108
path: "/tmp/test".to_string(),

‎src/cli/ls.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub fn execute(
5252
#[cfg(test)]
5353
mod tests {
5454
use super::*;
55-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
55+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
5656
use tempfile::TempDir;
5757

5858
fn create_test_model(name: &str, uuid: &str, author: &str) -> ModelInfo {
@@ -64,15 +64,15 @@ mod tests {
6464
ModelInfo {
6565
uuid: uuid.to_string(),
6666
name: name.to_string(),
67+
provider: "huggingface".to_string(),
6768
author: Some(author.to_string()),
6869
task: Some("text-generation".to_string()),
6970
model_series: Some("gpt2".to_string()),
70-
provider: "huggingface".to_string(),
7171
license: Some("mit".to_string()),
7272
created_at: "2025-01-01T00:00:00Z".to_string(),
7373
updated_at: "2025-01-01T00:00:00Z".to_string(),
7474
metadata: ModelMetadata {
75-
artifact: ArtifactInfo {
75+
cache: CacheInfo {
7676
revision: uuid.to_string(),
7777
size: 1000,
7878
path: "/tmp/test".to_string(),

‎src/cli/rm.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ pub fn execute(registry: &ModelRegistry, model_name: &str) -> Result<(), String>
1414
#[cfg(test)]
1515
mod tests {
1616
use super::*;
17-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
17+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
1818
use tempfile::TempDir;
1919

2020
fn create_test_model(name: &str, uuid: &str) -> ModelInfo {
@@ -26,15 +26,15 @@ mod tests {
2626
ModelInfo {
2727
uuid: uuid.to_string(),
2828
name: name.to_string(),
29+
provider: "huggingface".to_string(),
2930
author: Some("test-author".to_string()),
3031
task: Some("text-generation".to_string()),
3132
model_series: Some("gpt2".to_string()),
32-
provider: "huggingface".to_string(),
3333
license: Some("mit".to_string()),
3434
created_at: "2025-01-01T00:00:00Z".to_string(),
3535
updated_at: "2025-01-01T00:00:00Z".to_string(),
3636
metadata: ModelMetadata {
37-
artifact: ArtifactInfo {
37+
cache: CacheInfo {
3838
revision: uuid.to_string(),
3939
size: 1000,
4040
path: "/tmp/test".to_string(),
@@ -55,7 +55,7 @@ mod tests {
5555
std::fs::write(cache_dir.join("model.safetensors"), "fake data").unwrap();
5656

5757
let mut model = create_test_model("test/remove-model", "abc123");
58-
model.metadata.artifact.path = cache_dir.to_string_lossy().to_string();
58+
model.metadata.cache.path = cache_dir.to_string_lossy().to_string();
5959

6060
registry.register_model(model).unwrap();
6161
assert!(registry.get_model("test/remove-model").unwrap().is_some());

‎src/downloader/huggingface.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use indicatif::{ProgressBar, ProgressStyle};
66

77
use crate::downloader::downloader::{DownloadError, Downloader};
88
use crate::downloader::progress::{DownloadProgressManager, FileProgress};
9-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry};
9+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
1010
use crate::utils::file::{self, format_model_name};
1111

1212
/// Adapter to bridge HuggingFace's Progress trait with our FileProgress
@@ -284,14 +284,14 @@ impl Downloader for HuggingFaceDownloader {
284284
let model_size =
285285
storage_from_api.unwrap_or_else(|| progress_manager.total_downloaded_bytes());
286286

287-
let artifact = ArtifactInfo {
287+
let cache = CacheInfo {
288288
revision: sha.clone(),
289289
size: model_size,
290290
path: model_cache_path.to_string_lossy().to_string(),
291291
};
292292

293293
let metadata = ModelMetadata {
294-
artifact,
294+
cache,
295295
context_window,
296296
safetensors: safetensors_from_api,
297297
};
@@ -300,10 +300,10 @@ impl Downloader for HuggingFaceDownloader {
300300
let model_info_record = ModelInfo {
301301
uuid: sha, // Use revision SHA as UUID for now
302302
name: name.to_string(),
303+
provider: "huggingface".to_string(),
303304
author: author_from_api,
304305
task: task_from_api,
305306
model_series: model_series_from_api,
306-
provider: "huggingface".to_string(),
307307
license: license_from_api,
308308
created_at: now.clone(),
309309
updated_at: now,

‎src/registry/model_registry.rs‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,15 @@ use crate::storage::{ModelStorage, SqliteStorage};
88
use crate::utils::file;
99

1010
#[derive(Debug, Serialize, Deserialize, Clone)]
11-
pub struct ArtifactInfo {
11+
pub struct CacheInfo {
1212
pub revision: String,
1313
pub size: u64,
1414
pub path: String,
1515
}
1616

1717
#[derive(Debug, Serialize, Deserialize, Clone)]
1818
pub struct ModelMetadata {
19-
pub artifact: ArtifactInfo,
19+
pub cache: CacheInfo,
2020
#[serde(skip_serializing_if = "Option::is_none")]
2121
pub context_window: Option<u32>,
2222
#[serde(skip_serializing_if = "Option::is_none")]
@@ -27,12 +27,12 @@ pub struct ModelMetadata {
2727
pub struct ModelInfo {
2828
pub uuid: String,
2929
pub name: String,
30+
pub provider: String,
3031
pub author: Option<String>,
3132
#[serde(skip_serializing_if = "Option::is_none")]
3233
pub task: Option<String>, // Task type (image-text-to-text, text-generation)
3334
#[serde(skip_serializing_if = "Option::is_none")]
3435
pub model_series: Option<String>, // Architecture series (qwen3_5, gpt2, llama3)
35-
pub provider: String,
3636
#[serde(skip_serializing_if = "Option::is_none")]
3737
pub license: Option<String>,
3838
pub metadata: ModelMetadata,
@@ -82,7 +82,7 @@ impl ModelRegistry {
8282

8383
if let Some(info) = model_info {
8484
// Delete artifact directory if it exists
85-
let artifact_path = std::path::Path::new(&info.metadata.artifact.path);
85+
let artifact_path = std::path::Path::new(&info.metadata.cache.path);
8686
if artifact_path.exists() {
8787
fs::remove_dir_all(artifact_path)?;
8888
}
@@ -119,15 +119,15 @@ mod tests {
119119
ModelInfo {
120120
uuid: revision.to_string(),
121121
name: name.to_string(),
122+
provider: "huggingface".to_string(),
122123
author: Some("test-author".to_string()),
123124
task: Some("text-generation".to_string()),
124125
model_series: Some("gpt2".to_string()),
125-
provider: "huggingface".to_string(),
126126
license: Some("mit".to_string()),
127127
created_at: "2025-01-01T00:00:00Z".to_string(),
128128
updated_at: "2025-01-01T00:00:00Z".to_string(),
129129
metadata: ModelMetadata {
130-
artifact: ArtifactInfo {
130+
cache: CacheInfo {
131131
revision: revision.to_string(),
132132
size: 1000,
133133
path: "/tmp/test".to_string(),
@@ -202,17 +202,17 @@ mod tests {
202202
registry.register_model(model1).unwrap();
203203

204204
let mut model2 = create_test_model("test/model", "def456");
205-
model2.metadata.artifact.size = 2000;
206-
model2.metadata.artifact.path = "/tmp/test2".to_string();
205+
model2.metadata.cache.size = 2000;
206+
model2.metadata.cache.path = "/tmp/test2".to_string();
207207
model2.created_at = "2025-01-02T00:00:00Z".to_string();
208208
model2.updated_at = "2025-01-02T00:00:00Z".to_string();
209209

210210
registry.register_model(model2).unwrap();
211211

212212
let models = registry.load_models(None).unwrap();
213213
assert_eq!(models.len(), 1);
214-
assert_eq!(models[0].metadata.artifact.revision, "def456");
215-
assert_eq!(models[0].metadata.artifact.size, 2000);
214+
assert_eq!(models[0].metadata.cache.revision, "def456");
215+
assert_eq!(models[0].metadata.cache.size, 2000);
216216
// created_at should be preserved from model1
217217
assert_eq!(models[0].created_at, "2025-01-01T00:00:00Z");
218218
// updated_at should be from model2
@@ -230,7 +230,7 @@ mod tests {
230230
fs::write(cache_dir.join("test.txt"), "test data").unwrap();
231231

232232
let mut model = create_test_model("test/model", "abc123");
233-
model.metadata.artifact.path = cache_dir.to_string_lossy().to_string();
233+
model.metadata.cache.path = cache_dir.to_string_lossy().to_string();
234234

235235
registry.register_model(model).unwrap();
236236
assert_eq!(registry.load_models(None).unwrap().len(), 1);
@@ -271,7 +271,7 @@ mod tests {
271271
let model_info = retrieved.unwrap();
272272
assert_eq!(model_info.name, "test/gpt-model");
273273
assert_eq!(model_info.provider, "huggingface");
274-
assert_eq!(model_info.metadata.artifact.revision, "abc123def456");
274+
assert_eq!(model_info.metadata.cache.revision, "abc123def456");
275275
assert_eq!(model_info.model_series, Some("gpt2".to_string()));
276276
assert_eq!(model_info.metadata.context_window, Some(2048));
277277
assert_eq!(

‎src/storage/sqlite.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -112,10 +112,10 @@ impl ModelStorage for SqliteStorage {
112112
Ok(ModelInfo {
113113
uuid: row.get(0)?,
114114
name: row.get(1)?,
115+
provider: row.get(5)?,
115116
author: row.get(2)?,
116117
task: row.get(3)?,
117118
model_series: row.get(4)?,
118-
provider: row.get(5)?,
119119
license: row.get(6)?,
120120
metadata,
121121
created_at: row.get(8)?,
@@ -202,10 +202,10 @@ impl ModelStorage for SqliteStorage {
202202
Ok(ModelInfo {
203203
uuid: row.get(0)?,
204204
name: row.get(1)?,
205+
provider: row.get(5)?,
205206
author: row.get(2)?,
206207
task: row.get(3)?,
207208
model_series: row.get(4)?,
208-
provider: row.get(5)?,
209209
license: row.get(6)?,
210210
created_at: row.get(8)?,
211211
updated_at: row.get(9)?,
@@ -225,7 +225,7 @@ impl ModelStorage for SqliteStorage {
225225
#[cfg(test)]
226226
mod tests {
227227
use super::*;
228-
use crate::registry::model_registry::{ArtifactInfo, ModelMetadata};
228+
use crate::registry::model_registry::{CacheInfo, ModelMetadata};
229229
use tempfile::TempDir;
230230

231231
fn create_test_model(name: &str, uuid: &str) -> ModelInfo {
@@ -237,15 +237,15 @@ mod tests {
237237
ModelInfo {
238238
uuid: uuid.to_string(),
239239
name: name.to_string(),
240+
provider: "huggingface".to_string(),
240241
author: Some("test-author".to_string()),
241242
task: Some("text-generation".to_string()),
242243
model_series: Some("gpt2".to_string()),
243-
provider: "huggingface".to_string(),
244244
license: Some("mit".to_string()),
245245
created_at: "2025-01-01T00:00:00Z".to_string(),
246246
updated_at: "2025-01-01T00:00:00Z".to_string(),
247247
metadata: ModelMetadata {
248-
artifact: ArtifactInfo {
248+
cache: CacheInfo {
249249
revision: "abc123".to_string(),
250250
size: 1000,
251251
path: "/tmp/test".to_string(),
@@ -336,8 +336,8 @@ mod tests {
336336
storage.register_model(model).unwrap();
337337

338338
let retrieved = storage.get_model("test/model").unwrap().unwrap();
339-
assert_eq!(retrieved.metadata.artifact.revision, "abc123");
340-
assert_eq!(retrieved.metadata.artifact.size, 1000);
339+
assert_eq!(retrieved.metadata.cache.revision, "abc123");
340+
assert_eq!(retrieved.metadata.cache.size, 1000);
341341
assert_eq!(retrieved.metadata.context_window, Some(2048));
342342
assert!(retrieved.metadata.safetensors.is_some());
343343

0 commit comments

Comments
 (0)