1use crate::records::{Capability, JsonValue, ModelRecord, ParamSpec, ParamType, RuntimeId};
5
6type Matcher = Box<dyn Fn(&ModelRecord) -> bool + Send + Sync>;
7
8pub struct ModelProfile {
11 pub id: String,
13 pub schema: Vec<ParamSpec>,
15 matcher: Matcher,
16}
17
18impl ModelProfile {
19 pub fn new(
21 id: &str,
22 schema: Vec<ParamSpec>,
23 matcher: impl Fn(&ModelRecord) -> bool + Send + Sync + 'static,
24 ) -> Self {
25 Self {
26 id: id.to_owned(),
27 schema,
28 matcher: Box::new(matcher),
29 }
30 }
31
32 pub fn matches(&self, record: &ModelRecord) -> bool {
34 (self.matcher)(record)
35 }
36}
37
38pub struct ProfileRegistry {
40 pub profiles: Vec<ModelProfile>,
42}
43
44impl ProfileRegistry {
45 pub fn new(profiles: Vec<ModelProfile>) -> Self {
47 Self { profiles }
48 }
49
50 pub fn schema(&self, record: &ModelRecord) -> Vec<ParamSpec> {
54 let mut specs: Vec<ParamSpec> = Vec::new();
55 for profile in &self.profiles {
56 if profile.matches(record) {
57 for spec in &profile.schema {
58 if !specs.iter().any(|kept| kept.key == spec.key) {
59 specs.push(spec.clone());
60 }
61 }
62 }
63 }
64 if !specs.iter().any(|spec| spec.key == "context_length")
65 && let Some(context) = context_length_spec(record)
66 {
67 specs.push(context);
68 }
69 specs
70 }
71
72 pub fn refreshed(&self, record: &ModelRecord) -> ModelRecord {
75 let schema = self.schema(record);
76 if schema.is_empty() {
77 return record.clone();
78 }
79 let mut updated = record.clone();
80 updated.params = schema;
81 updated
82 }
83
84 pub fn builtin() -> Self {
87 Self::new(vec![
88 ModelProfile::new(
89 "text-generation",
90 vec![temperature(), top_p(), max_tokens()],
91 is_text,
92 ),
93 runtime_extras(
94 "sampling-llama-cpp",
95 RuntimeId::llama_cpp(),
96 vec![
97 top_k(),
98 min_p(),
99 repeat_penalty(),
100 frequency_penalty(),
101 presence_penalty(),
102 seed(),
103 stop(),
104 ],
105 ),
106 runtime_extras(
107 "sampling-mlx-swift",
108 RuntimeId::mlx_swift(),
109 vec![repeat_penalty(), stop()],
110 ),
111 runtime_extras(
112 "sampling-mlx-lm",
113 RuntimeId::mlx_lm(),
114 vec![top_k(), min_p(), repeat_penalty(), seed(), stop()],
115 ),
116 runtime_extras(
117 "sampling-ollama",
118 RuntimeId::ollama(),
119 vec![
120 top_k(),
121 min_p(),
122 seed(),
123 repeat_penalty(),
124 frequency_penalty(),
125 presence_penalty(),
126 stop(),
127 ],
128 ),
129 runtime_extras(
130 "sampling-endpoint",
131 RuntimeId::openai_endpoint(),
132 vec![stop(), seed(), frequency_penalty(), presence_penalty()],
133 ),
134 runtime_extras(
135 "sampling-apple",
136 RuntimeId::apple_foundation(),
137 vec![top_k(), seed()],
138 ),
139 ModelProfile::new(
140 "speech-synthesis",
141 vec![plain("voice", ParamType::String), speed()],
142 |record| record.can(&Capability::speak()),
143 ),
144 ModelProfile::new(
145 "transcription",
146 vec![
147 plain("language", ParamType::String),
148 plain("translate", ParamType::Bool),
149 ],
150 |record| record.can(&Capability::transcribe()),
151 ),
152 ModelProfile::new(
153 "togglable-thinking",
154 vec![plain("thinking", ParamType::Bool)],
155 |record| {
156 record.can(&Capability::chat())
157 && record.runtime.id.as_ref().is_some_and(is_thinking_runtime)
158 },
159 ),
160 ])
161 }
162}
163
164pub fn context_length_spec(record: &ModelRecord) -> Option<ParamSpec> {
167 if !(record.can(&Capability::chat()) || record.can(&Capability::complete())) {
168 return None;
169 }
170 let runtime = record.runtime.id.as_ref()?;
171 if !is_context_honoring(runtime) {
172 return None;
173 }
174 match record.context_length {
175 Some(window) if window > 0 => Some(ParamSpec {
176 key: "context_length".to_owned(),
177 param_type: ParamType::Int,
178 default_value: Some(JsonValue::Int(window.min(32768))),
179 range: Some(vec![
180 JsonValue::Int(512.min(window)),
181 JsonValue::Int(window),
182 ]),
183 values: None,
184 }),
185 _ => Some(ParamSpec {
186 key: "context_length".to_owned(),
187 param_type: ParamType::Int,
188 default_value: None,
189 range: Some(vec![JsonValue::Int(512), JsonValue::Int(131072)]),
190 values: None,
191 }),
192 }
193}
194
195fn is_text(record: &ModelRecord) -> bool {
196 record.can(&Capability::chat()) || record.can(&Capability::complete())
197}
198
199fn is_context_honoring(runtime: &RuntimeId) -> bool {
200 *runtime == RuntimeId::ollama() || *runtime == RuntimeId::llama_cpp()
201}
202
203fn is_thinking_runtime(runtime: &RuntimeId) -> bool {
204 *runtime == RuntimeId::ollama() || *runtime == RuntimeId::mlx_lm()
205}
206
207fn runtime_extras(id: &str, runtime: RuntimeId, schema: Vec<ParamSpec>) -> ModelProfile {
208 ModelProfile::new(id, schema, move |record| {
209 is_text(record) && record.runtime.id.as_ref() == Some(&runtime)
210 })
211}
212
213fn float(key: &str, low: f64, high: f64) -> ParamSpec {
214 ParamSpec {
215 key: key.to_owned(),
216 param_type: ParamType::Float,
217 default_value: None,
218 range: Some(vec![JsonValue::Double(low), JsonValue::Double(high)]),
219 values: None,
220 }
221}
222
223fn int(key: &str, low: i64, high: i64) -> ParamSpec {
224 ParamSpec {
225 key: key.to_owned(),
226 param_type: ParamType::Int,
227 default_value: None,
228 range: Some(vec![JsonValue::Int(low), JsonValue::Int(high)]),
229 values: None,
230 }
231}
232
233fn plain(key: &str, param_type: ParamType) -> ParamSpec {
234 ParamSpec {
235 key: key.to_owned(),
236 param_type,
237 default_value: None,
238 range: None,
239 values: None,
240 }
241}
242
243fn temperature() -> ParamSpec {
244 float("temperature", 0.0, 2.0)
245}
246fn top_p() -> ParamSpec {
247 float("top_p", 0.0, 1.0)
248}
249fn top_k() -> ParamSpec {
250 int("top_k", 0, 100)
251}
252fn min_p() -> ParamSpec {
253 float("min_p", 0.0, 1.0)
254}
255fn max_tokens() -> ParamSpec {
256 int("max_tokens", 1, 32768)
257}
258fn repeat_penalty() -> ParamSpec {
259 float("repeat_penalty", 0.5, 2.0)
260}
261fn frequency_penalty() -> ParamSpec {
262 float("frequency_penalty", -2.0, 2.0)
263}
264fn presence_penalty() -> ParamSpec {
265 float("presence_penalty", -2.0, 2.0)
266}
267fn seed() -> ParamSpec {
268 plain("seed", ParamType::Int)
269}
270fn stop() -> ParamSpec {
271 plain("stop", ParamType::String)
272}
273fn speed() -> ParamSpec {
274 ParamSpec {
275 key: "speed".to_owned(),
276 param_type: ParamType::Float,
277 default_value: Some(JsonValue::Double(1.0)),
278 range: Some(vec![JsonValue::Double(0.5), JsonValue::Double(2.0)]),
279 values: None,
280 }
281}