gglib_core/services/
model_registrar.rs

1//! Model registrar service implementation.
2//!
3//! This service implements `ModelRegistrarPort` using the `ModelRepository`
4//! and `GgufParserPort` dependencies. It's used by the download manager
5//! to register completed downloads.
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use chrono::Utc;
11
12use super::model_import::fetch_published_sampling;
13use super::{HfOrigin, ModelOrigin, build_new_model};
14use crate::domain::{Model, NewModelFile};
15use crate::ports::huggingface::HfClientPort;
16use crate::ports::{
17    CompletedDownload, GgufParserPort, ModelRegistrarPort, ModelRepository, RepositoryError,
18};
19
20/// Repository trait for model files metadata.
21///
22/// We don't depend on `gglib_db` directly - adapters inject the implementation.
23/// This type is re-exported from `gglib_db` for use in adapters.
24#[async_trait]
25pub trait ModelFilesRepositoryPort: Send + Sync {
26    /// Insert a new model file record.
27    async fn insert(&self, model_file: &NewModelFile) -> anyhow::Result<()>;
28}
29
30/// Implementation of the model registrar port.
31///
32/// This service composes over `ModelRepository` for persistence and
33/// `GgufParserPort` for metadata extraction.
34pub struct ModelRegistrar {
35    /// Repository for persisting models.
36    model_repo: Arc<dyn ModelRepository>,
37    /// Parser for extracting GGUF metadata.
38    gguf_parser: Arc<dyn GgufParserPort>,
39    /// Repository for persisting model file metadata.
40    model_files_repo: Option<Arc<dyn ModelFilesRepositoryPort>>,
41    /// Used to look up the model author's published sampling recipe.
42    ///
43    /// Optional, and absent means "do not look" rather than "cannot register".
44    /// A registrar without one behaves exactly as it did before this existed:
45    /// the `reasoning` tag guess applies. That keeps the feature off in tests
46    /// and in any embedding that has no HF client, without either having to
47    /// know it exists.
48    hf_client: Option<Arc<dyn HfClientPort>>,
49}
50
51impl ModelRegistrar {
52    /// Create a new model registrar.
53    ///
54    /// # Arguments
55    ///
56    /// * `model_repo` - Repository for persisting models
57    /// * `gguf_parser` - Parser for extracting GGUF metadata
58    /// * `model_files_repo` - Optional repository for persisting model file metadata
59    pub fn new(
60        model_repo: Arc<dyn ModelRepository>,
61        gguf_parser: Arc<dyn GgufParserPort>,
62        model_files_repo: Option<Arc<dyn ModelFilesRepositoryPort>>,
63    ) -> Self {
64        Self {
65            model_repo,
66            gguf_parser,
67            model_files_repo,
68            hf_client: None,
69        }
70    }
71
72    /// Look up published sampling recipes at import, using `client`.
73    ///
74    /// A builder method rather than a fourth constructor parameter: every
75    /// existing call site wants the previous behaviour, and only the
76    /// application wiring has an HF client to give.
77    #[must_use]
78    pub fn with_hf_client(mut self, client: Arc<dyn HfClientPort>) -> Self {
79        self.hf_client = Some(client);
80        self
81    }
82}
83
84#[async_trait]
85impl ModelRegistrarPort for ModelRegistrar {
86    async fn register_model(&self, download: &CompletedDownload) -> Result<Model, RepositoryError> {
87        let file_path = download.db_path();
88
89        // Parse GGUF metadata from the downloaded file
90        let gguf_metadata = self.gguf_parser.parse(file_path).ok();
91
92        // Best-effort, and deliberately before the row is built: a recipe the
93        // author published is better evidence than the tag guess
94        // `build_new_model` would otherwise write. Returns `None` for every
95        // failure — gated repo, offline, nothing published — and the import
96        // proceeds unchanged. See `fetch_published_sampling`.
97        let published = match &self.hf_client {
98            Some(client) => {
99                fetch_published_sampling(client.as_ref(), &download.repo_id, &download.hf_tags)
100                    .await
101            }
102            None => None,
103        };
104
105        let origin = ModelOrigin::HuggingFace(HfOrigin {
106            repo_id: &download.repo_id,
107            commit_sha: &download.commit_sha,
108            hf_tags: &download.hf_tags,
109            quantization_fallback: download.quantization,
110            file_paths: download.file_paths.as_deref(),
111            published_sampling: published.as_ref(),
112        });
113        let model = build_new_model(
114            file_path,
115            gguf_metadata.as_ref(),
116            self.gguf_parser.as_ref(),
117            &origin,
118            Utc::now(),
119        );
120
121        let registered = self.model_repo.insert(&model).await?;
122
123        // Insert model_files records with OIDs for each shard (if repo is available)
124        if let Some(ref repo) = self.model_files_repo {
125            for (file_index, file_entry) in download.hf_file_entries.iter().enumerate() {
126                if let Some(size) = file_entry.size {
127                    #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
128                    let model_file = NewModelFile::new(
129                        registered.id,
130                        file_entry.path.clone(),
131                        file_index as i32,
132                        size as i64,
133                        file_entry.oid.clone(),
134                    );
135
136                    if let Err(e) = repo.insert(&model_file).await {
137                        // Soft fail - log but don't propagate error
138                        tracing::warn!(
139                            model_id = registered.id,
140                            file_path = %file_entry.path,
141                            error = %e,
142                            "Failed to insert model_files record - verification features may be unavailable"
143                        );
144                    }
145                }
146            }
147        }
148
149        Ok(registered)
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use crate::domain::{Model, NewModel};
157    use crate::download::Quantization;
158    use crate::ports::NoopGgufParser;
159    use std::path::PathBuf;
160    use std::sync::Mutex;
161
162    /// Mock model repository for testing.
163    struct MockModelRepo {
164        models: Mutex<Vec<Model>>,
165        next_id: Mutex<i64>,
166    }
167
168    impl MockModelRepo {
169        fn new() -> Self {
170            Self {
171                models: Mutex::new(Vec::new()),
172                next_id: Mutex::new(1),
173            }
174        }
175    }
176
177    #[async_trait]
178    impl ModelRepository for MockModelRepo {
179        async fn list(&self) -> Result<Vec<Model>, RepositoryError> {
180            Ok(self.models.lock().unwrap().clone())
181        }
182
183        async fn get_by_id(&self, id: i64) -> Result<Model, RepositoryError> {
184            self.models
185                .lock()
186                .unwrap()
187                .iter()
188                .find(|m| m.id == id)
189                .cloned()
190                .ok_or_else(|| RepositoryError::NotFound(format!("id={id}")))
191        }
192
193        async fn get_by_name(&self, name: &str) -> Result<Model, RepositoryError> {
194            self.models
195                .lock()
196                .unwrap()
197                .iter()
198                .find(|m| m.name == name)
199                .cloned()
200                .ok_or_else(|| RepositoryError::NotFound(format!("name={name}")))
201        }
202
203        async fn find_by_path(
204            &self,
205            path: &std::path::Path,
206        ) -> Result<Option<Model>, RepositoryError> {
207            Ok(self
208                .models
209                .lock()
210                .unwrap()
211                .iter()
212                .find(|m| m.file_path.as_path() == path)
213                .cloned())
214        }
215
216        async fn insert(&self, model: &NewModel) -> Result<Model, RepositoryError> {
217            let mut id = self.next_id.lock().unwrap();
218            let persisted = Model {
219                dialect_spec: None,
220                id: *id,
221                name: model.name.clone(),
222                model_key: String::new(),
223                file_path: model.file_path.clone(),
224                param_count_b: model.param_count_b,
225                architecture: model.architecture.clone(),
226                quantization: model.quantization.clone(),
227                context_length: model.context_length,
228                expert_count: model.expert_count,
229                expert_used_count: model.expert_used_count,
230                expert_shared_count: model.expert_shared_count,
231                metadata: model.metadata.clone(),
232                added_at: model.added_at,
233                hf_repo_id: model.hf_repo_id.clone(),
234                hf_commit_sha: model.hf_commit_sha.clone(),
235                hf_filename: model.hf_filename.clone(),
236                capabilities: model.capabilities,
237                download_date: model.download_date,
238                last_update_check: model.last_update_check,
239                tags: model.tags.clone(),
240                inference_defaults: model.inference_defaults.clone(),
241                defaults_origin: model.defaults_origin,
242                server_defaults: model.server_defaults.clone(),
243                template_caps: None,
244                benchmark_summary: None,
245            };
246            // Mirror the `SQLite` repository: a repeat registration of the
247            // same file updates that row and keeps its id. This double models
248            // registration-after-download, which is the very path the trait
249            // doc cites as the reason `insert` upserts — a double that appends
250            // contradicts the contract it exists to stand in for.
251            let mut models = self.models.lock().unwrap();
252            if let Some(index) = models
253                .iter()
254                .position(|m| m.file_path == persisted.file_path)
255            {
256                let mut updated = persisted.clone();
257                updated.id = models[index].id;
258                models[index] = updated.clone();
259                drop(models);
260                return Ok(updated);
261            }
262            *id += 1;
263            drop(id);
264            models.push(persisted.clone());
265            drop(models);
266            Ok(persisted)
267        }
268
269        async fn update(&self, _model: &Model) -> Result<(), RepositoryError> {
270            Ok(())
271        }
272
273        async fn delete(&self, _id: i64) -> Result<(), RepositoryError> {
274            Ok(())
275        }
276    }
277
278    #[tokio::test]
279    async fn test_register_model_basic() {
280        let repo = Arc::new(MockModelRepo::new());
281        let parser = Arc::new(NoopGgufParser);
282        let registrar = ModelRegistrar::new(repo.clone(), parser, None);
283
284        let download = CompletedDownload {
285            primary_path: PathBuf::from("/models/test-model-q4_k_m.gguf"),
286            all_paths: vec![PathBuf::from("/models/test-model-q4_k_m.gguf")],
287            quantization: Quantization::Q4KM,
288            repo_id: "test/model".to_string(),
289            commit_sha: "abc123".to_string(),
290            is_sharded: false,
291            file_paths: None,
292            hf_tags: vec![],
293            hf_file_entries: vec![],
294        };
295
296        let result = registrar.register_model(&download).await;
297        assert!(result.is_ok());
298
299        let model = result.unwrap();
300        assert_eq!(model.name, "model");
301        assert_eq!(model.hf_repo_id, Some("test/model".to_string()));
302        assert_eq!(model.hf_commit_sha, Some("abc123".to_string()));
303        assert_eq!(model.quantization, Some("Q4_K_M".to_string()));
304    }
305
306    #[tokio::test]
307    async fn test_register_sharded_model() {
308        let repo = Arc::new(MockModelRepo::new());
309        let parser = Arc::new(NoopGgufParser);
310        let registrar = ModelRegistrar::new(repo.clone(), parser, None);
311
312        let download = CompletedDownload {
313            primary_path: PathBuf::from("/models/llama-00001-of-00004.gguf"),
314            all_paths: vec![
315                PathBuf::from("/models/llama-00001-of-00004.gguf"),
316                PathBuf::from("/models/llama-00002-of-00004.gguf"),
317                PathBuf::from("/models/llama-00003-of-00004.gguf"),
318                PathBuf::from("/models/llama-00004-of-00004.gguf"),
319            ],
320            quantization: Quantization::Q8_0,
321            repo_id: "test/llama".to_string(),
322            commit_sha: "def456".to_string(),
323            is_sharded: true,
324            file_paths: None,
325            hf_tags: vec![],
326            hf_file_entries: vec![],
327        };
328
329        let result = registrar.register_model(&download).await;
330        assert!(result.is_ok());
331
332        let model = result.unwrap();
333        assert_eq!(model.quantization, Some("Q8_0".to_string()));
334        assert_eq!(model.name, "llama");
335    }
336}
gglib_core/services/model_registrar.rs

gglib_core/services/
model_registrar.rs
