Skip to main content

gglib_core/domain/
gguf.rs

1//! GGUF domain types.
2//!
3//! This module contains the domain-facing types for GGUF file metadata
4//! and model capabilities. Parsing logic lives in `gglib-gguf`.
5
6use std::collections::{BTreeSet, HashMap};
7use std::fmt;
8
9// =============================================================================
10// Capabilities (Structured, forward-compatible)
11// =============================================================================
12
13bitflags::bitflags! {
14    /// Known model capabilities detected from GGUF metadata.
15    ///
16    /// Uses bitflags for compile-time safety on stable capabilities.
17    /// Unknown/experimental capabilities go in `GgufCapabilities::extensions`.
18    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
19    pub struct CapabilityFlags: u32 {
20        /// Model supports reasoning/thinking (e.g., DeepSeek R1, QwQ).
21        const REASONING = 0b0000_0001;
22        /// Model supports tool/function calling (e.g., Hermes, Functionary).
23        const TOOL_CALLING = 0b0000_0010;
24        /// Model supports vision/image input.
25        const VISION = 0b0000_0100;
26        /// Model supports code generation.
27        const CODE = 0b0000_1000;
28        /// Model is a mixture-of-experts architecture.
29        const MOE = 0b0001_0000;
30        /// Model contains embedded MTP (Multi-Token Prediction) draft heads.
31        ///
32        /// Detected via the `{arch}.nextn_predict_layers > 0` GGUF metadata key.
33        /// Enables `--spec-type draft-mtp` speculative decoding in llama-server.
34        const MTP = 0b0010_0000;
35        /// Model produces embeddings rather than generated text.
36        ///
37        /// Detected from a non-none `{arch}.pooling_type` or an encoder-only
38        /// `general.architecture`. Enables `--embeddings` in llama-server,
39        /// which restricts that server to the embedding use case — it will
40        /// refuse chat completions.
41        const EMBEDDING = 0b0100_0000;
42    }
43}
44
45/// Model capabilities detected from GGUF metadata.
46///
47/// Combines stable known capabilities (bitflags) with forward-compatible
48/// extension strings for new/experimental capabilities.
49#[derive(Debug, Clone, Default, PartialEq, Eq)]
50pub struct GgufCapabilities {
51    /// Known stable capabilities (compile-time checked).
52    pub flags: CapabilityFlags,
53    /// Unknown/experimental capabilities (forward-compatible).
54    pub extensions: BTreeSet<String>,
55    /// Tool-call dialect identified at detection time, if any.
56    ///
57    /// `None` means no structured dialect is known — consumers fall back to
58    /// mapping `format:*` tags to a builtin spec.
59    pub dialect: Option<crate::domain::dialect::DialectSpec>,
60}
61
62impl GgufCapabilities {
63    /// Create empty capabilities.
64    #[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    /// Check if reasoning is supported.
74    #[must_use]
75    pub const fn has_reasoning(&self) -> bool {
76        self.flags.contains(CapabilityFlags::REASONING)
77    }
78
79    /// Check if tool calling is supported.
80    #[must_use]
81    pub const fn has_tool_calling(&self) -> bool {
82        self.flags.contains(CapabilityFlags::TOOL_CALLING)
83    }
84
85    /// Check if vision is supported.
86    #[must_use]
87    pub const fn has_vision(&self) -> bool {
88        self.flags.contains(CapabilityFlags::VISION)
89    }
90
91    /// Check if MTP (Multi-Token Prediction) draft heads are present.
92    #[must_use]
93    pub const fn has_mtp(&self) -> bool {
94        self.flags.contains(CapabilityFlags::MTP)
95    }
96
97    /// Check if this model produces embeddings rather than generated text.
98    #[must_use]
99    pub const fn has_embedding(&self) -> bool {
100        self.flags.contains(CapabilityFlags::EMBEDDING)
101    }
102
103    /// Convert capabilities to tag strings for model metadata.
104    ///
105    /// Returns tags like "reasoning", "agent" (for tool calling), etc.
106    #[must_use]
107    pub fn to_tags(&self) -> Vec<String> {
108        let mut tags = Vec::new();
109
110        // The producer side of `capability_tags`: every string here is read
111        // back somewhere by that module's constants, so both ends name the
112        // same thing rather than agreeing by convention.
113        if self.has_reasoning() {
114            tags.push(super::capability_tags::REASONING.to_string());
115        }
116        if self.has_tool_calling() {
117            // triggers --jinja auto-enable
118            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            // triggers --spec-type draft-mtp auto-enable
131            tags.push(super::capability_tags::MTP.to_string());
132        }
133        if self.has_embedding() {
134            // triggers --embeddings auto-enable
135            tags.push(super::capability_tags::EMBEDDING.to_string());
136        }
137
138        // Add extension tags
139        for ext in &self.extensions {
140            if !tags.contains(ext) {
141                tags.push(ext.clone());
142            }
143        }
144
145        tags
146    }
147}
148
149// =============================================================================
150// Metadata value types
151// =============================================================================
152
153/// GGUF metadata value types.
154///
155/// Represents all possible value types that can appear in GGUF metadata.
156#[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                // Limit array output to prevent massive tokenizer vocab dumps
190                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    /// Try to convert the value to a u64.
209    ///
210    /// Attempts to convert various numeric GGUF value types to u64.
211    /// Only converts non-negative values to avoid overflow issues.
212    #[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    /// Try to convert the value to a f64.
229    #[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    /// Try to get the value as a string reference.
248    #[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// =============================================================================
258// Metadata
259// =============================================================================
260
261/// Parsed metadata from a GGUF file.
262///
263/// This is the domain-facing type used by services and ports.
264/// Parsing logic that produces this type lives in `gglib-gguf`.
265#[derive(Debug, Clone, Default)]
266pub struct GgufMetadata {
267    /// Model architecture (e.g., "llama", "mistral").
268    pub architecture: Option<String>,
269    /// Quantization type (e.g., "`Q4_K_M`", "`Q8_0`").
270    pub quantization: Option<String>,
271    /// Number of parameters in billions.
272    pub param_count_b: Option<f64>,
273    /// Maximum context length.
274    pub context_length: Option<u64>,
275    /// Number of experts (for `MoE` models).
276    pub expert_count: Option<u32>,
277    /// Number of experts used during inference (for `MoE` models).
278    pub expert_used_count: Option<u32>,
279    /// Number of shared experts (for `MoE` models).
280    pub expert_shared_count: Option<u32>,
281    /// Additional key-value metadata from the file (string representation).
282    pub metadata: HashMap<String, String>,
283}
284
285/// Raw metadata from GGUF parsing (before string conversion).
286///
287/// Used internally by parsers; services typically use `GgufMetadata`.
288pub type RawMetadata = HashMap<String, GgufValue>;
289
290// =============================================================================
291// Tests
292// =============================================================================
293
294#[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}