1use crate::hardware::HardwareInfo;
12use crate::models::ModelRegistry;
13use serde::{Deserialize, Serialize};
14use tracing::debug;
15
16pub use crate::adaptive_router::TaskComplexity;
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct RoutingDecision {
22 pub model: String,
23 pub complexity: TaskComplexity,
24 pub reason: String,
25}
26
27pub struct ModelRouter {
29 hw: HardwareInfo,
30}
31
32const TIERS: &[&str] = &[
34 "Qwen3-0.6B",
35 "Qwen3-1.7B",
36 "Qwen3-4B",
37 "Qwen3-8B",
38 "Qwen3-30B-A3B",
39];
40
41pub(crate) fn tier_fallback_order(target: &str) -> impl Iterator<Item = &'static str> {
45 let tier = target
46 .len()
47 .checked_sub("-MLX".len())
48 .filter(|&at| target.is_char_boundary(at) && target[at..].eq_ignore_ascii_case("-MLX"))
49 .map_or(target, |at| &target[..at]);
50 let target_idx = TIERS.iter().position(|&t| t == tier).unwrap_or(0);
51 TIERS[target_idx..]
52 .iter()
53 .chain(TIERS[..target_idx].iter().rev())
54 .copied()
55}
56
57impl ModelRouter {
58 pub fn new(hw: HardwareInfo) -> Self {
59 Self { hw }
60 }
61
62 pub fn route_generate(&self, prompt: &str, registry: &ModelRegistry) -> RoutingDecision {
64 let complexity = Self::assess_complexity(prompt);
65 let target = self.target_for_complexity(complexity);
66 let model = self.best_available(&target, registry, |_| true);
67 let reason = format!("{:?} task -> {} (target: {})", complexity, model, target);
68 debug!(
69 complexity = ?complexity,
70 target = %target,
71 selected = %model,
72 "routed generation request"
73 );
74 RoutingDecision {
75 model,
76 complexity,
77 reason,
78 }
79 }
80
81 pub fn route_small(&self, registry: &ModelRegistry) -> String {
83 self.route_small_within(registry, |_| true)
84 }
85
86 pub fn route_small_within(
89 &self,
90 registry: &ModelRegistry,
91 listed: impl Fn(&str) -> bool,
92 ) -> String {
93 let model = self.best_available("Qwen3-0.6B", registry, listed);
94 debug!(selected = %model, "routed small task (classify)");
95 model
96 }
97
98 pub fn route_embedding(&self, _registry: &ModelRegistry) -> String {
100 "Qwen3-Embedding-0.6B".to_string()
101 }
102
103 pub fn assess_complexity(prompt: &str) -> TaskComplexity {
105 let lower = prompt.to_lowercase();
106 let word_count = prompt.split_whitespace().count();
107 let estimated_tokens = (word_count as f64 * 1.3) as usize;
108
109 let code_markers = [
111 "```",
112 "fn ",
113 "def ",
114 "class ",
115 "import ",
116 "require(",
117 "async fn",
118 "pub fn",
119 "function ",
120 "const ",
121 "let ",
122 "var ",
123 "#include",
124 "package ",
125 "impl ",
126 ];
127 let has_code = code_markers.iter().any(|m| prompt.contains(m));
128
129 let repair_markers = [
131 "fix", "repair", "debug", "refactor", "broken", "failing", "error", "bug",
132 ];
133 let has_repair = repair_markers.iter().any(|m| lower.contains(m));
134
135 let reasoning_markers = [
137 "analyze",
138 "compare",
139 "explain why",
140 "step by step",
141 "think through",
142 "evaluate",
143 "trade-off",
144 "tradeoff",
145 "pros and cons",
146 "architecture",
147 "design",
148 "strategy",
149 "optimize",
150 "comprehensive",
151 ];
152 let has_reasoning = reasoning_markers.iter().any(|m| lower.contains(m));
153
154 let simple_patterns = [
156 "what is",
157 "who is",
158 "when did",
159 "where is",
160 "how many",
161 "yes or no",
162 "true or false",
163 "name the",
164 "list the",
165 "define ",
166 ];
167 let is_simple_question = simple_patterns.iter().any(|p| lower.contains(p));
168
169 if has_code || has_repair {
171 TaskComplexity::Code
172 } else if has_reasoning || estimated_tokens > 500 {
173 TaskComplexity::Complex
174 } else if is_simple_question || estimated_tokens < 30 {
175 TaskComplexity::Simple
176 } else {
177 TaskComplexity::Medium
178 }
179 }
180
181 fn target_for_complexity(&self, complexity: TaskComplexity) -> String {
183 match complexity {
184 TaskComplexity::Simple => "Qwen3-0.6B".into(),
185 TaskComplexity::Medium => "Qwen3-1.7B".into(),
186 TaskComplexity::Code => "Qwen3-4B".into(),
187 TaskComplexity::Complex => self.hw.recommended_model.clone(),
188 }
189 }
190
191 fn fits_now(&self, size_mb: u64) -> bool {
203 const WORKING_RESERVE_MB: u64 = 1536;
205
206 if size_mb > self.hw.max_model_mb {
207 return false;
208 }
209 match crate::hardware::available_ram_mb() {
210 Some(avail) => size_mb.saturating_add(WORKING_RESERVE_MB) <= avail,
211 None => true,
212 }
213 }
214
215 fn best_available(
224 &self,
225 target: &str,
226 registry: &ModelRegistry,
227 listed: impl Fn(&str) -> bool,
228 ) -> String {
229 let models: Vec<_> = registry
230 .list_models()
231 .into_iter()
232 .filter(|m| listed(&m.name))
233 .collect();
234
235 if let Some(m) = models.iter().find(|m| m.name == target) {
237 if m.downloaded && self.fits_now(m.quantized_size_mb) {
238 return target.to_string();
239 }
240 }
241
242 let listed_tier = |name: &str| models.iter().find(|m| m.name == name);
245 if let Some(name) = tier_fallback_order(target).find(|&name| {
246 listed_tier(name).is_some_and(|m| m.downloaded && self.fits_now(m.quantized_size_mb))
247 }) {
248 return name.to_string();
249 }
250
251 if listed(target) {
257 return target.to_string();
258 }
259 tier_fallback_order(target)
260 .find(|&name| listed_tier(name).is_some())
261 .map_or_else(|| target.to_string(), str::to_string)
262 }
263}
264
265#[cfg(test)]
266mod tests {
267 use super::*;
268
269 #[test]
270 fn simple_question() {
271 assert_eq!(
272 TaskComplexity::assess("What is the capital of France?"),
273 TaskComplexity::Simple,
274 );
275 }
276
277 #[test]
278 fn code_with_function_keyword() {
279 assert_eq!(
280 TaskComplexity::assess("Write a function to sort a list"),
281 TaskComplexity::Code,
282 );
283 }
284
285 #[test]
286 fn complex_reasoning() {
287 assert_eq!(
288 TaskComplexity::assess(
289 "Analyze the trade-offs between microservices and monolithic architecture for our use case"
290 ),
291 TaskComplexity::Complex,
292 );
293 }
294
295 #[test]
296 fn medium_default() {
297 assert_eq!(
299 TaskComplexity::assess(
300 "Tell me about the history of computing and how it has evolved over the decades from early mechanical calculation devices through transistors to the modern digital systems we use today in everyday life"
301 ),
302 TaskComplexity::Medium,
303 );
304 }
305
306 #[test]
307 fn code_with_backticks() {
308 assert_eq!(
309 TaskComplexity::assess("Here is my code:\n```rust\nfn main() {}\n```\nFix it"),
310 TaskComplexity::Code,
311 );
312 }
313
314 #[test]
315 fn repair_marker() {
316 assert_eq!(
317 TaskComplexity::assess("Debug this failing test case"),
318 TaskComplexity::Code,
319 );
320 }
321
322 fn test_router() -> ModelRouter {
323 ModelRouter::new(HardwareInfo {
324 os: "macos".into(),
325 arch: "aarch64".into(),
326 cpu_cores: 10,
327 total_ram_mb: 32768,
328 gpu_backend: crate::hardware::GpuBackend::Metal,
329 gpu_memory_mb: Some(28672),
330 gpu_devices: Vec::new(),
331 recommended_model: "Qwen3-8B".into(),
332 recommended_context: 8192,
333 max_model_mb: 18000,
334 })
335 }
336
337 #[test]
342 fn undownloaded_fallback_stays_within_listed() {
343 let tmp = tempfile::TempDir::new().unwrap();
344 let registry = ModelRegistry::new(tmp.path().to_path_buf());
345 let router = test_router();
346
347 assert_eq!(
348 router.route_small_within(®istry, |name| name != "Qwen3-0.6B"),
349 "Qwen3-1.7B",
350 "the nearest listed tier above an unlisted target"
351 );
352 assert_eq!(
353 router.best_available("Qwen3-8B", ®istry, |name| name == "Qwen3-4B"),
354 "Qwen3-4B",
355 "a listed tier below the target when none above it is listed"
356 );
357 assert_eq!(
358 router.route_small_within(®istry, |_| true),
359 "Qwen3-0.6B",
360 "a listed target is still the one returned"
361 );
362 assert_eq!(
363 router.best_available("Qwen3-8B-MLX", ®istry, |_| true),
364 "Qwen3-8B-MLX",
365 "an admitted target this catalog does not carry is still returned"
366 );
367 assert_eq!(
368 router.route_small_within(®istry, |_| false),
369 "Qwen3-0.6B",
370 "with no tier listed, the target comes back for the caller to refuse"
371 );
372 }
373}