1use std::collections::{BTreeSet, HashMap};
7use std::fmt;
8
9bitflags::bitflags! {
14 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
19 pub struct CapabilityFlags: u32 {
20 const REASONING = 0b0000_0001;
22 const TOOL_CALLING = 0b0000_0010;
24 const VISION = 0b0000_0100;
26 const CODE = 0b0000_1000;
28 const MOE = 0b0001_0000;
30 const MTP = 0b0010_0000;
35 const EMBEDDING = 0b0100_0000;
42 }
43}
44
45#[derive(Debug, Clone, Default, PartialEq, Eq)]
50pub struct GgufCapabilities {
51 pub flags: CapabilityFlags,
53 pub extensions: BTreeSet<String>,
55 pub dialect: Option<crate::domain::dialect::DialectSpec>,
60}
61
62impl GgufCapabilities {
63 #[must_use]
65 pub const fn empty() -> Self {
66 Self {
67 flags: CapabilityFlags::empty(),
68 extensions: BTreeSet::new(),
69 dialect: None,
70 }
71 }
72
73 #[must_use]
75 pub const fn has_reasoning(&self) -> bool {
76 self.flags.contains(CapabilityFlags::REASONING)
77 }
78
79 #[must_use]
81 pub const fn has_tool_calling(&self) -> bool {
82 self.flags.contains(CapabilityFlags::TOOL_CALLING)
83 }
84
85 #[must_use]
87 pub const fn has_vision(&self) -> bool {
88 self.flags.contains(CapabilityFlags::VISION)
89 }
90
91 #[must_use]
93 pub const fn has_mtp(&self) -> bool {
94 self.flags.contains(CapabilityFlags::MTP)
95 }
96
97 #[must_use]
99 pub const fn has_embedding(&self) -> bool {
100 self.flags.contains(CapabilityFlags::EMBEDDING)
101 }
102
103 #[must_use]
107 pub fn to_tags(&self) -> Vec<String> {
108 let mut tags = Vec::new();
109
110 if self.has_reasoning() {
114 tags.push(super::capability_tags::REASONING.to_string());
115 }
116 if self.has_tool_calling() {
117 tags.push(super::capability_tags::AGENT.to_string());
119 }
120 if self.has_vision() {
121 tags.push(super::capability_tags::VISION.to_string());
122 }
123 if self.flags.contains(CapabilityFlags::CODE) {
124 tags.push(super::capability_tags::CODE.to_string());
125 }
126 if self.flags.contains(CapabilityFlags::MOE) {
127 tags.push(super::capability_tags::MOE.to_string());
128 }
129 if self.has_mtp() {
130 tags.push(super::capability_tags::MTP.to_string());
132 }
133 if self.has_embedding() {
134 tags.push(super::capability_tags::EMBEDDING.to_string());
136 }
137
138 for ext in &self.extensions {
140 if !tags.contains(ext) {
141 tags.push(ext.clone());
142 }
143 }
144
145 tags
146 }
147}
148
149#[derive(Debug, Clone)]
157pub enum GgufValue {
158 U8(u8),
159 I8(i8),
160 U16(u16),
161 I16(i16),
162 U32(u32),
163 I32(i32),
164 F32(f32),
165 Bool(bool),
166 String(String),
167 Array(Vec<Self>),
168 U64(u64),
169 I64(i64),
170 F64(f64),
171}
172
173impl fmt::Display for GgufValue {
174 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175 match self {
176 Self::U8(v) => write!(f, "{v}"),
177 Self::I8(v) => write!(f, "{v}"),
178 Self::U16(v) => write!(f, "{v}"),
179 Self::I16(v) => write!(f, "{v}"),
180 Self::U32(v) => write!(f, "{v}"),
181 Self::I32(v) => write!(f, "{v}"),
182 Self::F32(v) => write!(f, "{v}"),
183 Self::Bool(v) => write!(f, "{v}"),
184 Self::String(v) => write!(f, "{v}"),
185 Self::U64(v) => write!(f, "{v}"),
186 Self::I64(v) => write!(f, "{v}"),
187 Self::F64(v) => write!(f, "{v}"),
188 Self::Array(arr) => {
189 if arr.len() > 10 {
191 write!(f, "[Array with {} elements]", arr.len())
192 } else {
193 write!(
194 f,
195 "[{}]",
196 arr.iter()
197 .map(std::string::ToString::to_string)
198 .collect::<Vec<_>>()
199 .join(", ")
200 )
201 }
202 }
203 }
204 }
205}
206
207impl GgufValue {
208 #[must_use]
213 #[allow(clippy::cast_sign_loss)]
214 pub fn as_u64(&self) -> Option<u64> {
215 match self {
216 Self::U8(v) => Some(u64::from(*v)),
217 Self::U16(v) => Some(u64::from(*v)),
218 Self::U32(v) => Some(u64::from(*v)),
219 Self::U64(v) => Some(*v),
220 Self::I8(v) if *v >= 0 => Some(*v as u64),
221 Self::I16(v) if *v >= 0 => Some(*v as u64),
222 Self::I32(v) if *v >= 0 => Some(*v as u64),
223 Self::I64(v) if *v >= 0 => Some(*v as u64),
224 _ => None,
225 }
226 }
227
228 #[must_use]
230 #[allow(clippy::cast_precision_loss)]
231 pub fn as_f64(&self) -> Option<f64> {
232 match self {
233 Self::F32(v) => Some(f64::from(*v)),
234 Self::F64(v) => Some(*v),
235 Self::U8(v) => Some(f64::from(*v)),
236 Self::U16(v) => Some(f64::from(*v)),
237 Self::U32(v) => Some(f64::from(*v)),
238 Self::U64(v) => Some(*v as f64),
239 Self::I8(v) => Some(f64::from(*v)),
240 Self::I16(v) => Some(f64::from(*v)),
241 Self::I32(v) => Some(f64::from(*v)),
242 Self::I64(v) => Some(*v as f64),
243 _ => None,
244 }
245 }
246
247 #[must_use]
249 pub fn as_str(&self) -> Option<&str> {
250 match self {
251 Self::String(s) => Some(s),
252 _ => None,
253 }
254 }
255}
256
257#[derive(Debug, Clone, Default)]
266pub struct GgufMetadata {
267 pub architecture: Option<String>,
269 pub quantization: Option<String>,
271 pub param_count_b: Option<f64>,
273 pub context_length: Option<u64>,
275 pub expert_count: Option<u32>,
277 pub expert_used_count: Option<u32>,
279 pub expert_shared_count: Option<u32>,
281 pub metadata: HashMap<String, String>,
283}
284
285pub type RawMetadata = HashMap<String, GgufValue>;
289
290#[cfg(test)]
295mod tests {
296 use super::*;
297
298 #[test]
299 fn test_capabilities_empty() {
300 let caps = GgufCapabilities::empty();
301 assert!(!caps.has_reasoning());
302 assert!(!caps.has_tool_calling());
303 assert!(caps.to_tags().is_empty());
304 }
305
306 #[test]
307 fn test_capabilities_flags() {
308 let caps = GgufCapabilities {
309 flags: CapabilityFlags::REASONING | CapabilityFlags::TOOL_CALLING,
310 extensions: BTreeSet::new(),
311 dialect: None,
312 };
313 assert!(caps.has_reasoning());
314 assert!(caps.has_tool_calling());
315
316 let tags = caps.to_tags();
317 assert!(tags.contains(&"reasoning".to_string()));
318 assert!(tags.contains(&"agent".to_string()));
319 }
320
321 #[test]
322 fn test_capabilities_extensions() {
323 let mut extensions = BTreeSet::new();
324 extensions.insert("experimental-feature".to_string());
325
326 let caps = GgufCapabilities {
327 flags: CapabilityFlags::empty(),
328 extensions,
329 dialect: None,
330 };
331
332 let tags = caps.to_tags();
333 assert!(tags.contains(&"experimental-feature".to_string()));
334 }
335
336 #[test]
337 fn test_gguf_value_as_u64() {
338 assert_eq!(GgufValue::U32(4096).as_u64(), Some(4096));
339 assert_eq!(GgufValue::I32(-1).as_u64(), None);
340 assert_eq!(GgufValue::String("hello".to_string()).as_u64(), None);
341 assert_eq!(GgufValue::I32(100).as_u64(), Some(100));
342 }
343
344 #[test]
345 fn test_gguf_value_as_f64() {
346 assert!((GgufValue::F32(7.5).as_f64().unwrap() - 7.5).abs() < f64::EPSILON);
347 assert!((GgufValue::U64(1000).as_f64().unwrap() - 1000.0).abs() < f64::EPSILON);
348 assert_eq!(GgufValue::Bool(true).as_f64(), None);
349 }
350
351 #[test]
352 fn test_gguf_value_display() {
353 assert_eq!(GgufValue::U32(42).to_string(), "42");
354 assert_eq!(GgufValue::String("test".to_string()).to_string(), "test");
355
356 let large_array = GgufValue::Array(vec![GgufValue::U8(0); 100]);
357 assert!(large_array.to_string().contains("100 elements"));
358 }
359}