1use std::fs::File;
2use std::io::{self, Read};
3use std::path::{Path, PathBuf};
4
5use sha2::{Digest, Sha256};
6use thiserror::Error;
7
8use crate::inference::{ModelCapabilities, ModelQualification};
9
10pub const NEOHORSE_V1_PROFILE_NAME: &str = "neohorse-1-9b-q8";
11
12pub const NEOHORSE_V1: QualifiedModelProfile = QualifiedModelProfile {
13 name: NEOHORSE_V1_PROFILE_NAME,
14 identifier: "neohorse-1-9b-q8",
15 repository: "TokenRhythm/NeoHorse-1-9B-GGUF",
16 revision: "ddcb4c939b5392c86a9d2733c7c0ed30db2554fd",
17 artifact: "NeoHorse-1-9B-Q8_0.gguf",
18 artifact_sha256: "519869730bda973ec50bb3ac42cd874569e5f3a57c3e0d5d23a2b5040e93310f",
19 quantization: "Q8_0",
20 chat_template: "artifact-metadata",
21 context_window_tokens: 8_192,
22 runtime: "Ollama-bundled ROCm HIP backend",
23 runtime_version: "ROCm 7.2 / llama-server 0.3.0-dev",
24 runtime_commit: "d222767c7",
25 accelerator: "AMD Radeon RX 7900 XTX",
26 architecture: "gfx1100",
27 mtp_enabled: false,
28};
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct QualifiedModelProfile {
32 pub name: &'static str,
33 pub identifier: &'static str,
34 pub repository: &'static str,
35 pub revision: &'static str,
36 pub artifact: &'static str,
37 pub artifact_sha256: &'static str,
38 pub quantization: &'static str,
39 pub chat_template: &'static str,
40 pub context_window_tokens: u32,
41 pub runtime: &'static str,
42 pub runtime_version: &'static str,
43 pub runtime_commit: &'static str,
44 pub accelerator: &'static str,
45 pub architecture: &'static str,
46 pub mtp_enabled: bool,
47}
48
49impl QualifiedModelProfile {
50 pub fn validate_artifact(self, path: impl AsRef<Path>) -> Result<String, ProfileError> {
51 let path = path.as_ref();
52 if !path.is_file() {
53 return Err(ProfileError::MissingArtifact(path.to_path_buf()));
54 }
55 let filename = path.file_name().and_then(|name| name.to_str());
56 if filename != Some(self.artifact) {
57 return Err(ProfileError::ArtifactName {
58 expected: self.artifact.to_owned(),
59 actual: filename.unwrap_or("<non-UTF-8>").to_owned(),
60 });
61 }
62 let actual = sha256_file(path)?;
63 if actual != self.artifact_sha256 {
64 return Err(ProfileError::ArtifactDigest {
65 expected: self.artifact_sha256.to_owned(),
66 actual,
67 });
68 }
69 Ok(actual)
70 }
71
72 pub fn capabilities(self, path: impl AsRef<Path>) -> Result<ModelCapabilities, ProfileError> {
73 let path = path.as_ref();
74 let digest = self.validate_artifact(path)?;
75 self.capabilities_from_verified_artifact(path, &digest)
76 }
77
78 pub(crate) fn capabilities_from_verified_artifact(
79 self,
80 path: &Path,
81 verified_sha256: &str,
82 ) -> Result<ModelCapabilities, ProfileError> {
83 let filename = path.file_name().and_then(|name| name.to_str());
84 if filename != Some(self.artifact) {
85 return Err(ProfileError::ArtifactName {
86 expected: self.artifact.to_owned(),
87 actual: filename.unwrap_or("<non-UTF-8>").to_owned(),
88 });
89 }
90 if !verified_sha256.eq_ignore_ascii_case(self.artifact_sha256) {
91 return Err(ProfileError::ArtifactDigest {
92 expected: self.artifact_sha256.to_owned(),
93 actual: verified_sha256.to_owned(),
94 });
95 }
96 Ok(ModelCapabilities {
97 identifier: self.identifier.to_owned(),
98 repository: Some(self.repository.to_owned()),
99 artifact: Some(path.display().to_string()),
100 artifact_sha256: Some(self.artifact_sha256.to_owned()),
101 quantization: Some(self.quantization.to_owned()),
102 chat_template: Some(self.chat_template.to_owned()),
103 context_window_tokens: Some(self.context_window_tokens),
104 native_tools: true,
105 qualification: Some(ModelQualification {
106 profile_name: self.name.to_owned(),
107 revision: self.revision.to_owned(),
108 expected_artifact: self.artifact.to_owned(),
109 runtime: self.runtime.to_owned(),
110 runtime_version: self.runtime_version.to_owned(),
111 runtime_commit: self.runtime_commit.to_owned(),
112 accelerator: self.accelerator.to_owned(),
113 architecture: self.architecture.to_owned(),
114 mtp_enabled: self.mtp_enabled,
115 artifact_validated: true,
116 }),
117 })
118 }
119}
120
121#[derive(Debug, Error)]
122pub enum ProfileError {
123 #[error("qualified model artifact is missing or not a file: {0}")]
124 MissingArtifact(PathBuf),
125 #[error("qualified model artifact name mismatch: expected {expected}, found {actual}")]
126 ArtifactName { expected: String, actual: String },
127 #[error("qualified model artifact SHA-256 mismatch: expected {expected}, found {actual}")]
128 ArtifactDigest { expected: String, actual: String },
129 #[error("qualified model artifact could not be read: {0}")]
130 Io(#[from] io::Error),
131}
132
133fn sha256_file(path: &Path) -> Result<String, io::Error> {
134 let mut file = File::open(path)?;
135 let mut digest = Sha256::new();
136 let mut buffer = [0_u8; 1024 * 1024];
137 loop {
138 let count = file.read(&mut buffer)?;
139 if count == 0 {
140 break;
141 }
142 digest.update(&buffer[..count]);
143 }
144 Ok(format!("{:x}", digest.finalize()))
145}
146
147#[cfg(test)]
148mod tests {
149 use std::fs;
150
151 use super::{ProfileError, QualifiedModelProfile};
152
153 const TEST_PROFILE: QualifiedModelProfile = QualifiedModelProfile {
154 name: "test",
155 identifier: "test",
156 repository: "example/test",
157 revision: "revision",
158 artifact: "model.gguf",
159 artifact_sha256: "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824",
160 quantization: "test",
161 chat_template: "test",
162 context_window_tokens: 8_192,
163 runtime: "test",
164 runtime_version: "test",
165 runtime_commit: "test",
166 accelerator: "test",
167 architecture: "test",
168 mtp_enabled: false,
169 };
170
171 #[test]
172 fn accepts_only_the_exact_named_artifact_and_digest() {
173 let directory = tempfile::tempdir().expect("tempdir");
174 let artifact = directory.path().join("model.gguf");
175 fs::write(&artifact, b"hello").expect("artifact");
176 let capabilities = TEST_PROFILE.capabilities(&artifact).expect("qualified");
177 assert_eq!(
178 capabilities.artifact_sha256.as_deref(),
179 Some(TEST_PROFILE.artifact_sha256)
180 );
181 assert!(
182 capabilities
183 .qualification
184 .expect("qualification")
185 .artifact_validated
186 );
187 }
188
189 #[test]
190 fn refuses_missing_or_mismatched_artifacts() {
191 let directory = tempfile::tempdir().expect("tempdir");
192 let missing = directory.path().join("model.gguf");
193 assert!(matches!(
194 TEST_PROFILE.validate_artifact(&missing),
195 Err(ProfileError::MissingArtifact(_))
196 ));
197 fs::write(&missing, b"wrong").expect("artifact");
198 assert!(matches!(
199 TEST_PROFILE.validate_artifact(&missing),
200 Err(ProfileError::ArtifactDigest { .. })
201 ));
202 let wrong_name = directory.path().join("other.gguf");
203 fs::write(&wrong_name, b"hello").expect("artifact");
204 assert!(matches!(
205 TEST_PROFILE.validate_artifact(&wrong_name),
206 Err(ProfileError::ArtifactName { .. })
207 ));
208 }
209}