Skip to main content

mesh_client/network/
router.rs

1/// Smart model router — classifies requests and picks the best model.
2use serde_json::Value;
3
4// ── Request categories ──────────────────────────────────────────────
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
7pub enum Category {
8    Code,
9    Reasoning,
10    Chat,
11    ToolCall,
12    Creative,
13    /// Factual lookup, summarization, knowledge retrieval
14    Info,
15    /// Image generation or analysis (future: multimodal models)
16    Image,
17}
18
19/// How complex/heavy the request appears to be.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Complexity {
22    Quick,    // simple fact, short answer, casual
23    Moderate, // normal conversation, standard code
24    Deep,     // long reasoning, complex analysis, architecture
25}
26
27/// Full classification result.
28#[derive(Debug, Clone, PartialEq)]
29pub struct Classification {
30    pub category: Category,
31    pub complexity: Complexity,
32    pub needs_tools: bool,
33    pub has_media_inputs: bool,
34}
35
36#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
37pub struct MediaRequirements {
38    pub has_media: bool,
39    pub needs_vision: bool,
40    pub needs_audio: bool,
41}
42
43// ── Model profiles ──────────────────────────────────────────────────
44
45/// Quality tier: higher = better quality, slower.
46/// 1 = draft/tiny, 2 = good, 3 = strong, 4 = frontier
47pub type Tier = u8;
48
49pub struct ModelProfile {
50    pub name: &'static str,
51    pub strengths: &'static [Category],
52    pub tier: Tier,
53    /// Whether this model can handle tool-calling requests (function calling).
54    /// Models without this set to true are filtered out when tools are present.
55    pub tools: bool,
56}
57
58/// Static profiles for catalog models.
59/// Order of strengths matters — first entry is primary strength.
60pub static MODEL_PROFILES: &[ModelProfile] = &[
61    // ── Tier 4: Frontier ────────────────────────────────────────
62    ModelProfile {
63        name: "Qwen3-235B-A22B-Q4_K_M",
64        strengths: &[
65            Category::Code,
66            Category::Reasoning,
67            Category::Chat,
68            Category::Creative,
69        ],
70        tier: 4,
71        tools: true,
72    },
73    ModelProfile {
74        name: "Llama-3.1-405B-Instruct-Q2_K",
75        strengths: &[Category::Chat, Category::Reasoning, Category::Code],
76        tier: 4,
77        tools: true,
78    },
79    ModelProfile {
80        name: "MiniMax-M2.5-Q4_K_M",
81        strengths: &[
82            Category::Code,
83            Category::Reasoning,
84            Category::Chat,
85            Category::Creative,
86            Category::ToolCall,
87        ],
88        tier: 4,
89        tools: true,
90    },
91    // ── Tier 3: Strong ──────────────────────────────────────────
92    ModelProfile {
93        name: "Qwen2.5-72B-Instruct-Q4_K_M",
94        strengths: &[Category::Chat, Category::Reasoning, Category::Code],
95        tier: 3,
96        tools: true,
97    },
98    ModelProfile {
99        name: "Llama-3.3-70B-Instruct-Q4_K_M",
100        strengths: &[Category::Chat, Category::ToolCall, Category::Code],
101        tier: 3,
102        tools: true,
103    },
104    ModelProfile {
105        name: "DeepSeek-R1-Distill-70B-Q4_K_M",
106        strengths: &[Category::Reasoning],
107        tier: 3,
108        tools: false, // reasoning-only, no tool support
109    },
110    ModelProfile {
111        name: "Mixtral-8x22B-Instruct-Q4_K_M",
112        strengths: &[Category::Chat, Category::Code, Category::Reasoning],
113        tier: 3,
114        tools: true,
115    },
116    ModelProfile {
117        name: "Qwen3-32B-Q4_K_M",
118        strengths: &[Category::Reasoning, Category::Code, Category::Chat],
119        tier: 3,
120        tools: true,
121    },
122    ModelProfile {
123        name: "Qwen2.5-Coder-32B-Instruct-Q4_K_M",
124        strengths: &[Category::Code],
125        tier: 3,
126        tools: true,
127    },
128    ModelProfile {
129        name: "DeepSeek-R1-Distill-Qwen-32B-Q4_K_M",
130        strengths: &[Category::Reasoning],
131        tier: 3,
132        tools: false,
133    },
134    ModelProfile {
135        name: "Qwen3-30B-A3B-Q4_K_M",
136        strengths: &[Category::Chat, Category::Reasoning, Category::Code],
137        tier: 3,
138        tools: true,
139    },
140    ModelProfile {
141        name: "Qwen3-Coder-30B-A3B-Instruct-Q4_K_M",
142        strengths: &[Category::Code, Category::ToolCall],
143        tier: 3,
144        tools: true,
145    },
146    ModelProfile {
147        name: "Qwen2.5-32B-Instruct-Q4_K_M",
148        strengths: &[
149            Category::Chat,
150            Category::Reasoning,
151            Category::Code,
152            Category::ToolCall,
153        ],
154        tier: 3,
155        tools: true,
156    },
157    ModelProfile {
158        name: "Gemma-3-27B-it-Q4_K_M",
159        strengths: &[Category::Reasoning, Category::Chat],
160        tier: 3,
161        tools: false, // unreliable tool calling
162    },
163    ModelProfile {
164        name: "Qwen3.5-27B-Q4_K_M",
165        strengths: &[Category::Code, Category::Reasoning, Category::Chat],
166        tier: 3,
167        tools: true,
168    },
169    ModelProfile {
170        name: "Qwen3-Coder-Next-Q4_K_M",
171        strengths: &[Category::Code, Category::ToolCall, Category::Reasoning],
172        tier: 4,
173        tools: true,
174    },
175    // ── Tier 2: Good ────────────────────────────────────────────
176    ModelProfile {
177        name: "Qwen3.5-9B-Q4_K_M",
178        strengths: &[Category::Chat, Category::Code],
179        tier: 2,
180        tools: false,
181    },
182    ModelProfile {
183        name: "Mistral-Small-3.1-24B-Instruct-Q4_K_M",
184        strengths: &[Category::Chat, Category::ToolCall],
185        tier: 2,
186        tools: true,
187    },
188    ModelProfile {
189        name: "Devstral-Small-2505-Q4_K_M",
190        strengths: &[Category::Code, Category::ToolCall],
191        tier: 2,
192        tools: true,
193    },
194    ModelProfile {
195        name: "GLM-4.7-Flash-Q4_K_M",
196        strengths: &[Category::Chat, Category::ToolCall],
197        tier: 2,
198        tools: true,
199    },
200    ModelProfile {
201        name: "GLM-4-32B-0414-Q4_K_M",
202        strengths: &[Category::Chat, Category::ToolCall, Category::Code],
203        tier: 2,
204        tools: true,
205    },
206    ModelProfile {
207        name: "Llama-4-Scout-Q4_K_M",
208        strengths: &[Category::Chat, Category::ToolCall],
209        tier: 2,
210        tools: true,
211    },
212    ModelProfile {
213        name: "Qwen3-14B-Q4_K_M",
214        strengths: &[Category::Chat, Category::Reasoning],
215        tier: 2,
216        tools: true,
217    },
218    ModelProfile {
219        name: "Qwen2.5-14B-Instruct-Q4_K_M",
220        strengths: &[Category::Chat],
221        tier: 2,
222        tools: true,
223    },
224    ModelProfile {
225        name: "Qwen2.5-Coder-14B-Instruct-Q4_K_M",
226        strengths: &[Category::Code],
227        tier: 2,
228        tools: true,
229    },
230    ModelProfile {
231        name: "DeepSeek-R1-Distill-Qwen-14B-Q4_K_M",
232        strengths: &[Category::Reasoning],
233        tier: 2,
234        tools: false,
235    },
236    ModelProfile {
237        name: "Gemma-3-12B-it-Q4_K_M",
238        strengths: &[Category::Chat, Category::Reasoning],
239        tier: 2,
240        tools: false,
241    },
242    ModelProfile {
243        name: "Qwen3-8B-Q4_K_M",
244        strengths: &[Category::Chat, Category::Code],
245        tier: 2,
246        tools: true,
247    },
248    ModelProfile {
249        name: "Hermes-2-Pro-Mistral-7B-Q4_K_M",
250        strengths: &[Category::Chat],
251        tier: 2,
252        tools: false,
253    },
254    ModelProfile {
255        name: "Qwen2.5-Coder-7B-Instruct-Q4_K_M",
256        strengths: &[Category::Code],
257        tier: 2,
258        tools: true,
259    },
260    // ── Tier 1: Small / Draft ───────────────────────────────────
261    ModelProfile {
262        name: "Qwen3-4B-Q4_K_M",
263        strengths: &[Category::Chat],
264        tier: 1,
265        tools: true,
266    },
267    ModelProfile {
268        name: "Qwen2.5-3B-Instruct-Q4_K_M",
269        strengths: &[Category::Chat],
270        tier: 1,
271        tools: true,
272    },
273    ModelProfile {
274        name: "Llama-3.2-3B-Instruct-Q4_K_M",
275        strengths: &[Category::Chat, Category::ToolCall],
276        tier: 1,
277        tools: true,
278    },
279];
280
281pub fn profile_for(model_name: &str) -> Option<&'static ModelProfile> {
282    // Direct match first
283    if let Some(p) = MODEL_PROFILES.iter().find(|p| p.name == model_name) {
284        return Some(p);
285    }
286    // Strip split GGUF suffix: "Model-00001-of-00004" → "Model"
287    let clean = strip_split_suffix(model_name);
288    if clean != model_name {
289        return MODEL_PROFILES.iter().find(|p| p.name == clean);
290    }
291    None
292}
293
294/// Strip split GGUF suffix like "-00001-of-00004" from a model name.
295pub fn strip_split_suffix(name: &str) -> &str {
296    // Pattern: -NNNNN-of-NNNNN at the end
297    if let Some(idx) = name.rfind("-of-") {
298        // Check that what follows is digits and what precedes is -digits
299        let after = &name[idx + 4..];
300        if after.chars().all(|c| c.is_ascii_digit()) && !after.is_empty() {
301            // Find the preceding -NNNNN
302            if let Some(dash) = name[..idx].rfind('-') {
303                let between = &name[dash + 1..idx];
304                if between.chars().all(|c| c.is_ascii_digit()) && !between.is_empty() {
305                    return &name[..dash];
306                }
307            }
308        }
309    }
310    name
311}
312
313/// Owned version of strip_split_suffix for contexts that need a String.
314pub fn strip_split_suffix_owned(name: &str) -> String {
315    strip_split_suffix(name).to_string()
316}
317
318// ── Request classification ──────────────────────────────────────────
319
320/// Classify a chat completion request body using heuristics.
321/// No LLM call, just pattern matching on the request structure.
322/// Classify a request body into category + complexity + needs_tools.
323/// Tools presence is an attribute, not a category override — a code request
324/// with tools is still Code (with needs_tools=true), not ToolCall.
325pub fn classify(body: &Value) -> Classification {
326    // Collect all text from messages for keyword analysis
327    let text = collect_message_text(body);
328    let lower = text.to_lowercase();
329    let media = media_requirements(body);
330
331    // Check if the request actually needs tool execution.
332    // If the client sends a tools schema, this is an agentic session (Claude Code,
333    // Goose, etc.) — always prefer the strongest tool-capable model regardless of
334    // what the first message says.  Keyword matching on content is a secondary signal
335    // but not required when tools are present.
336    let has_tools_schema = body
337        .get("tools")
338        .and_then(|t| t.as_array())
339        .map(|a| !a.is_empty())
340        .unwrap_or(false);
341    // Anthropic-style requests may include structured content blocks with
342    // explicit tool_use/tool_result blocks — definitely tool-driven.
343    let has_tool_blocks = body
344        .get("messages")
345        .and_then(|m| m.as_array())
346        .map(|msgs| {
347            msgs.iter().any(|msg| {
348                msg.get("content")
349                    .and_then(|c| c.as_array())
350                    .map(|blocks| {
351                        blocks.iter().any(|b| {
352                            matches!(
353                                b.get("type").and_then(|t| t.as_str()),
354                                Some("tool_use") | Some("tool_result")
355                            )
356                        })
357                    })
358                    .unwrap_or(false)
359            })
360        })
361        .unwrap_or(false);
362    let needs_tools = has_tools_schema || has_tool_blocks;
363
364    // Count last user message tokens (rough proxy for complexity)
365    let last_user_len = last_user_message_len(body);
366
367    // Code signals
368    let code_signals = [
369        "```",
370        "def ",
371        "fn ",
372        "func ",
373        "class ",
374        "import ",
375        "function",
376        "const ",
377        "let ",
378        "var ",
379        "return ",
380        "write a program",
381        "write code",
382        "implement",
383        "refactor",
384        "debug",
385        "fix the bug",
386        "write a script",
387        "code review",
388        "pull request",
389        "git ",
390        "compile",
391        "syntax",
392        "python",
393        "javascript",
394        "typescript",
395        " rust ",
396        "golang",
397        "java ",
398        "c++",
399        " ruby ",
400        " swift ",
401        "kotlin",
402        "algorithm",
403        "binary search",
404        " sort ",
405        "regex",
406        " api ",
407        " http ",
408        " sql ",
409        "database",
410        " query ",
411    ];
412    let code_score: usize = code_signals.iter().filter(|s| lower.contains(*s)).count();
413
414    // Reasoning signals
415    let reasoning_signals = [
416        "prove",
417        "explain why",
418        "step by step",
419        "calculate",
420        "solve",
421        "derive",
422        "what is the probability",
423        "how many",
424        "analyze",
425        "compare and contrast",
426        "evaluate",
427        "mathematical",
428        "theorem",
429        "equation",
430        "logic",
431        "think carefully",
432        "reason about",
433    ];
434    let reasoning_score: usize = reasoning_signals
435        .iter()
436        .filter(|s| lower.contains(*s))
437        .count();
438
439    // Creative signals
440    let creative_signals = [
441        "write a story",
442        "write a poem",
443        "creative",
444        "imagine",
445        "fiction",
446        "narrative",
447        "compose",
448        "brainstorm",
449        "write a song",
450        "screenplay",
451        "dialogue",
452    ];
453    let creative_score: usize = creative_signals
454        .iter()
455        .filter(|s| lower.contains(*s))
456        .count();
457
458    // Info/knowledge signals — factual lookup, summarization
459    let info_signals = [
460        "what is",
461        "who is",
462        "when did",
463        "where is",
464        "how does",
465        "define ",
466        "explain ",
467        "summarize",
468        "summary",
469        "overview",
470        "tell me about",
471        "describe ",
472        "what are the",
473        "list the",
474        "difference between",
475        "compare ",
476        "history of",
477    ];
478    let info_score: usize = info_signals.iter().filter(|s| lower.contains(*s)).count();
479
480    // Image signals — generation or analysis (future)
481    let image_signals = [
482        "image",
483        "picture",
484        "photo",
485        "draw",
486        "generate an image",
487        "visualize",
488        "diagram",
489        "screenshot",
490        "describe this image",
491    ];
492    let image_score: usize = image_signals.iter().filter(|s| lower.contains(*s)).count();
493
494    // Deep-thinking signals (want the biggest brain)
495    let deep_signals = [
496        "architect",
497        "design a system",
498        "trade-off",
499        "tradeoff",
500        "in depth",
501        "comprehensive",
502        "thorough",
503        "detailed analysis",
504        "long-term",
505        "strategy",
506        "plan for",
507        "review this codebase",
508        "rewrite",
509        "from scratch",
510    ];
511    let deep_score: usize = deep_signals.iter().filter(|s| lower.contains(*s)).count();
512
513    // System prompt hints
514    let mut system_code = false;
515    if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
516        for msg in messages {
517            if msg.get("role").and_then(|r| r.as_str()) == Some("system")
518                && let Some(content) = msg.get("content").and_then(|c| c.as_str())
519            {
520                let sys = content.to_lowercase();
521                if sys.contains("developer") || sys.contains("coding") || sys.contains("programmer")
522                {
523                    system_code = true;
524                }
525            }
526        }
527    }
528
529    // Pick category — tools don't override, content wins
530    let category = if system_code
531        || code_score >= 2
532        || (code_score >= 1 && reasoning_score == 0 && creative_score == 0)
533    {
534        Category::Code
535    } else if reasoning_score >= 2 {
536        Category::Reasoning
537    } else if creative_score >= 1 {
538        Category::Creative
539    } else if media.needs_vision || image_score >= 1 {
540        Category::Image
541    } else if needs_tools && code_score == 0 && reasoning_score == 0 && creative_score == 0 {
542        // Only ToolCall if tools present AND no other signal dominates
543        Category::ToolCall
544    } else if info_score >= 2 && code_score == 0 {
545        Category::Info
546    } else {
547        Category::Chat
548    };
549
550    // Complexity: Quick / Moderate / Deep
551    let total_messages = body
552        .get("messages")
553        .and_then(|m| m.as_array())
554        .map(|a| a.len())
555        .unwrap_or(0);
556    let complexity = if deep_score >= 1 || last_user_len > 500 || total_messages > 10 {
557        Complexity::Deep
558    } else if last_user_len < 60 && total_messages <= 2 && reasoning_score == 0 && deep_score == 0 {
559        Complexity::Quick
560    } else {
561        Complexity::Moderate
562    };
563
564    Classification {
565        category,
566        complexity,
567        needs_tools,
568        has_media_inputs: media.has_media,
569    }
570}
571
572pub fn media_requirements(body: &Value) -> MediaRequirements {
573    let mut requirements = MediaRequirements::default();
574    let Some(messages) = body.get("messages").and_then(|m| m.as_array()) else {
575        return requirements;
576    };
577
578    for msg in messages {
579        let Some(blocks) = msg.get("content").and_then(|c| c.as_array()) else {
580            continue;
581        };
582        for block in blocks {
583            let block_type = block
584                .get("type")
585                .and_then(|t| t.as_str())
586                .unwrap_or_default();
587            match block_type {
588                "image_url" | "input_image" | "image" => {
589                    requirements.has_media = true;
590                    requirements.needs_vision = true;
591                }
592                "audio_url" | "input_audio" | "audio" => {
593                    requirements.has_media = true;
594                    requirements.needs_audio = true;
595                }
596                "file" | "input_file" => {
597                    requirements.has_media = true;
598                }
599                _ => {
600                    if block.get("image_url").is_some() || block.get("image").is_some() {
601                        requirements.has_media = true;
602                        requirements.needs_vision = true;
603                    }
604                    if block.get("audio_url").is_some() || block.get("audio").is_some() {
605                        requirements.has_media = true;
606                        requirements.needs_audio = true;
607                    }
608                }
609            }
610        }
611    }
612
613    requirements
614}
615
616/// Length of last user message in characters (rough complexity proxy).
617fn last_user_message_len(body: &Value) -> usize {
618    body.get("messages")
619        .and_then(|m| m.as_array())
620        .and_then(|msgs| {
621            msgs.iter()
622                .rev()
623                .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
624        })
625        .map(message_text)
626        .map(|s| s.len())
627        .unwrap_or(0)
628}
629
630fn collect_message_text(body: &Value) -> String {
631    let mut text = String::new();
632    if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
633        for msg in messages {
634            let content = message_text(msg);
635            if !content.is_empty() {
636                text.push_str(&content);
637                text.push('\n');
638            }
639        }
640    }
641    text
642}
643
644/// Extract message text for both OpenAI-style and Anthropic-style payloads.
645fn message_text(msg: &Value) -> String {
646    if let Some(s) = msg.get("content").and_then(|c| c.as_str()) {
647        return s.to_string();
648    }
649
650    // Anthropic content blocks: [{"type":"text","text":"..."}, ...]
651    if let Some(blocks) = msg.get("content").and_then(|c| c.as_array()) {
652        let mut out = String::new();
653        for b in blocks {
654            if let Some(t) = b.get("text").and_then(|t| t.as_str()) {
655                out.push_str(t);
656                out.push('\n');
657            }
658        }
659        return out;
660    }
661
662    String::new()
663}
664
665// ── Model selection ─────────────────────────────────────────────────
666
667/// Pick the best model using full classification (category + complexity + tools).
668pub fn pick_model_classified<'a>(
669    classification: &Classification,
670    available_models: &[(&'a str, f64)],
671) -> Option<&'a str> {
672    if available_models.is_empty() {
673        return None;
674    }
675
676    // Filter for tool-capable models if tools are required
677    let filtered: Vec<(&str, f64)> = if classification.needs_tools {
678        available_models
679            .iter()
680            .filter(|(name, _)| profile_for(name).map(|p| p.tools).unwrap_or(false))
681            .copied()
682            .collect()
683    } else {
684        available_models.to_vec()
685    };
686    // Fall back to all models if no tool-capable model found
687    let candidates = if filtered.is_empty() {
688        available_models
689    } else {
690        &filtered
691    };
692
693    let category = classification.category;
694
695    // Score each available model
696    let mut scored: Vec<(&str, i32)> = candidates
697        .iter()
698        .map(|(name, tok_s)| {
699            let profile = profile_for(name);
700            let tier = profile.map(|p| p.tier).unwrap_or(1) as i32;
701
702            // Task match is the primary signal.
703            let has_match = profile
704                .map(|p| p.strengths.contains(&category))
705                .unwrap_or(false);
706
707            let match_bonus = if has_match { 1000 } else { 0 };
708
709            // Within matched models: primary > secondary > listed
710            let position_bonus = profile
711                .map(|p| {
712                    p.strengths
713                        .iter()
714                        .position(|s| *s == category)
715                        .map(|i| match i {
716                            0 => 20,
717                            1 => 10,
718                            _ => 5,
719                        })
720                        .unwrap_or(0)
721                })
722                .unwrap_or(0);
723
724            // Agentic vs chat scoring:
725            // When tools are needed, strongly prefer the most capable model.
726            // For chat, prefer the fastest model that matches.
727            let tier_bonus = if classification.needs_tools {
728                // Agentic: always prefer strongest. Tier dominates.
729                // tier 1→20, tier 2→40, tier 3→60, tier 4→80
730                tier * 20
731            } else {
732                // Chat/no-tools: always prefer bigger models, but less aggressively
733                // than agentic. Small models are fallbacks, not first choice.
734                match classification.complexity {
735                    Complexity::Quick => tier * 5,     // tier 2→10, tier 3→15, tier 4→20
736                    Complexity::Moderate => tier * 10, // tier 2→20, tier 3→30, tier 4→40
737                    Complexity::Deep => tier * 15,     // tier 2→30, tier 3→45, tier 4→60
738                }
739            };
740
741            // Speed bonus: higher for chat (speed matters), lower for agentic (quality matters)
742            let speed_bonus = if classification.needs_tools {
743                // Agentic: speed is a tiebreaker only
744                (tok_s / 20.0).min(5.0) as i32
745            } else {
746                // Chat: speed matters more
747                (tok_s / 5.0).min(20.0) as i32
748            };
749
750            let score = match_bonus + tier_bonus + position_bonus + speed_bonus;
751            (*name, score)
752        })
753        .collect();
754
755    scored.sort_by_key(|entry| std::cmp::Reverse(entry.1));
756
757    // For non-agentic requests, spread load across top-scoring models.
758    // Pick randomly among candidates within 15 points of the best score.
759    // This avoids queueing all concurrent chat users on the same model
760    // while keeping weak models as fallbacks, not equal contenders.
761    if !classification.needs_tools && scored.len() > 1 {
762        let best_score = scored[0].1;
763        let top_tier: Vec<&(&str, i32)> = scored
764            .iter()
765            .filter(|(_, s)| best_score - s <= 15)
766            .collect();
767        if top_tier.len() > 1 {
768            // Simple pseudo-random: use current time nanos to pick
769            let nanos = std::time::SystemTime::now()
770                .duration_since(std::time::UNIX_EPOCH)
771                .unwrap_or_default()
772                .subsec_nanos() as usize;
773            let pick = top_tier[nanos % top_tier.len()];
774            return Some(pick.0);
775        }
776    }
777
778    scored.first().map(|(name, _)| *name)
779}
780
781/// Legacy wrapper for tests that have category + tools but no complexity.
782#[cfg(test)]
783pub fn pick_model_with_tools<'a>(
784    category: Category,
785    available_models: &[(&'a str, f64)],
786    tools_required: bool,
787) -> Option<&'a str> {
788    pick_model_classified(
789        &Classification {
790            category,
791            complexity: Complexity::Moderate,
792            needs_tools: tools_required,
793            has_media_inputs: false,
794        },
795        available_models,
796    )
797}
798
799// ── Tests ───────────────────────────────────────────────────────────
800
801#[cfg(test)]
802mod tests {
803    use super::*;
804    use serde_json::json;
805
806    #[test]
807    fn test_classify_tool_call() {
808        // Content that implies tool use + tools schema = ToolCall
809        let body = json!({
810            "messages": [{"role": "user", "content": "Run the tests and check the output"}],
811            "tools": [{"type": "function", "function": {"name": "bash"}}]
812        });
813        assert_eq!(classify(&body).category, Category::ToolCall);
814    }
815
816    #[test]
817    fn test_classify_code() {
818        let body = json!({
819            "messages": [
820                {"role": "user", "content": "Write a Python function to implement binary search and debug any issues"}
821            ]
822        });
823        assert_eq!(classify(&body).category, Category::Code);
824    }
825
826    #[test]
827    fn test_classify_reasoning() {
828        let body = json!({
829            "messages": [
830                {"role": "user", "content": "Prove that the square root of 2 is irrational. Explain step by step."}
831            ]
832        });
833        assert_eq!(classify(&body).category, Category::Reasoning);
834    }
835
836    #[test]
837    fn test_classify_creative() {
838        let body = json!({
839            "messages": [
840                {"role": "user", "content": "Write a story about a robot who learns to paint"}
841            ]
842        });
843        assert_eq!(classify(&body).category, Category::Creative);
844    }
845
846    #[test]
847    fn test_classify_chat_default() {
848        let body = json!({
849            "messages": [
850                {"role": "user", "content": "What's the capital of France?"}
851            ]
852        });
853        let cl = classify(&body);
854        assert_eq!(cl.category, Category::Chat);
855        assert_eq!(cl.complexity, Complexity::Quick); // short simple question
856        assert!(!cl.needs_tools);
857        assert!(!cl.has_media_inputs);
858    }
859
860    #[test]
861    fn test_classify_deep_analysis() {
862        let body = json!({
863            "messages": [
864                {"role": "user", "content": "Design a system architecture for a distributed database with strong consistency guarantees. Provide a detailed analysis of the trade-offs between CAP theorem constraints and explain how to handle network partitions in depth."}
865            ]
866        });
867        let cl = classify(&body);
868        assert_eq!(cl.complexity, Complexity::Deep);
869    }
870
871    #[test]
872    fn test_classify_code_with_tools() {
873        // Code request that happens to have tools — should be Code, not ToolCall
874        let body = json!({
875            "messages": [{"role": "user", "content": "Write a Python function to sort a list and debug it"}],
876            "tools": [{"type": "function", "function": {"name": "bash"}}]
877        });
878        let cl = classify(&body);
879        assert_eq!(cl.category, Category::Code);
880        assert!(cl.needs_tools);
881    }
882
883    #[test]
884    fn test_classify_tools_schema_always_needs_tools() {
885        // Tools schema present = agentic session, always needs_tools
886        // even if the message content is plain chat
887        let body = json!({
888            "messages": [{"role": "user", "content": "hello"}],
889            "tools": [{"type": "function", "function": {"name": "bash"}}]
890        });
891        let cl = classify(&body);
892        assert!(cl.needs_tools);
893    }
894
895    #[test]
896    fn test_classify_tools_schema_with_tool_content() {
897        // Tools in schema AND content implies tool use — needs tools
898        let body = json!({
899            "messages": [{"role": "user", "content": "Read the file and fix the bug"}],
900            "tools": [{"type": "function", "function": {"name": "read"}}]
901        });
902        let cl = classify(&body);
903        assert!(cl.needs_tools);
904    }
905
906    #[test]
907    fn test_classify_anthropic_text_blocks_with_tools() {
908        // Anthropic-style content blocks should still be parsed as text
909        // and trigger needs_tools when tool-intent is present.
910        let body = json!({
911            "messages": [
912                {
913                    "role": "user",
914                    "content": [
915                        {"type": "text", "text": "List files in this directory and read README.md"}
916                    ]
917                }
918            ],
919            "tools": [{"name": "shell"}]
920        });
921        let cl = classify(&body);
922        assert!(cl.needs_tools);
923        assert!(matches!(cl.category, Category::Code | Category::ToolCall));
924    }
925
926    #[test]
927    fn test_classify_anthropic_tool_use_block_sets_needs_tools() {
928        // If an explicit tool_use/tool_result block is present, mark as needs_tools.
929        let body = json!({
930            "messages": [
931                {
932                    "role": "assistant",
933                    "content": [
934                        {"type": "tool_use", "id": "toolu_123", "name": "shell", "input": {"command": "ls"}}
935                    ]
936                }
937            ]
938        });
939        let cl = classify(&body);
940        assert!(cl.needs_tools);
941    }
942
943    #[test]
944    fn test_anthropic_tool_request_prefers_stronger_tool_model() {
945        // Reproduces Claude-like tool request shape and verifies needs_tools=true
946        // pushes selection toward the stronger tool-capable model.
947        let body = json!({
948            "messages": [
949                {
950                    "role": "user",
951                    "content": [
952                        {"type": "text", "text": "List files in this directory and read README.md"}
953                    ]
954                }
955            ],
956            "tools": [{"name": "shell"}]
957        });
958        let cl = classify(&body);
959        assert!(cl.needs_tools);
960
961        let available = vec![("Qwen3-8B-Q4_K_M", 40.0), ("MiniMax-M2.5-Q4_K_M", 20.0)];
962        let picked = pick_model_classified(&cl, &available);
963        assert_eq!(picked, Some("MiniMax-M2.5-Q4_K_M"));
964    }
965
966    #[test]
967    fn test_classify_system_prompt_code() {
968        let body = json!({
969            "messages": [
970                {"role": "system", "content": "You are a senior developer and coding assistant."},
971                {"role": "user", "content": "Help me with this."}
972            ]
973        });
974        assert_eq!(classify(&body).category, Category::Code);
975    }
976
977    #[test]
978    fn test_media_requirements_detect_audio_block() {
979        let body = json!({
980            "messages": [
981                {
982                    "role": "user",
983                    "content": [
984                        {"type": "text", "text": "Transcribe this clip"},
985                        {"type": "audio_url", "audio_url": {"url": "mesh://blob/client-1/example"}}
986                    ]
987                }
988            ]
989        });
990        let media = media_requirements(&body);
991        assert!(media.has_media);
992        assert!(media.needs_audio);
993        assert!(!media.needs_vision);
994        assert!(classify(&body).has_media_inputs);
995    }
996
997    #[test]
998    fn test_media_requirements_detect_image_block() {
999        let body = json!({
1000            "messages": [
1001                {
1002                    "role": "user",
1003                    "content": [
1004                        {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}
1005                    ]
1006                }
1007            ]
1008        });
1009        let media = media_requirements(&body);
1010        assert!(media.has_media);
1011        assert!(media.needs_vision);
1012        assert!(!media.needs_audio);
1013        assert!(classify(&body).has_media_inputs);
1014    }
1015
1016    #[test]
1017    fn test_pick_model_primary_strength_wins() {
1018        // Qwen3-8B (tier 2, Chat primary) and 235B (tier 4, Chat 3rd) score within
1019        // 15 points at Moderate complexity, so either is a valid pick (load spread).
1020        let available = vec![("Qwen3-8B-Q4_K_M", 50.0), ("Qwen3-235B-A22B-Q4_K_M", 20.0)];
1021        let result = pick_model_classified(
1022            &Classification {
1023                category: Category::Chat,
1024                complexity: Complexity::Moderate,
1025                needs_tools: false,
1026                has_media_inputs: false,
1027            },
1028            &available,
1029        );
1030        assert!(result == Some("Qwen3-8B-Q4_K_M") || result == Some("Qwen3-235B-A22B-Q4_K_M"));
1031    }
1032
1033    #[test]
1034    fn test_deep_complexity_prefers_bigger() {
1035        // Deep complexity amplifies tier bonus, but scores are within 15 points
1036        // so load spread makes either a valid pick.
1037        let available = vec![("Qwen3-8B-Q4_K_M", 50.0), ("Qwen3-235B-A22B-Q4_K_M", 20.0)];
1038        let result = pick_model_classified(
1039            &Classification {
1040                category: Category::Chat,
1041                complexity: Complexity::Deep,
1042                needs_tools: false,
1043                has_media_inputs: false,
1044            },
1045            &available,
1046        );
1047        assert!(result == Some("Qwen3-8B-Q4_K_M") || result == Some("Qwen3-235B-A22B-Q4_K_M"));
1048    }
1049
1050    #[test]
1051    fn test_quick_complexity_prefers_smaller() {
1052        // Quick complexity: both score within 15 points (load spread applies).
1053        // Either is a valid pick — the key property is neither is excluded.
1054        let available = vec![
1055            ("Qwen3-8B-Q4_K_M", 50.0),
1056            ("Qwen2.5-72B-Instruct-Q4_K_M", 10.0),
1057        ];
1058        let result = pick_model_classified(
1059            &Classification {
1060                category: Category::Chat,
1061                complexity: Complexity::Quick,
1062                needs_tools: false,
1063                has_media_inputs: false,
1064            },
1065            &available,
1066        );
1067        assert!(result == Some("Qwen3-8B-Q4_K_M") || result == Some("Qwen2.5-72B-Instruct-Q4_K_M"));
1068    }
1069
1070    #[test]
1071    fn test_pick_model_prefers_strength_match() {
1072        // Same tier, same speed — scores within 15 points, load spread applies.
1073        let available = vec![
1074            ("DeepSeek-R1-Distill-70B-Q4_K_M", 10.0), // tier 3, reasoning specialist
1075            ("Qwen2.5-72B-Instruct-Q4_K_M", 10.0),    // tier 3, chat primary
1076        ];
1077        let result = pick_model_classified(
1078            &Classification {
1079                category: Category::Reasoning,
1080                complexity: Complexity::Moderate,
1081                needs_tools: false,
1082                has_media_inputs: false,
1083            },
1084            &available,
1085        );
1086        assert!(
1087            result == Some("DeepSeek-R1-Distill-70B-Q4_K_M")
1088                || result == Some("Qwen2.5-72B-Instruct-Q4_K_M")
1089        );
1090    }
1091
1092    #[test]
1093    fn test_pick_model_code_specialist() {
1094        // Same tier, same speed — scores within 15 points, load spread applies.
1095        let available = vec![
1096            ("Qwen2.5-Coder-32B-Instruct-Q4_K_M", 15.0),
1097            ("Qwen2.5-32B-Instruct-Q4_K_M", 15.0),
1098        ];
1099        let result = pick_model_classified(
1100            &Classification {
1101                category: Category::Code,
1102                complexity: Complexity::Moderate,
1103                needs_tools: false,
1104                has_media_inputs: false,
1105            },
1106            &available,
1107        );
1108        assert!(
1109            result == Some("Qwen2.5-Coder-32B-Instruct-Q4_K_M")
1110                || result == Some("Qwen2.5-32B-Instruct-Q4_K_M")
1111        );
1112    }
1113
1114    #[test]
1115    fn test_pick_model_empty() {
1116        let available: Vec<(&str, f64)> = vec![];
1117        assert_eq!(
1118            pick_model_classified(
1119                &Classification {
1120                    category: Category::Chat,
1121                    complexity: Complexity::Moderate,
1122                    needs_tools: false,
1123                    has_media_inputs: false,
1124                },
1125                &available
1126            ),
1127            None
1128        );
1129    }
1130
1131    #[test]
1132    fn test_pick_model_unknown_model_still_works() {
1133        let available = vec![("SomeUnknownModel", 30.0)];
1134        let result = pick_model_classified(
1135            &Classification {
1136                category: Category::Chat,
1137                complexity: Complexity::Moderate,
1138                needs_tools: false,
1139                has_media_inputs: false,
1140            },
1141            &available,
1142        );
1143        assert_eq!(result, Some("SomeUnknownModel"));
1144    }
1145
1146    #[test]
1147    fn test_profile_lookup() {
1148        assert!(profile_for("Qwen3-235B-A22B-Q4_K_M").is_some());
1149        assert_eq!(profile_for("Qwen3-235B-A22B-Q4_K_M").unwrap().tier, 4);
1150        assert!(profile_for("nonexistent").is_none());
1151    }
1152
1153    #[test]
1154    fn test_all_profiles_have_strengths() {
1155        for p in MODEL_PROFILES {
1156            assert!(!p.strengths.is_empty(), "{} has no strengths", p.name);
1157        }
1158    }
1159
1160    #[test]
1161    fn test_classify_empty_tools_is_not_tool_call() {
1162        let body = json!({
1163            "messages": [{"role": "user", "content": "hello"}],
1164            "tools": []
1165        });
1166        assert_eq!(classify(&body).category, Category::Chat);
1167    }
1168
1169    #[test]
1170    fn test_strip_split_suffix() {
1171        assert_eq!(
1172            strip_split_suffix("MiniMax-M2.5-Q4_K_M-00001-of-00004"),
1173            "MiniMax-M2.5-Q4_K_M"
1174        );
1175        assert_eq!(
1176            strip_split_suffix("Qwen3-Coder-Next-Q4_K_M-00001-of-00004"),
1177            "Qwen3-Coder-Next-Q4_K_M"
1178        );
1179        assert_eq!(
1180            strip_split_suffix("Hermes-2-Pro-Mistral-7B-Q4_K_M"),
1181            "Hermes-2-Pro-Mistral-7B-Q4_K_M"
1182        );
1183        assert_eq!(strip_split_suffix(""), "");
1184    }
1185
1186    #[test]
1187    fn test_profile_for_split_gguf() {
1188        let p = profile_for("MiniMax-M2.5-Q4_K_M-00001-of-00004");
1189        assert!(p.is_some());
1190        assert_eq!(p.unwrap().name, "MiniMax-M2.5-Q4_K_M");
1191        assert_eq!(p.unwrap().tier, 4);
1192    }
1193}
1194
1195#[test]
1196fn test_tools_filter_prefers_capable() {
1197    let available = vec![
1198        ("DeepSeek-R1-Distill-Qwen-32B-Q4_K_M", 10.0), // tools: false, Reasoning only
1199        ("Qwen2.5-32B-Instruct-Q4_K_M", 50.0),         // tools: true, Chat+Reasoning+Code
1200    ];
1201    // Without tools, Reasoning request: scores within 15 points, either valid
1202    let result = pick_model_with_tools(Category::Reasoning, &available, false);
1203    assert!(
1204        result == Some("DeepSeek-R1-Distill-Qwen-32B-Q4_K_M")
1205            || result == Some("Qwen2.5-32B-Instruct-Q4_K_M")
1206    );
1207    // With tools, Reasoning request: Qwen wins (DeepSeek filtered out — can't do tools)
1208    let result = pick_model_with_tools(Category::Reasoning, &available, true);
1209    assert_eq!(result, Some("Qwen2.5-32B-Instruct-Q4_K_M"));
1210}
1211
1212#[test]
1213fn test_tools_filter_fallback_when_none_capable() {
1214    let available = vec![
1215        ("DeepSeek-R1-Distill-Qwen-32B-Q4_K_M", 10.0), // tools: false
1216    ];
1217    // With tools required but nothing capable: falls back to available
1218    let result = pick_model_with_tools(Category::Reasoning, &available, true);
1219    assert_eq!(result, Some("DeepSeek-R1-Distill-Qwen-32B-Q4_K_M"));
1220}
1221
1222#[test]
1223fn test_agentic_prefers_strongest_model() {
1224    // Agentic (needs_tools=true): 32B (tier 3) should beat 7B (tier 2) even though 7B is faster
1225    let available = vec![
1226        ("Hermes-2-Pro-Mistral-7B-Q4_K_M", 87.0), // tier 2, tools: false
1227        ("Qwen2.5-Coder-7B-Instruct-Q4_K_M", 85.0), // tier 2, tools: true
1228        ("Qwen2.5-32B-Instruct-Q4_K_M", 18.0),    // tier 3, tools: true
1229    ];
1230    let cl = Classification {
1231        category: Category::Code,
1232        complexity: Complexity::Moderate,
1233        needs_tools: true,
1234        has_media_inputs: false,
1235    };
1236    let result = pick_model_classified(&cl, &available);
1237    // 32B should win: tier 3×20=60 beats Coder tier 2×20=40, despite lower speed
1238    assert_eq!(result, Some("Qwen2.5-32B-Instruct-Q4_K_M"));
1239}
1240
1241#[test]
1242fn test_chat_prefers_fastest_model() {
1243    // Chat (needs_tools=false, Quick): scores within 15 points, load spread applies.
1244    let available = vec![
1245        ("Hermes-2-Pro-Mistral-7B-Q4_K_M", 87.0), // tier 2, fast
1246        ("Qwen2.5-32B-Instruct-Q4_K_M", 18.0),    // tier 3, slow
1247    ];
1248    let cl = Classification {
1249        category: Category::Chat,
1250        complexity: Complexity::Quick,
1251        needs_tools: false,
1252        has_media_inputs: false,
1253    };
1254    let result = pick_model_classified(&cl, &available);
1255    assert!(
1256        result == Some("Hermes-2-Pro-Mistral-7B-Q4_K_M")
1257            || result == Some("Qwen2.5-32B-Instruct-Q4_K_M")
1258    );
1259}
1260
1261#[test]
1262fn test_agentic_deep_strongly_prefers_biggest() {
1263    // Deep agentic: tier 4 should massively beat tier 2
1264    let available = vec![
1265        ("Qwen2.5-Coder-7B-Instruct-Q4_K_M", 85.0), // tier 2
1266        ("MiniMax-M2.5-Q4_K_M", 21.0),              // tier 4
1267    ];
1268    let cl = Classification {
1269        category: Category::Code,
1270        complexity: Complexity::Deep,
1271        needs_tools: true,
1272        has_media_inputs: false,
1273    };
1274    let result = pick_model_classified(&cl, &available);
1275    assert_eq!(result, Some("MiniMax-M2.5-Q4_K_M"));
1276}