use std::path::Path;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ModelSourceInfo {
path: Option<String>,
size_bytes: Option<u64>,
format: Option<String>,
quantization: Option<String>,
architecture: Option<String>,
context_length: Option<usize>,
model_max_context_length: Option<usize>,
content_hash: Option<String>,
parameter_count: Option<u64>,
}
#[must_use]
pub fn detect_format_from_magic(magic: &[u8]) -> Option<&'static str> {
if magic.starts_with(b"GGUF") {
return Some("gguf");
}
if magic.starts_with(&crate::apr::MAGIC_PREFIX) {
return Some("apr");
}
if magic.len() > 8 && magic[8] == b'{' {
return Some("safetensors");
}
None
}
impl ModelSourceInfo {
#[must_use]
pub fn from_path(path: &Path) -> Self {
let size_bytes = std::fs::metadata(path).ok().map(|m| m.len());
let format = read_magic(path)
.as_deref()
.and_then(detect_format_from_magic)
.map(str::to_string);
let path_str = std::fs::canonicalize(path)
.unwrap_or_else(|_| path.to_path_buf())
.to_string_lossy()
.into_owned();
Self {
path: Some(path_str),
size_bytes,
format,
..Self::default()
}
}
#[must_use]
pub fn with_quantization(mut self, quantization: impl Into<String>) -> Self {
self.quantization = Some(quantization.into());
self
}
#[must_use]
pub fn with_architecture(mut self, architecture: impl Into<String>) -> Self {
self.architecture = Some(architecture.into());
self
}
#[must_use]
pub fn with_context_length(mut self, context_length: usize) -> Self {
self.context_length = Some(context_length);
self
}
#[must_use]
pub fn with_model_max_context_length(mut self, context_length: usize) -> Self {
self.model_max_context_length = Some(context_length);
self
}
#[must_use]
pub fn with_content_hash(mut self, content_hash: impl Into<String>) -> Self {
self.content_hash = Some(content_hash.into());
self
}
#[must_use]
pub fn with_parameter_count(mut self, parameter_count: u64) -> Self {
self.parameter_count = Some(parameter_count);
self
}
#[must_use]
pub fn path(&self) -> Option<&str> {
self.path.as_deref()
}
#[must_use]
pub fn size_bytes(&self) -> Option<u64> {
self.size_bytes
}
#[must_use]
pub fn format(&self) -> Option<&str> {
self.format.as_deref()
}
#[must_use]
pub fn quantization(&self) -> Option<&str> {
self.quantization.as_deref()
}
#[must_use]
pub fn architecture(&self) -> Option<&str> {
self.architecture.as_deref()
}
#[must_use]
pub fn context_length(&self) -> Option<usize> {
self.context_length
}
#[must_use]
pub fn model_max_context_length(&self) -> Option<usize> {
self.model_max_context_length
}
#[must_use]
pub fn content_hash(&self) -> Option<&str> {
self.content_hash.as_deref()
}
#[must_use]
pub fn parameter_count(&self) -> Option<u64> {
self.parameter_count
}
}
fn read_magic(path: &Path) -> Option<Vec<u8>> {
use std::io::Read;
let mut file = std::fs::File::open(path).ok()?;
let mut buf = [0u8; 16];
let n = file.read(&mut buf).ok()?;
Some(buf[..n].to_vec())
}
#[must_use]
pub fn gguf_qtype_name(qtype: u32) -> Option<&'static str> {
Some(match qtype {
0 => "F32",
1 => "F16",
2 => "Q4_0",
3 => "Q4_1",
6 => "Q5_0",
7 => "Q5_1",
8 => "Q8_0",
9 => "Q8_1",
10 => "Q2_K",
11 => "Q3_K",
12 => "Q4_K",
13 => "Q5_K",
14 => "Q6_K",
15 => "Q8_K",
30 => "BF16",
_ => return None,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_source_knows_nothing_and_admits_it() {
let src = ModelSourceInfo::default();
assert_eq!(src.size_bytes(), None);
assert_eq!(src.format(), None);
assert_eq!(src.quantization(), None);
assert_eq!(src.context_length(), None);
assert_eq!(src.content_hash(), None, "never invent a content hash");
}
#[test]
fn detect_format_recognises_the_three_containers() {
assert_eq!(detect_format_from_magic(b"GGUF\0\0\0\0"), Some("gguf"));
assert_eq!(detect_format_from_magic(b"APR\0____"), Some("apr"));
assert_eq!(detect_format_from_magic(b"APRN____"), Some("apr"));
let st = [8u8, 0, 0, 0, 0, 0, 0, 0, b'{', b'"'];
assert_eq!(detect_format_from_magic(&st), Some("safetensors"));
}
#[test]
fn unrecognised_magic_is_unknown_not_gguf() {
assert_eq!(detect_format_from_magic(b"\x7fELF\0\0\0\0"), None);
assert_eq!(detect_format_from_magic(b""), None);
}
#[test]
fn from_path_measures_size_and_format_of_a_real_file() {
let dir = std::env::temp_dir().join(format!(
"apr-model-source-{}-{}",
std::process::id(),
line!()
));
std::fs::create_dir_all(&dir).expect("mkdir");
let path = dir.join("tiny.gguf");
std::fs::write(&path, b"GGUF\x03\0\0\0\0\0\0\0\0\0\0\0").expect("write");
let src = ModelSourceInfo::from_path(&path);
assert_eq!(
src.size_bytes(),
Some(16),
"size must be the real file size"
);
assert_eq!(src.format(), Some("gguf"));
assert!(src.path().is_some_and(|p| p.ends_with("tiny.gguf")));
assert_eq!(src.content_hash(), None);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn from_path_on_a_missing_file_reports_absence_not_zero() {
let src = ModelSourceInfo::from_path(Path::new("/nonexistent/model.gguf"));
assert_eq!(src.size_bytes(), None, "a missing file is unknown, not 0");
assert_eq!(src.format(), None);
}
#[test]
fn builders_record_measured_values() {
let src = ModelSourceInfo::default()
.with_quantization("Q4_K")
.with_architecture("qwen2")
.with_context_length(128)
.with_model_max_context_length(32768)
.with_content_hash("blake3:deadbeef")
.with_parameter_count(1_500_000_000);
assert_eq!(src.quantization(), Some("Q4_K"));
assert_eq!(src.architecture(), Some("qwen2"));
assert_eq!(src.context_length(), Some(128));
assert_eq!(src.model_max_context_length(), Some(32768));
assert_eq!(src.content_hash(), Some("blake3:deadbeef"));
assert_eq!(src.parameter_count(), Some(1_500_000_000));
}
#[test]
fn qtype_names_cover_the_shipped_quantizations() {
assert_eq!(gguf_qtype_name(12), Some("Q4_K"));
assert_eq!(gguf_qtype_name(14), Some("Q6_K"));
assert_eq!(gguf_qtype_name(0), Some("F32"));
assert_eq!(gguf_qtype_name(9999), None);
}
}