oxirs_graphrag/model_loader/
registry.rs

1//! Thread-safe model registry for GGUF model metadata.
2//!
3//! The registry stores [`ModelInfo`] records (path + parsed metadata) keyed by
4//! name-based [`ModelHandle`]s.  No tensor weights are held in memory.
5
6use std::collections::HashMap;
7use std::path::PathBuf;
8use std::sync::RwLock;
9
10use super::gguf_parser::{GgufMetadata, GgufParseError, GgufParser};
11
12// ─── ModelHandle ─────────────────────────────────────────────────────────────
13
14/// An opaque handle to a registered model, identified by name.
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
16pub struct ModelHandle(String);
17
18impl ModelHandle {
19    /// Return the model name.
20    pub fn name(&self) -> &str {
21        &self.0
22    }
23}
24
25// ─── ModelInfo ───────────────────────────────────────────────────────────────
26
27/// Lightweight model record held by the registry.
28///
29/// Contains only the file path and the parsed GGUF metadata header; tensor
30/// weight data is never loaded.
31#[derive(Debug, Clone)]
32pub struct ModelInfo {
33    /// Opaque handle for this model.
34    pub handle: ModelHandle,
35    /// Absolute path to the `.gguf` file.
36    pub path: PathBuf,
37    /// Parsed GGUF header metadata.
38    pub metadata: GgufMetadata,
39}
40
41impl ModelInfo {
42    /// Return the model's hidden / embedding dimension, if declared.
43    pub fn embedding_dim(&self) -> Option<usize> {
44        self.metadata.arch.embedding_length.map(|v| v as usize)
45    }
46
47    /// Return the vocabulary size, if declared.
48    pub fn vocab_size(&self) -> Option<usize> {
49        self.metadata.arch.vocab_size.map(|v| v as usize)
50    }
51}
52
53// ─── RegistryError ───────────────────────────────────────────────────────────
54
55/// Errors produced by the model registry.
56#[derive(Debug, thiserror::Error)]
57pub enum RegistryError {
58    /// A model with this name is already registered.
59    #[error("model '{0}' already registered")]
60    AlreadyRegistered(String),
61    /// No model with this name was found.
62    #[error("model '{0}' not found")]
63    NotFound(String),
64    /// Failed to parse the GGUF file.
65    #[error("GGUF parse error: {0}")]
66    ParseError(#[from] GgufParseError),
67}
68
69// ─── ModelRegistry ───────────────────────────────────────────────────────────
70
71/// Thread-safe registry of loaded model metadata.
72///
73/// Models are keyed by name; each registration parses (or accepts) GGUF
74/// metadata without loading tensor weights.
75///
76/// # Example
77///
78/// ```rust
79/// # #[cfg(feature = "gguf-loader")]
80/// # {
81/// use oxirs_graphrag::model_loader::{ModelRegistry, GgufMetadata, GgufParser};
82/// use std::path::PathBuf;
83///
84/// let registry = ModelRegistry::new();
85/// // In real usage you would supply a path to a .gguf file:
86/// // let handle = registry.register("llama-3-8b", PathBuf::from("model.gguf")).unwrap();
87/// assert!(registry.is_empty());
88/// # }
89/// ```
90#[derive(Default)]
91pub struct ModelRegistry {
92    models: RwLock<HashMap<String, ModelInfo>>,
93}
94
95impl ModelRegistry {
96    /// Create an empty registry.
97    pub fn new() -> Self {
98        Self::default()
99    }
100
101    /// Register a model by parsing its GGUF file.
102    ///
103    /// Returns an error if `name` is already registered or the file cannot be
104    /// parsed.
105    pub fn register(&self, name: &str, path: PathBuf) -> Result<ModelHandle, RegistryError> {
106        let metadata = GgufParser::parse_file(&path)?;
107        self.register_with_metadata(name, path, metadata)
108    }
109
110    /// Register a model with pre-parsed metadata.
111    ///
112    /// Useful for tests that construct synthetic metadata without a real file.
113    pub fn register_with_metadata(
114        &self,
115        name: &str,
116        path: PathBuf,
117        metadata: GgufMetadata,
118    ) -> Result<ModelHandle, RegistryError> {
119        let mut map = self.models.write().expect("registry lock poisoned");
120        if map.contains_key(name) {
121            return Err(RegistryError::AlreadyRegistered(name.to_owned()));
122        }
123        let handle = ModelHandle(name.to_owned());
124        let info = ModelInfo {
125            handle: handle.clone(),
126            path,
127            metadata,
128        };
129        map.insert(name.to_owned(), info);
130        Ok(handle)
131    }
132
133    /// Look up a model by its handle.
134    pub fn get(&self, handle: &ModelHandle) -> Option<ModelInfo> {
135        let map = self.models.read().expect("registry lock poisoned");
136        map.get(handle.name()).cloned()
137    }
138
139    /// Look up a model by name.
140    pub fn get_by_name(&self, name: &str) -> Option<ModelInfo> {
141        let map = self.models.read().expect("registry lock poisoned");
142        map.get(name).cloned()
143    }
144
145    /// List all registered model handles.
146    pub fn list(&self) -> Vec<ModelHandle> {
147        let map = self.models.read().expect("registry lock poisoned");
148        map.values().map(|info| info.handle.clone()).collect()
149    }
150
151    /// Remove a model from the registry.
152    ///
153    /// Returns `true` if the model existed and was removed.
154    pub fn remove(&self, handle: &ModelHandle) -> bool {
155        let mut map = self.models.write().expect("registry lock poisoned");
156        map.remove(handle.name()).is_some()
157    }
158
159    /// Return the number of registered models.
160    pub fn len(&self) -> usize {
161        let map = self.models.read().expect("registry lock poisoned");
162        map.len()
163    }
164
165    /// Return `true` if no models are registered.
166    pub fn is_empty(&self) -> bool {
167        self.len() == 0
168    }
169}
oxirs_graphrag/model_loader/registry.rs

oxirs_graphrag/model_loader/
registry.rs
