Skip to main content

kernel/records/
model_record.rs

1//! The central data model: a `ModelRecord` and the value types it composes.
2
3use 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/// Where a model's weights live and how it is identified.
16#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub struct ModelSource {
18    /// The kind of store the model came from.
19    pub kind: SourceKind,
20    /// The on-disk path (or endpoint identifier) of the model.
21    pub path: String,
22    /// The Hugging Face / Ollama repository, when applicable.
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub repo: Option<String>,
25    /// The revision/ref within the repository, when applicable.
26    #[serde(rename = "ref", default, skip_serializing_if = "Option::is_none")]
27    pub reference: Option<String>,
28}
29
30impl ModelSource {
31    /// Create a source with just a kind and path.
32    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    /// The identity string a stable id is derived from: `kind|path|repo`.
42    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/// How a model was resolved to its runtime.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
54#[serde(rename_all = "lowercase")]
55pub enum Resolution {
56    /// Chosen automatically by the runtime auction.
57    Auto,
58    /// Pinned by the user.
59    User,
60    /// Not yet resolved.
61    #[default]
62    Unresolved,
63}
64
65/// The runtime a model resolves to, plus how that choice was made.
66#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
67pub struct RuntimeRef {
68    /// The winning runtime, if resolved.
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub id: Option<RuntimeId>,
71    /// How the runtime was chosen.
72    #[serde(default)]
73    pub resolved: Resolution,
74    /// How much support the runtime needs.
75    #[serde(default)]
76    pub tier: RunTier,
77    /// Other runtimes that could also serve this model.
78    #[serde(default, skip_serializing_if = "Vec::is_empty")]
79    pub alternatives: Vec<RuntimeId>,
80    /// When the user confirmed the runtime, epoch millis.
81    #[serde(default, skip_serializing_if = "Option::is_none")]
82    pub confirmed_at: Option<i64>,
83}
84
85/// The type of a tunable parameter.
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
87#[serde(rename_all = "lowercase")]
88pub enum ParamType {
89    /// An integer.
90    Int,
91    /// A floating-point number.
92    Float,
93    /// A boolean.
94    Bool,
95    /// A free string.
96    String,
97    /// One of a fixed set of string values.
98    Enum,
99}
100
101/// The schema of one tunable parameter a model exposes.
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct ParamSpec {
104    /// The parameter's key, e.g. `temperature`.
105    pub key: String,
106    /// The parameter's type.
107    #[serde(rename = "type")]
108    pub param_type: ParamType,
109    /// The default value, if any.
110    #[serde(rename = "default", default, skip_serializing_if = "Option::is_none")]
111    pub default_value: Option<JsonValue>,
112    /// The `[min, max]` range for numeric parameters, if any.
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub range: Option<Vec<JsonValue>>,
115    /// The allowed values for an enum parameter, if any.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub values: Option<Vec<String>>,
118}
119
120/// A single model the kernel knows about: what it is, where it lives, how it runs,
121/// and how the user has configured it.
122#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
123pub struct ModelRecord {
124    /// Stable identity derived from the source (see [`stable_id`]).
125    pub id: String,
126    /// The model's display name.
127    pub name: String,
128    /// Its primary modality.
129    pub modality: Modality,
130    /// Everything it can be asked to do.
131    pub capabilities: Vec<Capability>,
132    /// Where its weights live.
133    pub source: ModelSource,
134    /// The runtime it resolves to.
135    #[serde(default)]
136    pub runtime: RuntimeRef,
137    /// The parameter schema it exposes.
138    #[serde(default)]
139    pub params: Vec<ParamSpec>,
140    /// User-set parameter values.
141    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
142    pub param_values: BTreeMap<String, JsonValue>,
143    /// A user-set system prompt override.
144    #[serde(default, skip_serializing_if = "Option::is_none")]
145    pub system_prompt: Option<String>,
146    /// A user-set display alias.
147    #[serde(default, skip_serializing_if = "Option::is_none")]
148    pub alias: Option<String>,
149    /// How the runtime delivers output.
150    #[serde(default)]
151    pub execution: ExecutionMode,
152    /// What the model's files take on disk, in bytes, as the store's scanner
153    /// measured them. The memory a run needs is estimated from it, not stored.
154    #[serde(default, skip_serializing_if = "Option::is_none")]
155    pub footprint_bytes: Option<i64>,
156    /// The whole-mebibyte figure a record written before sizes were exact
157    /// carried. Read so a shelf already on disk keeps its sizes through the
158    /// upgrade, folded into `footprint_bytes` when the registry loads, and
159    /// never written back.
160    #[serde(default, rename = "footprint_mb", skip_serializing)]
161    pub(crate) legacy_footprint_mb: Option<i64>,
162    /// The record's lifecycle state.
163    #[serde(default)]
164    pub state: ModelState,
165    /// When the record was first registered, epoch millis.
166    pub registered_at: i64,
167    /// The primary weight file, when one file dominates.
168    #[serde(default, skip_serializing_if = "Option::is_none")]
169    pub primary_weight_path: Option<String>,
170    /// The model's context window, when known.
171    #[serde(default, skip_serializing_if = "Option::is_none")]
172    pub context_length: Option<i64>,
173    /// Whether the model ships a chat template.
174    #[serde(default, skip_serializing_if = "Option::is_none")]
175    pub has_chat_template: Option<bool>,
176    /// Whether the model's chat template supports tool calling, when it could be
177    /// read. `None` means undetermined — treated as capable, gated by a request.
178    #[serde(default, skip_serializing_if = "Option::is_none")]
179    pub supports_tools: Option<bool>,
180    /// Stop tokens the model declares.
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub stop_tokens: Option<Vec<String>>,
183    /// Whether the model is still downloading.
184    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
185    pub downloading: bool,
186    /// A content fingerprint used to follow a model across moves.
187    #[serde(default, skip_serializing_if = "Option::is_none")]
188    pub content_fingerprint: Option<String>,
189}
190
191impl ModelRecord {
192    /// Create a record for a model, deriving its stable id from `source` and
193    /// stamping `registered_at` with the current time.
194    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    /// The name to show the user: the alias if set, otherwise the name.
227    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    /// The id this model has on the wire: what the gateway's model listings
235    /// advertise and what clients send back. One definition, so a launcher that
236    /// configures a client and the listing it will call cannot disagree. Same
237    /// rule as [`display_name`](Self::display_name), including the empty-alias
238    /// guard.
239    pub fn wire_id(&self) -> &str {
240        self.display_name()
241    }
242
243    /// The size on disk, when one is recorded and positive.
244    pub fn size_on_disk(&self) -> Option<i64> {
245        self.footprint_bytes.filter(|bytes| *bytes > 0)
246    }
247
248    /// Fold a size carried in whole mebibytes by a record written before sizes
249    /// were exact into the byte figure, so an upgraded shelf still shows a
250    /// size before its next scan measures one. The mebibyte figure is dropped
251    /// either way, so this runs once per record however often it is called.
252    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    /// The size on disk in whole mebibytes, the unit the memory governor
261    /// budgets in.
262    pub fn footprint_mib(&self) -> Option<i64> {
263        self.size_on_disk()
264            .map(|bytes| bytes / byte_format::BYTES_PER_MIB)
265    }
266
267    /// Whether the model can perform `capability`.
268    pub fn can(&self, capability: &Capability) -> bool {
269        self.capabilities.contains(capability)
270    }
271}
272
273/// The stable identity of a model: the first eight bytes of the SHA-256 of the
274/// source's `(kind, path, repo)` fields, hex-encoded. The fields are hashed with
275/// a length prefix so the mapping is injective — a `|` (or any byte) inside a
276/// path or repo cannot make two distinct sources collide. Two sources with the
277/// same kind, path, and repo map to the same id.
278pub 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}