use std::collections::BTreeSet;
use crate::ontology::extract::cli_surface::{Arg, Command, Surface, UNSET};
use crate::ontology::extract::covering::{self, Pair};
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum CellKind {
Matrix,
Probe,
Base,
Effect {
arg: String,
level: String,
base: String,
},
Sampling { control: String },
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct CellSpec {
pub id: String,
pub host: String,
pub command: String,
pub generates: bool,
pub model_sha256: Option<String>,
pub model_file: Option<String>,
pub args: Vec<(String, String)>,
pub shape: Option<String>,
pub thinking: Option<String>,
pub rung: Option<String>,
pub rung_tokens: Option<u64>,
#[serde(flatten)]
pub kind: CellKind,
}
#[derive(Debug, Clone)]
pub struct ModelRef<'a> {
pub sha: &'a str,
pub file: &'a str,
pub arch: Option<&'a str>,
pub modes: Vec<&'static str>,
pub rungs: Vec<(&'a str, Option<u64>)>,
}
#[derive(Debug, Clone)]
pub struct HostModels<'a> {
pub host: &'a str,
pub models: Vec<ModelRef<'a>>,
}
struct Factor {
kind: FactorKind,
levels: Vec<String>,
}
#[derive(Clone, PartialEq, Eq)]
enum FactorKind {
Arg(String),
Shape,
Thinking,
Rung,
}
#[must_use]
pub fn derive(surface: &Surface, hosts: &[HostModels<'_>]) -> Vec<CellSpec> {
let mut out = Vec::new();
for c in surface.commands.iter().filter(|c| c.leaf) {
for h in hosts {
if c.takes_model() {
for m in &h.models {
out.extend(matrix(surface, c, h.host, m));
}
out.extend(effects(surface, c, h));
out.extend(sampling(c, h));
} else {
out.push(probe(c, h.host));
}
}
}
out
}
fn mode_factors<'a>(surface: &'a Surface, c: &'a Command) -> Vec<(&'a Arg, Factor)> {
let mut seen = BTreeSet::new();
c.mode_args(&surface.global_args)
.filter(|a| seen.insert(a.spelling()))
.map(|a| {
(
a,
Factor {
kind: FactorKind::Arg(a.spelling()),
levels: a.levels(),
},
)
})
.collect()
}
fn forbidden(args: &[&Arg], factors: &[Factor]) -> BTreeSet<Pair> {
let mut out = BTreeSet::new();
for (i, a) in args.iter().enumerate() {
for (j, b) in args.iter().enumerate().skip(i + 1) {
let clash = a.conflicts_with.contains(&b.id) || b.conflicts_with.contains(&a.id);
if !clash {
continue;
}
for la in 1..factors[i].levels.len() {
for lb in 1..factors[j].levels.len() {
out.insert((i, la, j, lb));
}
}
}
}
out
}
fn matrix(surface: &Surface, c: &Command, host: &str, m: &ModelRef<'_>) -> Vec<CellSpec> {
let modes = mode_factors(surface, c);
let arg_refs: Vec<&Arg> = modes.iter().map(|(a, _)| *a).collect();
let mut factors: Vec<Factor> = modes.into_iter().map(|(_, f)| f).collect();
let fb = forbidden(&arg_refs, &factors);
let shapes: Vec<String> = c.input_shapes().iter().map(|a| a.spelling()).collect();
if !shapes.is_empty() {
factors.push(Factor {
kind: FactorKind::Shape,
levels: shapes,
});
}
let generates = c.generates == Some(true);
if generates {
factors.push(Factor {
kind: FactorKind::Thinking,
levels: m.modes.iter().map(|t| format!("think-{t}")).collect(),
});
if !m.rungs.is_empty() {
factors.push(Factor {
kind: FactorKind::Rung,
levels: m.rungs.iter().map(|(r, _)| (*r).to_string()).collect(),
});
}
}
let levels: Vec<usize> = factors.iter().map(|f| f.levels.len()).collect();
covering::pairwise(&levels, &fb)
.into_iter()
.map(|row| cell_from_row(c, host, m, &factors, &row, CellKind::Matrix))
.collect()
}
fn cell_from_row(
c: &Command,
host: &str,
m: &ModelRef<'_>,
factors: &[Factor],
row: &[usize],
kind: CellKind,
) -> CellSpec {
let mut spec = CellSpec {
id: String::new(),
host: host.to_string(),
command: c.key.clone(),
generates: c.generates == Some(true),
model_sha256: Some(m.sha.to_string()),
model_file: Some(m.file.to_string()),
args: Vec::new(),
shape: None,
thinking: None,
rung: None,
rung_tokens: None,
kind,
};
for (f, &l) in factors.iter().zip(row) {
let level = f.levels[l].clone();
match &f.kind {
FactorKind::Arg(sp) if level != UNSET => spec.args.push((sp.clone(), level)),
FactorKind::Arg(_) => {}
FactorKind::Shape => spec.shape = Some(level),
FactorKind::Thinking => {
spec.thinking = Some(level.trim_start_matches("think-").to_string())
}
FactorKind::Rung => {
spec.rung_tokens = m
.rungs
.iter()
.find(|(r, _)| *r == level)
.and_then(|(_, t)| *t);
spec.rung = Some(level);
}
}
}
spec.id = cell_id(&spec);
spec
}
#[must_use]
pub fn cell_id(s: &CellSpec) -> String {
let args = if s.args.is_empty() {
"defaults".to_string()
} else {
s.args
.iter()
.map(|(a, l)| {
if l == "true" {
a.clone()
} else {
format!("{a}={l}")
}
})
.collect::<Vec<_>>()
.join("+")
};
let prefix = match &s.kind {
CellKind::Base => "base:",
CellKind::Effect { .. } => "effect:",
CellKind::Sampling { .. } => "sampling:",
CellKind::Matrix | CellKind::Probe => "",
};
[
format!("{prefix}{}", s.command),
s.host.clone(),
s.model_file.clone().unwrap_or_else(|| "-".into()),
s.thinking
.as_ref()
.map_or_else(|| "-".into(), |t| format!("think-{t}")),
s.rung.clone().unwrap_or_else(|| "-".into()),
s.shape.clone().unwrap_or_else(|| "-".into()),
args,
]
.join("/")
}
fn probe(c: &Command, host: &str) -> CellSpec {
let mut spec = CellSpec {
id: String::new(),
host: host.to_string(),
command: c.key.clone(),
generates: false,
model_sha256: None,
model_file: None,
args: Vec::new(),
shape: None,
thinking: None,
rung: None,
rung_tokens: None,
kind: CellKind::Probe,
};
spec.id = cell_id(&spec);
spec
}
fn effects(surface: &Surface, c: &Command, h: &HostModels<'_>) -> Vec<CellSpec> {
let mut archs = BTreeSet::new();
let reps = h
.models
.iter()
.filter(|m| archs.insert(m.arch.unwrap_or("-")));
let modes = mode_factors(surface, c);
let mut out = Vec::new();
for m in reps {
let mut base = cell_from_row(c, h.host, m, &[], &[], CellKind::Base);
base.shape = c.input_shapes().first().map(|a| a.spelling());
if c.generates == Some(true) {
let off = m
.modes
.iter()
.find(|t| **t == "off")
.or_else(|| m.modes.first());
base.thinking = off.map(|t| (*t).to_string());
if let Some((r, t)) = m.rungs.first() {
base.rung = Some((*r).to_string());
base.rung_tokens = *t;
}
}
base.id = cell_id(&base);
for (_, f) in &modes {
let FactorKind::Arg(sp) = &f.kind else {
continue;
};
for level in f.levels.iter().filter(|l| l.as_str() != UNSET) {
let mut e = base.clone();
e.args = vec![(sp.clone(), level.clone())];
e.kind = CellKind::Effect {
arg: sp.clone(),
level: level.clone(),
base: base.id.clone(),
};
e.id = cell_id(&e);
out.push(e);
}
}
out.push(base);
}
out
}
pub const SAMPLING_CONTROLS: [(&str, &[(&str, &str)]); 5] = [
("t0", &[("temperature", "0")]),
("topk1", &[("top_k", "1")]),
("seed-a", &[("temperature", "0.8"), ("seed", "1")]),
("seed-a-again", &[("temperature", "0.8"), ("seed", "1")]),
("seed-b", &[("temperature", "0.8"), ("seed", "2")]),
];
pub const SAMPLING_KINDS_NEEDED: [&str; 3] = ["temperature", "top_k", "seed"];
fn sampling(c: &Command, h: &HostModels<'_>) -> Vec<CellSpec> {
if c.generates != Some(true) {
return Vec::new();
}
let knob = |kind: &str| {
c.args
.iter()
.find(|a| a.role == "sampling" && a.sampling_kind.as_deref() == Some(kind))
.map(Arg::spelling)
};
if !c.args.iter().any(|a| a.role == "sampling") {
return Vec::new();
}
let mut archs = BTreeSet::new();
let reps = h
.models
.iter()
.filter(|m| archs.insert(m.arch.unwrap_or("-")));
let mut out = Vec::new();
for m in reps {
for (control, knobs) in SAMPLING_CONTROLS {
let set: Option<Vec<(String, String)>> = knobs
.iter()
.map(|(k, v)| knob(k).map(|sp| (sp, (*v).to_string())))
.collect();
let Some(args) = set else { continue };
let mut cell = cell_from_row(
c,
h.host,
m,
&[],
&[],
CellKind::Sampling {
control: control.to_string(),
},
);
cell.shape = c.input_shapes().first().map(|a| a.spelling());
let off = m
.modes
.iter()
.find(|t| **t == "off")
.or_else(|| m.modes.first());
cell.thinking = off.map(|t| (*t).to_string());
if let Some((r, t)) = m.rungs.first() {
cell.rung = Some((*r).to_string());
cell.rung_tokens = *t;
}
cell.args = args;
cell.id = format!("{}#{control}", cell_id(&cell));
out.push(cell);
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ontology::extract::cli_surface;
const SURFACE: &str = r#"{"schema":"apr-cli-surface/v1","binary":{"version":"0.0.0","git_sha":"t"},
"global_args":[{"id":"json","long":"json","value_type":"flag","values":["true","false"],"role":"mode"}],
"commands":[
{"path":["gen"],"key":"gen","leaf":true,"generates":true,"args":[
{"id":"model","positional":true,"required":true,"value_type":"path","role":"model"},
{"id":"prompt","long":"prompt","value_type":"text","role":"prompt"},
{"id":"input","short":"i","value_type":"path","role":"input-file"},
{"id":"chat","long":"chat","value_type":"flag","values":["true","false"],"role":"mode","conflicts_with":["raw"]},
{"id":"raw","long":"raw","value_type":"flag","values":["true","false"],"role":"mode","conflicts_with":["chat"]},
{"id":"format","long":"format","value_type":"enum","values":["text","json"],"role":"mode"}]},
{"path":["look"],"key":"look","leaf":true,"generates":false,"args":[
{"id":"file","positional":true,"required":true,"value_type":"path","role":"model"}]},
{"path":["group"],"key":"group","leaf":false,"args":[]},
{"path":["list"],"key":"list","leaf":true,"args":[]}]}"#;
fn hosts() -> Vec<HostModels<'static>> {
let m = |sha, file, arch| ModelRef {
sha,
file,
arch: Some(arch),
modes: vec!["on", "off"],
rungs: vec![("4k", Some(4096)), ("declared", Some(10_000))],
};
vec![HostModels {
host: "h1",
models: vec![m("aa", "a.gguf", "qwen35"), m("bb", "b.gguf", "qwen35")],
}]
}
#[test]
fn the_matrix_covers_every_pair_including_thinking_by_rung_and_never_a_conflict() {
let s = cli_surface::parse("t.json", SURFACE).expect("parses");
let cells = derive(&s, &hosts());
let gen_a: Vec<&CellSpec> = cells
.iter()
.filter(|c| {
c.command == "gen"
&& c.kind == CellKind::Matrix
&& c.model_file.as_deref() == Some("a.gguf")
})
.collect();
for t in ["on", "off"] {
for r in ["4k", "declared"] {
assert!(
gen_a
.iter()
.any(|c| c.thinking.as_deref() == Some(t) && c.rung.as_deref() == Some(r)),
"(think-{t}, {r}) is covered"
);
}
}
for sh in ["--prompt", "-i"] {
assert!(gen_a.iter().any(|c| c.shape.as_deref() == Some(sh)), "{sh}");
assert!(
gen_a
.iter()
.any(|c| c.shape.as_deref() == Some(sh)
&& c.args.iter().any(|(a, _)| a == "--chat")),
"--chat with {sh}: the #3743 pair (run --prompt --chat) is a cell"
);
}
assert!(
!gen_a
.iter()
.any(|c| c.args.iter().any(|(a, _)| a == "--chat")
&& c.args.iter().any(|(a, _)| a == "--raw")),
"conflicts_with pairs never share a cell"
);
let decl = gen_a
.iter()
.find(|c| c.rung.as_deref() == Some("declared"))
.expect("declared");
assert_eq!(decl.rung_tokens, Some(10_000), "the model's own length");
}
#[test]
fn a_non_generating_model_command_owes_no_thinking_or_rung_and_a_modelless_one_owes_a_probe() {
let s = cli_surface::parse("t.json", SURFACE).expect("parses");
let cells = derive(&s, &hosts());
let look: Vec<&CellSpec> = cells
.iter()
.filter(|c| c.command == "look" && c.kind == CellKind::Matrix)
.collect();
assert!(!look.is_empty());
assert!(look
.iter()
.all(|c| c.thinking.is_none() && c.rung.is_none()));
let list: Vec<&CellSpec> = cells.iter().filter(|c| c.command == "list").collect();
assert_eq!(list.len(), 1, "one probe per host");
assert_eq!(list[0].kind, CellKind::Probe);
assert_eq!(list[0].id, "list/h1/-/-/-/-/defaults");
assert!(
cells.iter().all(|c| c.command != "group"),
"a non-leaf owes nothing itself"
);
}
#[test]
fn effect_cells_flip_exactly_one_arg_against_one_base_per_architecture() {
let s = cli_surface::parse("t.json", SURFACE).expect("parses");
let cells = derive(&s, &hosts());
let bases: Vec<&CellSpec> = cells
.iter()
.filter(|c| c.command == "gen" && c.kind == CellKind::Base)
.collect();
assert_eq!(
bases.len(),
1,
"one qwen35 representative (a.gguf), not one per model"
);
assert!(bases[0].args.is_empty());
let effects: Vec<&CellSpec> = cells
.iter()
.filter(|c| c.command == "gen" && matches!(c.kind, CellKind::Effect { .. }))
.collect();
assert_eq!(effects.len(), 5);
assert!(effects.iter().all(|e| e.args.len() == 1));
assert_eq!(derive(&s, &hosts()), cells, "deterministic (R-15)");
let mut tiny = hosts();
tiny[0].models[1].rungs.clear();
let tiny_cells = derive(&s, &tiny);
assert!(
tiny_cells
.iter()
.any(|c| c.model_file.as_deref() == Some("b.gguf") && c.command == "gen"),
"a model owing no rung keeps its cells (a zero-level factor must not empty the array)"
);
let ids: BTreeSet<&str> = cells.iter().map(|c| c.id.as_str()).collect();
assert_eq!(ids.len(), cells.len(), "every id names one cell");
}
const KNOBS: &str = r#"{"schema":"apr-cli-surface/v1","binary":{"version":"0.0.0","git_sha":"t"},
"global_args":[],
"commands":[
{"path":["samp"],"key":"samp","leaf":true,"generates":true,"args":[
{"id":"model","positional":true,"required":true,"value_type":"path","role":"model"},
{"id":"decoy","long":"decoy","value_type":"text","role":"mode","sampling_kind":"temperature"},
{"id":"topk","long":"top-k","value_type":"int","role":"sampling","sampling_kind":"top_k"},
{"id":"temp","long":"temp","value_type":"float","role":"sampling","sampling_kind":"temperature"},
{"id":"seed","long":"seed","value_type":"int","role":"sampling","sampling_kind":"seed"}]},
{"path":["only"],"key":"only","leaf":true,"generates":true,"args":[
{"id":"topk","long":"top-k","value_type":"int","role":"sampling","sampling_kind":"top_k"},
{"id":"temp","long":"temp","value_type":"float","role":"sampling","sampling_kind":"temperature"},
{"id":"seed","long":"seed","value_type":"int","role":"sampling","sampling_kind":"seed"}]},
{"path":["quiet"],"key":"quiet","leaf":true,"generates":false,"args":[
{"id":"temp","long":"temp","value_type":"float","role":"sampling","sampling_kind":"temperature"}]},
{"path":["plain"],"key":"plain","leaf":true,"generates":true,"args":[
{"id":"model","positional":true,"required":true,"value_type":"path","role":"model"}]},
{"path":["oneside"],"key":"oneside","leaf":true,"generates":false,"args":[
{"id":"a","long":"a","value_type":"flag","values":["true","false"],"role":"mode","conflicts_with":["b","a"]},
{"id":"b","long":"b","value_type":"flag","values":["true","false"],"role":"mode"},
{"id":"c","long":"c","value_type":"flag","values":["true","false"],"role":"mode"},
{"id":"d","long":"d","value_type":"flag","values":["true","false"],"role":"mode","conflicts_with":["c"]}]}]}"#;
fn host_one() -> HostModels<'static> {
hosts().remove(0)
}
fn command<'s>(s: &'s Surface, key: &str) -> &'s Command {
s.commands
.iter()
.find(|c| c.key == key)
.expect("command present")
}
fn flag_levels() -> Vec<Factor> {
(0..4)
.map(|_| Factor {
kind: FactorKind::Shape,
levels: vec!["unset".into(), "true".into()],
})
.collect()
}
#[test]
fn forbidden_takes_a_conflict_declared_on_either_side_and_never_a_self_pair() {
let s = cli_surface::parse("t.json", KNOBS).expect("parses");
let c = command(&s, "oneside");
let args: Vec<&Arg> = c.args.iter().collect();
let fb = forbidden(&args, &flag_levels());
assert!(fb.contains(&(0, 1, 1, 1)), "a declares b, b does not");
assert!(fb.contains(&(2, 1, 3, 1)), "d declares c, c does not");
assert_eq!(fb.len(), 2, "no self pair, no unset pair: {fb:?}");
}
#[test]
fn effect_base_owes_thinking_off_and_the_smallest_rung_only_when_the_command_generates() {
let s = cli_surface::parse("t.json", KNOBS).expect("parses");
let cells = derive(&s, &hosts());
let base = |cmd: &str| {
cells
.iter()
.find(|c| c.command == cmd && c.kind == CellKind::Base)
.unwrap_or_else(|| panic!("{cmd} has a base"))
};
let g = base("samp");
assert_eq!(
g.thinking.as_deref(),
Some("off"),
"off, not the first mode"
);
assert_eq!(g.rung.as_deref(), Some("4k"));
assert_eq!(g.rung_tokens, Some(4096));
let mut modes_no_off = hosts();
modes_no_off[0].models[0].modes = vec!["on", "auto"];
let first = derive(&s, &modes_no_off);
let b = first
.iter()
.find(|c| c.command == "samp" && c.kind == CellKind::Base)
.expect("base");
assert_eq!(b.thinking.as_deref(), Some("on"), "no off: the first mode");
let p = base("plain");
assert_eq!(p.thinking.as_deref(), Some("off"));
}
#[test]
fn effect_base_of_a_non_generating_command_owes_no_thinking_or_rung() {
let s = cli_surface::parse("t.json", KNOBS).expect("parses");
let h = host_one();
let mut with_model = s.clone();
let q = with_model
.commands
.iter_mut()
.find(|c| c.key == "quiet")
.expect("quiet");
q.args.push(Arg {
id: "m".into(),
long: None,
short: None,
positional: true,
required: true,
value_type: "path".into(),
values: Vec::new(),
role: "model".into(),
marker: None,
hidden: false,
conflicts_with: Vec::new(),
sampling_kind: None,
});
let cells = effects(&with_model, command(&with_model, "quiet"), &h);
let base = cells
.iter()
.find(|c| c.kind == CellKind::Base)
.expect("base");
assert!(base.thinking.is_none() && base.rung.is_none() && base.rung_tokens.is_none());
}
#[test]
fn sampling_controls_pick_each_knob_by_role_and_kind_on_the_off_small_base() {
let s = cli_surface::parse("t.json", KNOBS).expect("parses");
let cells = sampling(command(&s, "samp"), &host_one());
assert_eq!(cells.len(), 5, "{cells:?}");
let by = |control: &str| {
cells
.iter()
.find(|c| matches!(&c.kind, CellKind::Sampling { control: k } if k == control))
.unwrap_or_else(|| panic!("control {control}"))
};
assert_eq!(
by("t0").args,
vec![("--temp".to_string(), "0".to_string())],
"the sampling-role temperature knob, not the decoy mode arg or another kind"
);
assert_eq!(
by("topk1").args,
vec![("--top-k".to_string(), "1".to_string())]
);
assert_eq!(
by("seed-b").args,
vec![
("--temp".to_string(), "0.8".to_string()),
("--seed".to_string(), "2".to_string())
]
);
let t0 = by("t0");
assert_eq!(
t0.thinking.as_deref(),
Some("off"),
"off, not the first mode"
);
assert_eq!(t0.rung.as_deref(), Some("4k"));
assert_eq!(t0.rung_tokens, Some(4096));
assert_eq!(t0.model_file.as_deref(), Some("a.gguf"));
assert!(t0.id.ends_with("#t0") && t0.id.starts_with("sampling:samp/"));
}
#[test]
fn sampling_is_owed_by_generating_commands_with_sampling_args_only() {
let s = cli_surface::parse("t.json", KNOBS).expect("parses");
let h = host_one();
assert!(
sampling(command(&s, "quiet"), &h).is_empty(),
"a non-generating command owes none"
);
assert!(
sampling(command(&s, "plain"), &h).is_empty(),
"a generating command with no sampling args owes none"
);
assert_eq!(
sampling(command(&s, "only"), &h).len(),
5,
"a command whose every arg is a sampling arg still owes its controls"
);
}
}