tokenmiser_router/
tier1.rs

1//! Tier 1 semantic classifier: difficulty exemplars are embedded at startup
2//! with the same `bge-small-en-v1.5` model the L2 cache loads, and a request
3//! takes the difficulty of its nearest exemplar.
4//!
5//! Less accurate than a fine-tuned classifier, but it adds no dependencies and
6//! reuses the existing embedder. Swapping in trained weights is a change to
7//! `classify()` alone.
8
9use anyhow::{anyhow, Result};
10use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
11use parking_lot::Mutex;
12use tokenmiser_providers::ChatRequest;
13use tracing::info;
14
15use crate::Difficulty;
16
17struct Exemplar {
18    embedding: Vec<f32>,
19    difficulty: Difficulty,
20}
21
22pub struct Tier1Classifier {
23    embedder: Mutex<TextEmbedding>,
24    exemplars: Vec<Exemplar>,
25}
26
27impl Tier1Classifier {
28    /// Build the classifier and pre-embed the default exemplar set.
29    pub fn new() -> Result<Self> {
30        let opts = InitOptions::new(EmbeddingModel::BGESmallENV15);
31        let mut embedder = TextEmbedding::try_new(opts)
32            .map_err(|e| anyhow!("bge-small init for Tier1 failed: {e}"))?;
33
34        let raw = default_exemplars();
35        let texts: Vec<&str> = raw.iter().map(|(t, _)| *t).collect();
36        let embeddings = embedder
37            .embed(texts, None)
38            .map_err(|e| anyhow!("Tier1 exemplar embed failed: {e}"))?;
39
40        let exemplars = embeddings
41            .into_iter()
42            .zip(raw.iter())
43            .map(|(emb, (_, d))| Exemplar {
44                embedding: emb,
45                difficulty: *d,
46            })
47            .collect::<Vec<_>>();
48
49        info!(
50            exemplars = exemplars.len(),
51            "Tier1 semantic classifier ready"
52        );
53        Ok(Self {
54            embedder: Mutex::new(embedder),
55            exemplars,
56        })
57    }
58
59    /// Classify a request by finding the nearest exemplar.
60    pub fn classify(&self, req: &ChatRequest) -> Difficulty {
61        let text = req
62            .messages
63            .iter()
64            .filter(|m| m.role == "user")
65            .filter_map(|m| match &m.content {
66                serde_json::Value::String(s) => Some(s.clone()),
67                _ => None,
68            })
69            .collect::<Vec<_>>()
70            .join("\n");
71
72        if text.trim().is_empty() {
73            return Difficulty::Medium;
74        }
75
76        let mut e = self.embedder.lock();
77        let emb = match e.embed(vec![text.as_str()], None) {
78            Ok(mut v) => match v.pop() {
79                Some(x) => x,
80                None => return Difficulty::Medium,
81            },
82            Err(_) => return Difficulty::Medium,
83        };
84        drop(e);
85
86        let mut best: Option<(f32, Difficulty)> = None;
87        for ex in &self.exemplars {
88            let sim = cosine(&emb, &ex.embedding);
89            if best.map(|(b, _)| sim > b).unwrap_or(true) {
90                best = Some((sim, ex.difficulty));
91            }
92        }
93        best.map(|(_, d)| d).unwrap_or(Difficulty::Medium)
94    }
95}
96
97fn cosine(a: &[f32], b: &[f32]) -> f32 {
98    let mut dot = 0.0;
99    let mut na = 0.0;
100    let mut nb = 0.0;
101    for i in 0..a.len().min(b.len()) {
102        dot += a[i] * b[i];
103        na += a[i] * a[i];
104        nb += b[i] * b[i];
105    }
106    if na == 0.0 || nb == 0.0 {
107        return 0.0;
108    }
109    dot / (na.sqrt() * nb.sqrt())
110}
111
112/// Exemplars for the three difficulty bands, kept small to bound startup
113/// embedding time.
114fn default_exemplars() -> Vec<(&'static str, Difficulty)> {
115    vec![
116        // EASY: short factual queries, simple transforms.
117        ("what is the capital of france", Difficulty::Easy),
118        ("translate 'hello' to spanish", Difficulty::Easy),
119        ("define photosynthesis in one sentence", Difficulty::Easy),
120        (
121            "summarize: the cat sat on the mat. the mat was red.",
122            Difficulty::Easy,
123        ),
124        ("is 17 a prime number?", Difficulty::Easy),
125        (
126            "classify this sentence as positive or negative: the food was great",
127            Difficulty::Easy,
128        ),
129        ("convert 32 fahrenheit to celsius", Difficulty::Easy),
130        // MEDIUM: light reasoning, longer prose, structured outputs.
131        ("write a haiku about programming", Difficulty::Medium),
132        (
133            "compare REST and GraphQL in three bullet points",
134            Difficulty::Medium,
135        ),
136        (
137            "explain how DNS resolution works to a junior dev",
138            Difficulty::Medium,
139        ),
140        (
141            "draft an email apologizing for a missed deadline",
142            Difficulty::Medium,
143        ),
144        (
145            "what are the trade-offs between SQL and NoSQL databases",
146            Difficulty::Medium,
147        ),
148        // HARD: code edits, architectural reasoning, multi-step proofs.
149        (
150            "refactor this 400-line authentication middleware to use JWT tokens",
151            Difficulty::Hard,
152        ),
153        (
154            "design a distributed rate limiter that handles 1M requests per second",
155            Difficulty::Hard,
156        ),
157        (
158            "prove that this sorting algorithm terminates in O(n log n) time",
159            Difficulty::Hard,
160        ),
161        (
162            "debug this race condition in my concurrent queue implementation",
163            Difficulty::Hard,
164        ),
165        (
166            "implement a B-tree with concurrent insertions",
167            Difficulty::Hard,
168        ),
169        (
170            "optimize this SQL query that joins seven tables for sub-100ms latency",
171            Difficulty::Hard,
172        ),
173    ]
174}
tokenmiser_router/tier1.rs

tokenmiser_router/
tier1.rs
