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 }
31}
32
33#[derive(Debug, Clone, Default, PartialEq, Eq)]
38pub struct GgufCapabilities {
39 pub flags: CapabilityFlags,
41 pub extensions: BTreeSet<String>,
43}
44
45impl GgufCapabilities {
46 #[must_use]
48 pub const fn empty() -> Self {
49 Self {
50 flags: CapabilityFlags::empty(),
51 extensions: BTreeSet::new(),
52 }
53 }
54
55 #[must_use]
57 pub const fn has_reasoning(&self) -> bool {
58 self.flags.contains(CapabilityFlags::REASONING)
59 }
60
61 #[must_use]
63 pub const fn has_tool_calling(&self) -> bool {
64 self.flags.contains(CapabilityFlags::TOOL_CALLING)
65 }
66
67 #[must_use]
69 pub const fn has_vision(&self) -> bool {
70 self.flags.contains(CapabilityFlags::VISION)
71 }
72
73 #[must_use]
77 pub fn to_tags(&self) -> Vec<String> {
78 let mut tags = Vec::new();
79
80 if self.has_reasoning() {
81 tags.push("reasoning".to_string());
82 }
83 if self.has_tool_calling() {
84 tags.push("agent".to_string());
86 }
87 if self.has_vision() {
88 tags.push("vision".to_string());
89 }
90 if self.flags.contains(CapabilityFlags::CODE) {
91 tags.push("code".to_string());
92 }
93 if self.flags.contains(CapabilityFlags::MOE) {
94 tags.push("moe".to_string());
95 }
96
97 for ext in &self.extensions {
99 if !tags.contains(ext) {
100 tags.push(ext.clone());
101 }
102 }
103
104 tags
105 }
106}
107
108#[derive(Debug, Clone)]
116pub enum GgufValue {
117 U8(u8),
118 I8(i8),
119 U16(u16),
120 I16(i16),
121 U32(u32),
122 I32(i32),
123 F32(f32),
124 Bool(bool),
125 String(String),
126 Array(Vec<GgufValue>),
127 U64(u64),
128 I64(i64),
129 F64(f64),
130}
131
132impl fmt::Display for GgufValue {
133 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
134 match self {
135 Self::U8(v) => write!(f, "{v}"),
136 Self::I8(v) => write!(f, "{v}"),
137 Self::U16(v) => write!(f, "{v}"),
138 Self::I16(v) => write!(f, "{v}"),
139 Self::U32(v) => write!(f, "{v}"),
140 Self::I32(v) => write!(f, "{v}"),
141 Self::F32(v) => write!(f, "{v}"),
142 Self::Bool(v) => write!(f, "{v}"),
143 Self::String(v) => write!(f, "{v}"),
144 Self::U64(v) => write!(f, "{v}"),
145 Self::I64(v) => write!(f, "{v}"),
146 Self::F64(v) => write!(f, "{v}"),
147 Self::Array(arr) => {
148 if arr.len() > 10 {
150 write!(f, "[Array with {} elements]", arr.len())
151 } else {
152 write!(
153 f,
154 "[{}]",
155 arr.iter()
156 .map(std::string::ToString::to_string)
157 .collect::<Vec<_>>()
158 .join(", ")
159 )
160 }
161 }
162 }
163 }
164}
165
166impl GgufValue {
167 #[must_use]
172 #[allow(clippy::cast_sign_loss)]
173 pub fn as_u64(&self) -> Option<u64> {
174 match self {
175 Self::U8(v) => Some(u64::from(*v)),
176 Self::U16(v) => Some(u64::from(*v)),
177 Self::U32(v) => Some(u64::from(*v)),
178 Self::U64(v) => Some(*v),
179 Self::I8(v) if *v >= 0 => Some(*v as u64),
180 Self::I16(v) if *v >= 0 => Some(*v as u64),
181 Self::I32(v) if *v >= 0 => Some(*v as u64),
182 Self::I64(v) if *v >= 0 => Some(*v as u64),
183 _ => None,
184 }
185 }
186
187 #[must_use]
189 #[allow(clippy::cast_precision_loss)]
190 pub fn as_f64(&self) -> Option<f64> {
191 match self {
192 Self::F32(v) => Some(f64::from(*v)),
193 Self::F64(v) => Some(*v),
194 Self::U8(v) => Some(f64::from(*v)),
195 Self::U16(v) => Some(f64::from(*v)),
196 Self::U32(v) => Some(f64::from(*v)),
197 Self::U64(v) => Some(*v as f64),
198 Self::I8(v) => Some(f64::from(*v)),
199 Self::I16(v) => Some(f64::from(*v)),
200 Self::I32(v) => Some(f64::from(*v)),
201 Self::I64(v) => Some(*v as f64),
202 _ => None,
203 }
204 }
205
206 #[must_use]
208 pub fn as_str(&self) -> Option<&str> {
209 match self {
210 Self::String(s) => Some(s),
211 _ => None,
212 }
213 }
214}
215
216#[derive(Debug, Clone, Default)]
225pub struct GgufMetadata {
226 pub name: Option<String>,
228 pub architecture: Option<String>,
230 pub quantization: Option<String>,
232 pub param_count_b: Option<f64>,
234 pub context_length: Option<u64>,
236 pub metadata: HashMap<String, String>,
238}
239
240pub type RawMetadata = HashMap<String, GgufValue>;
244
245#[derive(Debug, Clone, Default)]
251pub struct ReasoningDetection {
252 pub supports_reasoning: bool,
254 pub confidence: f32,
256 pub matched_patterns: Vec<String>,
258 pub suggested_format: Option<String>,
260}
261
262#[derive(Debug, Clone, Default)]
264pub struct ToolCallingDetection {
265 pub supports_tool_calling: bool,
267 pub confidence: f32,
269 pub matched_patterns: Vec<String>,
271 pub detected_format: Option<String>,
273}
274
275#[cfg(test)]
280mod tests {
281 use super::*;
282
283 #[test]
284 fn test_capabilities_empty() {
285 let caps = GgufCapabilities::empty();
286 assert!(!caps.has_reasoning());
287 assert!(!caps.has_tool_calling());
288 assert!(caps.to_tags().is_empty());
289 }
290
291 #[test]
292 fn test_capabilities_flags() {
293 let caps = GgufCapabilities {
294 flags: CapabilityFlags::REASONING | CapabilityFlags::TOOL_CALLING,
295 extensions: BTreeSet::new(),
296 };
297 assert!(caps.has_reasoning());
298 assert!(caps.has_tool_calling());
299
300 let tags = caps.to_tags();
301 assert!(tags.contains(&"reasoning".to_string()));
302 assert!(tags.contains(&"agent".to_string()));
303 }
304
305 #[test]
306 fn test_capabilities_extensions() {
307 let mut extensions = BTreeSet::new();
308 extensions.insert("experimental-feature".to_string());
309
310 let caps = GgufCapabilities {
311 flags: CapabilityFlags::empty(),
312 extensions,
313 };
314
315 let tags = caps.to_tags();
316 assert!(tags.contains(&"experimental-feature".to_string()));
317 }
318
319 #[test]
320 fn test_gguf_value_as_u64() {
321 assert_eq!(GgufValue::U32(4096).as_u64(), Some(4096));
322 assert_eq!(GgufValue::I32(-1).as_u64(), None);
323 assert_eq!(GgufValue::String("hello".to_string()).as_u64(), None);
324 assert_eq!(GgufValue::I32(100).as_u64(), Some(100));
325 }
326
327 #[test]
328 fn test_gguf_value_as_f64() {
329 assert!((GgufValue::F32(7.5).as_f64().unwrap() - 7.5).abs() < f64::EPSILON);
330 assert!((GgufValue::U64(1000).as_f64().unwrap() - 1000.0).abs() < f64::EPSILON);
331 assert_eq!(GgufValue::Bool(true).as_f64(), None);
332 }
333
334 #[test]
335 fn test_gguf_value_display() {
336 assert_eq!(GgufValue::U32(42).to_string(), "42");
337 assert_eq!(GgufValue::String("test".to_string()).to_string(), "test");
338
339 let large_array = GgufValue::Array(vec![GgufValue::U8(0); 100]);
340 assert!(large_array.to_string().contains("100 elements"));
341 }
342}