1use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6use sha2::{Digest, Sha256};
7
8use crate::records::byte_format;
9use crate::records::identifiers::{
10 Capability, ExecutionMode, Modality, ModelState, RunTier, RuntimeId, SourceKind,
11};
12use crate::records::json_value::JsonValue;
13use crate::time::now_millis;
14
15#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub struct ModelSource {
18 pub kind: SourceKind,
20 pub path: String,
22 #[serde(default, skip_serializing_if = "Option::is_none")]
24 pub repo: Option<String>,
25 #[serde(rename = "ref", default, skip_serializing_if = "Option::is_none")]
27 pub reference: Option<String>,
28}
29
30impl ModelSource {
31 pub fn new(kind: SourceKind, path: &str) -> Self {
33 Self {
34 kind,
35 path: path.to_owned(),
36 repo: None,
37 reference: None,
38 }
39 }
40
41 pub fn identity(&self) -> String {
43 format!(
44 "{}|{}|{}",
45 self.kind.as_str(),
46 self.path,
47 self.repo.as_deref().unwrap_or("")
48 )
49 }
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
54#[serde(rename_all = "lowercase")]
55pub enum Resolution {
56 Auto,
58 User,
60 #[default]
62 Unresolved,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
67pub struct RuntimeRef {
68 #[serde(default, skip_serializing_if = "Option::is_none")]
70 pub id: Option<RuntimeId>,
71 #[serde(default)]
73 pub resolved: Resolution,
74 #[serde(default)]
76 pub tier: RunTier,
77 #[serde(default, skip_serializing_if = "Vec::is_empty")]
79 pub alternatives: Vec<RuntimeId>,
80 #[serde(default, skip_serializing_if = "Option::is_none")]
82 pub confirmed_at: Option<i64>,
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
87#[serde(rename_all = "lowercase")]
88pub enum ParamType {
89 Int,
91 Float,
93 Bool,
95 String,
97 Enum,
99}
100
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct ParamSpec {
104 pub key: String,
106 #[serde(rename = "type")]
108 pub param_type: ParamType,
109 #[serde(rename = "default", default, skip_serializing_if = "Option::is_none")]
111 pub default_value: Option<JsonValue>,
112 #[serde(default, skip_serializing_if = "Option::is_none")]
114 pub range: Option<Vec<JsonValue>>,
115 #[serde(default, skip_serializing_if = "Option::is_none")]
117 pub values: Option<Vec<String>>,
118}
119
120#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
123pub struct ModelRecord {
124 pub id: String,
126 pub name: String,
128 pub modality: Modality,
130 pub capabilities: Vec<Capability>,
132 pub source: ModelSource,
134 #[serde(default)]
136 pub runtime: RuntimeRef,
137 #[serde(default)]
139 pub params: Vec<ParamSpec>,
140 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
142 pub param_values: BTreeMap<String, JsonValue>,
143 #[serde(default, skip_serializing_if = "Option::is_none")]
145 pub system_prompt: Option<String>,
146 #[serde(default, skip_serializing_if = "Option::is_none")]
148 pub alias: Option<String>,
149 #[serde(default)]
151 pub execution: ExecutionMode,
152 #[serde(default, skip_serializing_if = "Option::is_none")]
155 pub footprint_bytes: Option<i64>,
156 #[serde(default, rename = "footprint_mb", skip_serializing)]
161 pub(crate) legacy_footprint_mb: Option<i64>,
162 #[serde(default)]
164 pub state: ModelState,
165 pub registered_at: i64,
167 #[serde(default, skip_serializing_if = "Option::is_none")]
169 pub primary_weight_path: Option<String>,
170 #[serde(default, skip_serializing_if = "Option::is_none")]
172 pub context_length: Option<i64>,
173 #[serde(default, skip_serializing_if = "Option::is_none")]
175 pub has_chat_template: Option<bool>,
176 #[serde(default, skip_serializing_if = "Option::is_none")]
179 pub supports_tools: Option<bool>,
180 #[serde(default, skip_serializing_if = "Option::is_none")]
182 pub stop_tokens: Option<Vec<String>>,
183 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
185 pub downloading: bool,
186 #[serde(default, skip_serializing_if = "Option::is_none")]
188 pub content_fingerprint: Option<String>,
189}
190
191impl ModelRecord {
192 pub fn new(
195 name: &str,
196 modality: Modality,
197 capabilities: Vec<Capability>,
198 source: ModelSource,
199 ) -> Self {
200 Self {
201 id: stable_id(&source),
202 name: name.to_owned(),
203 modality,
204 capabilities,
205 source,
206 runtime: RuntimeRef::default(),
207 params: Vec::new(),
208 param_values: BTreeMap::new(),
209 system_prompt: None,
210 alias: None,
211 execution: ExecutionMode::Sync,
212 footprint_bytes: None,
213 legacy_footprint_mb: None,
214 state: ModelState::Unresolved,
215 registered_at: now_millis(),
216 primary_weight_path: None,
217 context_length: None,
218 has_chat_template: None,
219 supports_tools: None,
220 stop_tokens: None,
221 downloading: false,
222 content_fingerprint: None,
223 }
224 }
225
226 pub fn display_name(&self) -> &str {
228 match &self.alias {
229 Some(alias) if !alias.is_empty() => alias,
230 _ => &self.name,
231 }
232 }
233
234 pub fn wire_id(&self) -> &str {
240 self.display_name()
241 }
242
243 pub fn size_on_disk(&self) -> Option<i64> {
245 self.footprint_bytes.filter(|bytes| *bytes > 0)
246 }
247
248 pub(crate) fn adopt_legacy_footprint(&mut self) {
253 if let Some(mebibytes) = self.legacy_footprint_mb.take()
254 && self.footprint_bytes.is_none()
255 {
256 self.footprint_bytes = Some(mebibytes.saturating_mul(byte_format::BYTES_PER_MIB));
257 }
258 }
259
260 pub fn footprint_mib(&self) -> Option<i64> {
263 self.size_on_disk()
264 .map(|bytes| bytes / byte_format::BYTES_PER_MIB)
265 }
266
267 pub fn can(&self, capability: &Capability) -> bool {
269 self.capabilities.contains(capability)
270 }
271}
272
273pub fn stable_id(source: &ModelSource) -> String {
279 let mut hasher = Sha256::new();
280 for field in [
281 source.kind.as_str(),
282 source.path.as_str(),
283 source.repo.as_deref().unwrap_or(""),
284 ] {
285 hasher.update((field.len() as u64).to_le_bytes());
286 hasher.update(field.as_bytes());
287 }
288 let digest = hasher.finalize();
289 hex::encode(&digest[..8])
290}