1use serde_json::Value;
4
5use crate::DaemonError;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum Cap<'a> {
9 Yes,
10 No(&'a str),
11 Unknown,
12}
13
14#[derive(Debug, Clone, Default, PartialEq)]
15pub struct Capabilities {
16 record: Option<Value>,
17}
18
19impl Capabilities {
20 pub fn parse(json: &str) -> Capabilities {
21 Capabilities { record: serde_json::from_str::<Value>(json).ok().filter(Value::is_object) }
22 }
23
24 pub fn record(&self) -> Option<&Value> {
25 self.record.as_ref()
26 }
27
28 pub fn get(&self, section: &str, key: &str) -> Cap<'_> {
29 match self.record.as_ref().and_then(|r| r.get(section)).and_then(|s| s.get(key)) {
30 Some(Value::Bool(true)) => Cap::Yes,
31 Some(Value::String(why)) => Cap::No(why),
32 Some(Value::Bool(false)) => Cap::No("not supported"),
33 _ => Cap::Unknown,
34 }
35 }
36
37 pub fn require(&self, section: &str, key: &str, code: &'static str, what: &str) -> Result<(), DaemonError> {
38 match self.get(section, key) {
39 Cap::No(why) => Err(DaemonError::Unsupported(crate::Refusal::new(code, None, format!("{what}: {why}")))),
40 Cap::Yes | Cap::Unknown => Ok(()),
41 }
42 }
43
44 pub fn limit(&self, key: &str) -> Option<u64> {
45 self.record.as_ref()?.get("limits")?.get(key)?.as_u64().filter(|n| *n > 0)
46 }
47
48 pub fn changes(&self, other: &Capabilities) -> Vec<String> {
49 fn walk(path: &str, a: &Value, b: &Value, out: &mut Vec<String>) {
50 match (a, b) {
51 (Value::Object(x), Value::Object(y)) => {
52 let mut keys: Vec<&String> = x.keys().chain(y.keys()).collect();
53 keys.sort();
54 keys.dedup();
55 for k in keys {
56 if path.is_empty() && k == "limits" {
57 continue;
58 }
59 let p = if path.is_empty() { k.clone() } else { format!("{path}.{k}") };
60 walk(&p, x.get(k).unwrap_or(&Value::Null), y.get(k).unwrap_or(&Value::Null), out);
61 }
62 }
63 _ if a != b => out.push(path.to_string()),
64 _ => {}
65 }
66 }
67 let mut out = Vec::new();
68 match (&self.record, &other.record) {
69 (Some(a), Some(b)) => walk("", a, b, &mut out),
70 (None, None) => {}
71 _ => out.push("the whole record".to_string()),
72 }
73 out
74 }
75
76 pub fn with_runtime(&self, picked: Option<&str>) -> Option<Value> {
77 let mut r = self.record.clone()?;
78 if let (Some(id), Some(obj)) = (picked, r.as_object_mut()) {
79 obj.entry("runtime").or_insert_with(|| serde_json::json!({ "id": id }));
80 }
81 Some(r)
82 }
83
84 pub fn check_startup_flags(&self, park_lossy: bool, kv_bits: Option<i32>, speculate: Option<&str>) -> Result<(), String> {
85 let generates = !matches!(self.get("workload", "causal_generation"), Cap::No(_));
86 if park_lossy && generates {
87 if let Cap::No(why) = self.get("serving", "park_lossy") {
88 return Err(format!("--park-lossy: {why}"));
89 }
90 }
91 if let Some(bits) = kv_bits.filter(|b| *b != 0) {
92 if let Cap::No(why) = self.get("load", "kv_bits") {
93 return Err(format!("--kv-bits {bits}: {why}"));
94 }
95 }
96 if let Some(d) = speculate.filter(|d| !matches!(d.trim(), "off" | "auto")).filter(|_| generates) {
97 if let Cap::No(why) = self.get("serving", "multi_row_verification") {
98 return Err(format!("--speculate {d}: {why}"));
99 }
100 }
101 Ok(())
102 }
103
104 pub fn declines_auto_speculation(&self) -> Option<&str> {
105 match self.get("serving", "multi_row_verification") {
106 Cap::No(why) => Some(why),
107 Cap::Yes | Cap::Unknown => None,
108 }
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115
116 const EXECUTOR: &str = r#"{"descriptor_version":1,"runtime":{"id":"llamacpp","version":"v"},
117 "load":{"kv_bits":"no such knob"},
118 "modalities":{"image_encode":"no media path","gemma_audio_encode":"no media path"},
119 "serving":{"park_lossy":"no lossy encoding","multi_row_verification":"no strategy","park_lossless":true}}"#;
120
121 #[test]
122 fn a_leaf_is_yes_a_reason_or_unknown() {
123 let c = Capabilities::parse(EXECUTOR);
124 assert_eq!(c.get("serving", "park_lossless"), Cap::Yes);
125 assert_eq!(c.get("modalities", "image_encode"), Cap::No("no media path"));
126 assert_eq!(c.get("modalities", "whisper_transcribe"), Cap::Unknown);
127 assert_eq!(Capabilities::parse("").get("serving", "park_lossy"), Cap::Unknown);
128 assert_eq!(Capabilities::parse("[1]").record(), None, "only an object is a record");
129 }
130
131 #[test]
132 fn require_refuses_in_the_runtimes_words_and_passes_the_unknown() {
133 let c = Capabilities::parse(EXECUTOR);
134 let e = c.require("modalities", "image_encode", "unsupported_input", "this model does not accept images").unwrap_err();
135 assert!(
136 matches!(&e, DaemonError::Unsupported(r) if r.message == "this model does not accept images: no media path" && r.code == "unsupported_input"),
137 "{e:?}"
138 );
139 assert!(c.require("modalities", "whisper_transcribe", "unsupported_model", "x").is_ok(), "an undescribed key is not a refusal");
140 assert!(Capabilities::default().require("modalities", "image_encode", "unsupported_input", "x").is_ok());
141 }
142
143 #[test]
144 fn startup_flags_are_checked_against_the_record() {
145 let c = Capabilities::parse(EXECUTOR);
146 assert_eq!(c.check_startup_flags(true, None, None).unwrap_err(), "--park-lossy: no lossy encoding");
147 assert_eq!(c.check_startup_flags(false, Some(8), None).unwrap_err(), "--kv-bits 8: no such knob");
148 assert_eq!(c.check_startup_flags(false, None, Some("prompt-lookup")).unwrap_err(), "--speculate prompt-lookup: no strategy");
149 assert!(c.check_startup_flags(false, Some(0), Some("off")).is_ok(), "0 and off are no request");
150 let native = Capabilities::parse(r#"{"descriptor_version":1,"serving":{"park_lossy":true,"multi_row_verification":true}}"#);
151 assert!(native.check_startup_flags(true, Some(8), Some("prompt-lookup")).is_ok());
152 }
153
154 #[test]
155 fn auto_speculation_is_a_preference_and_a_model_that_generates_nothing_takes_no_generation_flags() {
156 let hybrid = Capabilities::parse(
157 r#"{"descriptor_version":1,"workload":{"causal_generation":true},
158 "serving":{"park_lossy":true,"multi_row_verification":"Mamba-2 hybrid: no speculative verification"}}"#,
159 );
160 assert_eq!(hybrid.check_startup_flags(false, None, Some("auto")), Ok(()));
161 assert_eq!(hybrid.check_startup_flags(false, None, Some(" auto ")), Ok(()), "as the resolver reads it, trimmed");
162 assert_eq!(
163 hybrid.check_startup_flags(false, None, Some("prompt-lookup")).unwrap_err(),
164 "--speculate prompt-lookup: Mamba-2 hybrid: no speculative verification"
165 );
166 assert_eq!(hybrid.declines_auto_speculation(), Some("Mamba-2 hybrid: no speculative verification"));
167 let whisper = Capabilities::parse(
168 r#"{"descriptor_version":1,"workload":{"causal_generation":"speech encoder-decoder, no autoregressive text API"},
169 "serving":{"park_lossy":"no autoregressive decode","multi_row_verification":"no autoregressive decode"}}"#,
170 );
171 assert_eq!(whisper.check_startup_flags(true, None, Some("auto")), Ok(()));
172 assert_eq!(whisper.check_startup_flags(true, None, Some("dflash:/d.base")), Ok(()));
173 let c = Capabilities::parse(EXECUTOR);
174 assert_eq!(c.check_startup_flags(true, None, Some("auto")).unwrap_err(), "--park-lossy: no lossy encoding");
175 assert_eq!(c.check_startup_flags(false, None, Some("auto")), Ok(()));
176 assert_eq!(c.declines_auto_speculation(), Some("no strategy"));
177 assert_eq!(Capabilities::default().declines_auto_speculation(), None, "an engine that describes nothing speculates as before");
178 }
179
180 #[test]
181 fn a_respawn_that_resizes_is_the_same_model_one_that_upgrades_is_not() {
182 let a = Capabilities::parse(r#"{"runtime":{"id":"llamacpp","version":"1"},"serving":{"park_lossy":"no"},"limits":{"max_seq_len":8192,"cells_total":65536}}"#);
183 let resized = Capabilities::parse(r#"{"runtime":{"id":"llamacpp","version":"1"},"serving":{"park_lossy":"no"},"limits":{"max_seq_len":4096,"cells_total":32768}}"#);
184 assert!(a.changes(&resized).is_empty(), "limits follow the memory free at load");
185 assert_eq!((a.limit("max_seq_len"), resized.limit("cells_total")), (Some(8192), Some(32768)));
186 assert_eq!(a.limit("max_seqs"), None);
187 let upgraded = Capabilities::parse(r#"{"runtime":{"id":"llamacpp","version":"2"},"serving":{"park_lossy":true},"limits":{"max_seq_len":8192}}"#);
188 assert_eq!(a.changes(&upgraded), ["runtime.version", "serving.park_lossy"]);
189 assert_eq!(Capabilities::default().changes(&Capabilities::default()), Vec::<String>::new());
190 assert_eq!(Capabilities::default().changes(&a), ["the whole record"], "a record where there was none");
191 }
192
193 #[test]
194 fn the_runtime_is_named_once() {
195 let exec = Capabilities::parse(EXECUTOR).with_runtime(Some("basert")).unwrap();
196 assert_eq!(exec["runtime"]["id"], "llamacpp", "the record's own runtime wins");
197 let native = Capabilities::parse(r#"{"descriptor_version":1}"#).with_runtime(Some("basert")).unwrap();
198 assert_eq!(native["runtime"], serde_json::json!({"id": "basert"}));
199 assert_eq!(Capabilities::default().with_runtime(Some("basert")), None);
200 }
201}