vex_router/classifier/
mod.rs1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
6pub struct QueryComplexity {
7 pub score: f64,
8 pub capabilities: Vec<String>,
9 pub estimated_tokens: usize,
10 pub confidence: f64,
11}
12
13#[derive(Debug)]
14pub struct QueryClassifier;
15
16impl QueryClassifier {
17 pub fn new() -> Self {
18 Self
19 }
20
21 pub fn classify(&self, query: &str) -> QueryComplexity {
22 let query_lower = query.to_lowercase();
23 let word_count = query_lower.split_whitespace().count();
24 let estimated_tokens = (word_count as f64 * 1.3) as usize;
25
26 let mut capabilities = vec!["general".to_string()];
27
28 if query_lower.contains("code")
29 || query_lower.contains("function")
30 || query_lower.contains("implement")
31 {
32 capabilities.push("code".to_string());
33 }
34 if query_lower.contains("math") || query_lower.contains("calculate") {
35 capabilities.push("math".to_string());
36 }
37 if query_lower.contains("analyze") || query_lower.contains("compare") {
38 capabilities.push("analysis".to_string());
39 }
40
41 let score = self.calculate_complexity(query_lower.len(), word_count);
42
43 QueryComplexity {
44 score,
45 capabilities,
46 estimated_tokens,
47 confidence: 0.7,
48 }
49 }
50
51 fn calculate_complexity(&self, char_count: usize, word_count: usize) -> f64 {
52 let mut score: f64 = 0.1;
53
54 if word_count > 50 {
55 score += 0.3;
56 } else if word_count > 20 {
57 score += 0.2;
58 } else if word_count > 10 {
59 score += 0.1;
60 }
61
62 if char_count > 500 {
63 score += 0.2;
64 }
65
66 score = score.clamp(0.05, 1.0);
67
68 score
69 }
70}
71
72impl Default for QueryClassifier {
73 fn default() -> Self {
74 Self::new()
75 }
76}