1use crate::{
2 advice::{throughput, unified},
3 catalog::{BACKENDS, CAPABILITIES, Catalog, CatalogModel, PerfDataset, date, require},
4 sizing::{
5 Architecture, Hardware, KvCacheType, SizingRequest, ValidationError, evaluate,
6 parse_param_count, quant_bits,
7 },
8};
9use serde::{Deserialize, Serialize};
10use serde_json::{Value, json};
11use time::{OffsetDateTime, format_description::well_known::Rfc3339};
12
13#[derive(Debug, Default, Deserialize, Serialize)]
14#[serde(rename_all = "camelCase", deny_unknown_fields)]
15pub struct AdviceOptions {
16 pub task: Option<String>,
17 pub context: Option<f64>,
18 pub context_percent: Option<u8>,
19 #[serde(default)]
20 pub max_context: bool,
21 pub backend: Option<String>,
22 pub available_backends: Option<Vec<String>>,
23 #[serde(default, skip_serializing_if = "Option::is_none")]
24 pub kv_cache: Option<KvCacheType>,
25}
26impl AdviceOptions {
27 pub fn validate(&self) -> Result<(), ValidationError> {
28 require(
29 self.kv_cache.is_none()
30 || self.context.is_some()
31 || self.context_percent.is_some()
32 || self.max_context,
33 "a KV cache type requires context, contextPercent, or maxContext",
34 )?;
35 require(
36 usize::from(self.context.is_some())
37 + usize::from(self.context_percent.is_some())
38 + usize::from(self.max_context)
39 <= 1,
40 "context, contextPercent, and maxContext are mutually exclusive",
41 )?;
42 require(
43 self.task
44 .as_ref()
45 .is_none_or(|task| CAPABILITIES.contains(&task.as_str())),
46 "invalid task",
47 )?;
48 require(
49 self.backend
50 .as_ref()
51 .is_none_or(|backend| BACKENDS.contains(&backend.as_str())),
52 "invalid backend",
53 )?;
54 require(
55 self.context_percent
56 .is_none_or(|percent| [25, 50, 75, 100].contains(&percent)),
57 "invalid context percentage",
58 )?;
59 require(
60 self.context.is_none_or(|context| {
61 context.is_finite()
62 && context.fract() == 0.0
63 && (1.0..=10_000_000.0).contains(&context)
64 }),
65 "invalid context",
66 )?;
67 require(
68 self.available_backends.as_ref().is_none_or(|backends| {
69 backends
70 .iter()
71 .all(|backend| BACKENDS.contains(&backend.as_str()))
72 }),
73 "invalid available backend",
74 )
75 }
76 pub fn tokens(&self, model: &CatalogModel) -> Option<f64> {
77 self.context_percent
78 .map(|percent| {
79 (model.context_length * f64::from(percent) / 100.0)
80 .floor()
81 .max(1.0)
82 })
83 .or(self.context)
84 }
85}
86
87pub fn backends(model: &CatalogModel, hardware: &Hardware) -> Vec<&'static str> {
88 if model.supports_backend("bitnet") {
89 return vec!["bitnet"];
90 }
91 let mut result = Vec::new();
92 if model.supports_backend("ollama") {
93 result.push("ollama");
94 }
95 if model.supports_backend("llamacpp") {
96 result.push("llamacpp");
97 }
98 if unified(hardware) && model.supports_backend("mlx") {
99 result.push("mlx");
100 }
101 if model.supports_backend("lmstudio") && (model.source.gguf.is_some() || unified(hardware)) {
102 result.push("lmstudio");
103 }
104 result
105}
106
107pub fn recommend(
108 catalog: &Catalog,
109 hardware: &Hardware,
110 perf: &PerfDataset,
111 options: &AdviceOptions,
112) -> Result<Value, ValidationError> {
113 recommend_inner(catalog, hardware, perf, options, false)
114}
115pub fn recommend_detailed(
116 catalog: &Catalog,
117 hardware: &Hardware,
118 perf: &PerfDataset,
119 options: &AdviceOptions,
120) -> Result<Value, ValidationError> {
121 recommend_inner(catalog, hardware, perf, options, true)
122}
123fn recommend_inner(
124 catalog: &Catalog,
125 hardware: &Hardware,
126 perf: &PerfDataset,
127 options: &AdviceOptions,
128 detailed: bool,
129) -> Result<Value, ValidationError> {
130 options.validate()?;
131 let reference = OffsetDateTime::parse(&catalog.generated_at, &Rfc3339)
132 .map_err(|_| ValidationError("invalid reference date".into()))?;
133 let backend = options.backend.as_deref().unwrap_or("ollama");
134 let mut entries: Vec<(Value, f64, f64, String, String)> = Vec::new();
135 let mut wont_fit = Vec::new();
136 let mut usable = 0.0;
137 let mut kind = "ram";
138 for model in &catalog.models {
139 let resized = options
140 .kv_cache
141 .map(|kind| model.with_kv_cache(kind))
142 .transpose()?;
143 let model = resized.as_ref().unwrap_or(model);
144 let context = options.tokens(model);
145 let sized = evaluate(&SizingRequest {
146 model: model.sizing(),
147 hardware: hardware.clone(),
148 context,
149 })?;
150 usable = sized.usable_bytes;
151 kind = sized.memory_kind;
152 let supported = backends(model, hardware);
153 if options
154 .available_backends
155 .as_ref()
156 .is_some_and(|available| {
157 !supported
158 .iter()
159 .any(|backend| available.iter().any(|name| name == backend))
160 })
161 {
162 continue;
163 }
164 let fit = sized.fit;
165 if !fit.fits {
166 wont_fit.push(json!({"id":model.id,"reason":fit.reason}));
167 continue;
168 }
169 let quant = fit
170 .quant
171 .ok_or_else(|| ValidationError("missing fitting quant".into()))?;
172 let required = fit
173 .required_bytes
174 .ok_or_else(|| ValidationError("missing fitting size".into()))?;
175 let param_class = ((parse_param_count(&model.params)?.log10() - 9.0) / 3.0).clamp(0.0, 1.0);
176 let quality = model
177 .benchmark_proxy
178 .map_or(param_class, |proxy| 0.5 * param_class + 0.5 * proxy);
179 let fit_score = (1.0 - (required / usable - 0.6).abs() / 0.6).clamp(0.0, 1.0);
180 let active = if matches!(model.architecture, Architecture::Moe) {
181 model.active_params.as_deref().unwrap_or(&model.params)
182 } else {
183 &model.params
184 };
185 let bandwidth = if kind == "vram" {
186 1.0
187 } else if unified(hardware) {
188 0.5
189 } else {
190 0.2
191 };
192 let speed = (bandwidth
193 * (4.7 / quant_bits(&quant.name).unwrap_or(4.7)).clamp(0.0, 1.0)
194 * (7e9 / parse_param_count(active)?))
195 .clamp(0.0, 1.0);
196 let (day, _) = model
197 .recency()
198 .ok_or_else(|| ValidationError("catalog entry has no recency date".into()))?;
199 let released = date(day)?.midnight().assume_utc();
200 let age = (reference - released).as_seconds_f64() / 86400.0;
201 let recency = (1.0 - age / 730.0).clamp(0.0, 1.0);
202 let capability = if options
203 .task
204 .as_ref()
205 .is_none_or(|task| model.capabilities.contains(task))
206 {
207 1.0
208 } else {
209 0.0
210 };
211 let score =
212 (0.3 * quality + 0.2 * fit_score + 0.2 * speed + 0.15 * recency + 0.15 * capability)
213 .clamp(0.0, 1.0);
214 let estimate = throughput(model, &quant, hardware, perf, backend)?;
215 let verdict = if estimate.known
216 && (estimate.low_tok_per_sec + estimate.high_tok_per_sec) / 2.0 >= 10.0
217 {
218 "yes"
219 } else {
220 "slow"
221 };
222 let mut entry = json!({"rank":0,"id":model.id,"family":model.family,"params":model.params,"quant":quant.name,
223 "requiredBytes":required,"license":model.license,"capabilities":model.capabilities,"score":score,"verdict":verdict,
224 "estTokPerSec": if estimate.known { json!({"lowTokPerSec":estimate.low_tok_per_sec,"highTokPerSec":estimate.high_tok_per_sec}) } else { Value::Null },
225 "backends":supported,"throughputBackend":backend});
226 if model.provenance == crate::catalog::EntryProvenance::Auto {
227 entry["provenance"] = json!("auto");
228 }
229 if detailed {
230 entry["throughputEvidence"] = json!({"backend":backend,"source":"offline-estimate","unknownReason":if estimate.known{None}else{Some("no-sourced-performance-profile")}});
231 entry["scores"] = json!({"quality":quality,"fit":fit_score,"speed":speed,"recency":recency,"capability":capability});
232 entry["usableBytes"] = json!(usable);
233 entry["throughput"] = serde_json::to_value(&estimate)
234 .map_err(|_| ValidationError("invalid throughput".into()))?;
235 }
236 let index = model
237 .quantizations
238 .iter()
239 .position(|candidate| candidate.name == quant.name)
240 .ok_or_else(|| ValidationError("selected quant missing".into()))?;
241 if let Some(tokens) = context {
242 entry["context"] = json!(tokens);
243 entry["weightsBytes"] = json!(sized.weights[index]);
244 entry["kvCacheBytes"] = json!(model.kv_bytes_per_token.map(|rate| rate * tokens));
245 entry["kvPrecision"] = json!(KvCacheType::label(options.kv_cache));
246 } else if options.max_context {
247 let maximum = sized.max_context[index];
248 entry["maxContextTokens"] =
249 json!(maximum.map(|tokens| tokens.min(model.context_length)));
250 entry["boundBy"] =
251 json!(
252 maximum.map_or("unknown", |tokens| if tokens < model.context_length {
253 "hardware"
254 } else {
255 "model"
256 })
257 );
258 entry["kvPrecision"] = json!(KvCacheType::label(options.kv_cache));
259 }
260 entries.push((
261 entry,
262 score,
263 model.benchmark_proxy.unwrap_or(-1.0),
264 day.to_string(),
265 model.id.clone(),
266 ));
267 }
268 entries.sort_by(|left, right| {
269 right
270 .1
271 .total_cmp(&left.1)
272 .then_with(|| right.2.total_cmp(&left.2))
273 .then_with(|| right.3.cmp(&left.3))
274 .then_with(|| left.4.cmp(&right.4))
275 });
276 let ranked: Vec<Value> = entries
277 .into_iter()
278 .enumerate()
279 .map(|(index, (mut entry, ..))| {
280 entry["rank"] = json!(index + 1);
281 entry
282 })
283 .collect();
284 let command = ranked
285 .first()
286 .and_then(|entry| entry["id"].as_str())
287 .map(|id| format!("rigspark up {id}"));
288 Ok(
289 json!({"hardware":{"arch":hardware.arch,"platform":hardware.platform,"usableMemoryBytes":usable,"memoryKind":kind},"ranked":ranked,"wontFit":wont_fit,"command":command}),
290 )
291}