lean_ctx/core/triage/
model_loader.rs

1//! Optional local ONNX classifier for task triage.
2//!
3//! A missing model is normal: callers receive `None` and rule-based triage
4//! continues without any download or runtime dependency.
5
6use std::path::{Path, PathBuf};
7
8use super::{ProfileHypothesis, TaskAnalysisInput};
9#[cfg(feature = "onnx")]
10use super::{
11    TriageBackendLocal,
12    profile::{TaskProfileLocal, TaskScopeLocal},
13};
14
15pub const TRIAGE_MODEL_FILE: &str = "triage-v1.onnx";
16#[cfg(feature = "onnx")]
17const MAX_SEQUENCE_LENGTH: usize = 64;
18
19pub fn model_path(data_dir: &Path) -> PathBuf {
20    data_dir.join("models").join(TRIAGE_MODEL_FILE)
21}
22
23pub fn is_managed_model_path(data_dir: &Path, candidate: &Path) -> bool {
24    candidate == model_path(data_dir)
25}
26
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct TriageLogits {
29    pub intent: f32,
30    pub complexity: f32,
31    pub scope: f32,
32    pub reasoning_need: f32,
33    pub risk: f32,
34}
35
36/// Minimal whitespace/subword-compatible BERT input encoder.
37/// Unknown pieces use the standard BERT `[UNK]` ID; this avoids inventing IDs
38/// for a model vocabulary that is not part of the managed model artifact.
39#[cfg(feature = "onnx")]
40#[derive(Debug, Clone, Copy)]
41struct BertTinyTokenizer;
42
43#[cfg(feature = "onnx")]
44impl BertTinyTokenizer {
45    fn encode(self, text: &str) -> (Vec<i64>, Vec<i64>, Vec<i64>) {
46        const CLS: i64 = 101;
47        const SEP: i64 = 102;
48        const PAD: i64 = 0;
49        const UNK: i64 = 100;
50        let mut ids = vec![CLS];
51        for token in text
52            .split(|character: char| character.is_whitespace() || character.is_ascii_punctuation())
53            .filter(|token| !token.is_empty())
54        {
55            if ids.len() + 1 >= MAX_SEQUENCE_LENGTH {
56                break;
57            }
58            let id = match token.to_ascii_lowercase().as_str() {
59                "fix" => 8081,
60                "bug" => 11829,
61                "test" => 3231,
62                "refactor" => 10788,
63                "config" => 6149,
64                "deploy" => 21296,
65                "review" => 3319,
66                "debug" => 8567,
67                _ => UNK,
68            };
69            ids.push(id);
70        }
71        ids.push(SEP);
72        let mut mask = vec![1; ids.len()];
73        let mut types = vec![0; ids.len()];
74        ids.resize(MAX_SEQUENCE_LENGTH, PAD);
75        mask.resize(MAX_SEQUENCE_LENGTH, 0);
76        types.resize(MAX_SEQUENCE_LENGTH, 0);
77        (ids, mask, types)
78    }
79}
80
81#[derive(Debug)]
82pub struct ModelLoader {
83    model_path: PathBuf,
84    #[cfg(feature = "onnx")]
85    session: std::sync::Mutex<ort::session::Session>,
86    #[cfg(feature = "onnx")]
87    input_names: Vec<String>,
88    #[cfg(feature = "onnx")]
89    output_name: String,
90}
91
92impl ModelLoader {
93    pub fn load(data_dir: &Path) -> anyhow::Result<Option<Self>> {
94        Self::from_model_path(model_path(data_dir))
95    }
96
97    pub fn from_model_path(model_path: impl Into<PathBuf>) -> anyhow::Result<Option<Self>> {
98        let model_path = model_path.into();
99        if !model_path.is_file() {
100            return Ok(None);
101        }
102        #[cfg(feature = "onnx")]
103        {
104            let execution_providers = crate::core::ort_execution_providers::execution_providers();
105            crate::core::ort_environment::ensure_ort_env(&execution_providers)?;
106            let session = ort::session::Session::builder()
107                .map_err(|error| anyhow::anyhow!("ORT builder: {error}"))?
108                .commit_from_file(&model_path)
109                .map_err(|error| anyhow::anyhow!("ORT load triage model: {error}"))?;
110            let input_names: Vec<_> = session
111                .inputs()
112                .iter()
113                .map(|input| input.name().to_owned())
114                .collect();
115            let output_name = session
116                .outputs()
117                .first()
118                .map(|output| output.name().to_owned())
119                .ok_or_else(|| anyhow::anyhow!("triage model has no outputs"))?;
120            if input_names.len() < 2 {
121                anyhow::bail!("triage model has an unsupported graph signature");
122            }
123            Ok(Some(Self {
124                model_path,
125                session: std::sync::Mutex::new(session),
126                input_names,
127                output_name,
128            }))
129        }
130        #[cfg(not(feature = "onnx"))]
131        {
132            let _ = model_path;
133            Ok(None)
134        }
135    }
136
137    pub fn path(&self) -> &Path {
138        &self.model_path
139    }
140
141    pub fn infer(&self, input: &TaskAnalysisInput) -> anyhow::Result<ProfileHypothesis> {
142        #[cfg(feature = "onnx")]
143        {
144            let text = if input.query.is_empty() {
145                input.session_context.as_deref().unwrap_or_default()
146            } else {
147                &input.query
148            };
149            let (ids, mask, types) = BertTinyTokenizer.encode(text);
150            let ids = ort::value::Tensor::from_array(ndarray::Array2::from_shape_vec(
151                (1, MAX_SEQUENCE_LENGTH),
152                ids,
153            )?)?;
154            let mask = ort::value::Tensor::from_array(ndarray::Array2::from_shape_vec(
155                (1, MAX_SEQUENCE_LENGTH),
156                mask,
157            )?)?;
158            let types = ort::value::Tensor::from_array(ndarray::Array2::from_shape_vec(
159                (1, MAX_SEQUENCE_LENGTH),
160                types,
161            )?)?;
162            let mut session = self
163                .session
164                .lock()
165                .unwrap_or_else(std::sync::PoisonError::into_inner);
166            let output = if self.input_names.len() >= 3 {
167                session.run(ort::inputs![
168                    self.input_names[0].as_str() => ids,
169                    self.input_names[1].as_str() => mask,
170                    self.input_names[2].as_str() => types,
171                ])?
172            } else {
173                session.run(ort::inputs![
174                    self.input_names[0].as_str() => ids,
175                    self.input_names[1].as_str() => mask,
176                ])?
177            };
178            let (_, values) = output[self.output_name.as_str()].try_extract_tensor::<f32>()?;
179            let logits = TriageLogits {
180                intent: *values
181                    .first()
182                    .ok_or_else(|| anyhow::anyhow!("triage model returned no logits"))?,
183                complexity: *values
184                    .get(1)
185                    .ok_or_else(|| anyhow::anyhow!("triage model returned fewer than 5 logits"))?,
186                scope: *values
187                    .get(2)
188                    .ok_or_else(|| anyhow::anyhow!("triage model returned fewer than 5 logits"))?,
189                reasoning_need: *values
190                    .get(3)
191                    .ok_or_else(|| anyhow::anyhow!("triage model returned fewer than 5 logits"))?,
192                risk: *values
193                    .get(4)
194                    .ok_or_else(|| anyhow::anyhow!("triage model returned fewer than 5 logits"))?,
195            };
196            Ok(hypothesis_from_logits(logits))
197        }
198        #[cfg(not(feature = "onnx"))]
199        {
200            let _ = input;
201            anyhow::bail!("triage ONNX support is disabled")
202        }
203    }
204}
205
206#[cfg(feature = "onnx")]
207fn hypothesis_from_logits(logits: TriageLogits) -> ProfileHypothesis {
208    let confidence = probability(logits.intent).min(probability(logits.complexity));
209    let scope = probability(logits.scope);
210    let profile = TaskProfileLocal {
211        task_class: "coding".into(),
212        intent: if logits.intent >= 0.0 {
213            "coding_change"
214        } else {
215            "explore"
216        }
217        .into(),
218        complexity: if logits.complexity > 0.5 {
219            "high"
220        } else if logits.complexity > -0.5 {
221            "medium"
222        } else {
223            "low"
224        }
225        .into(),
226        scope: if scope > 0.8 {
227            TaskScopeLocal::CrossProject
228        } else if scope > 0.6 {
229            TaskScopeLocal::CrossModule
230        } else if scope > 0.4 {
231            TaskScopeLocal::MultiFile
232        } else {
233            TaskScopeLocal::SingleFile
234        },
235        context_need_milli: (scope * 1000.0).round() as u16,
236        reasoning_need_milli: (probability(logits.reasoning_need) * 1000.0).round() as u16,
237        risk_signal_milli: (probability(logits.risk) * 1000.0).round() as u16,
238        confidence_milli: (confidence * 1000.0).round() as u16,
239    };
240    ProfileHypothesis {
241        confidence_milli: profile.confidence_milli,
242        profile,
243        backend: TriageBackendLocal::Semantic,
244    }
245}
246
247#[cfg(feature = "onnx")]
248fn probability(logit: f32) -> f32 {
249    if logit.is_finite() {
250        1.0 / (1.0 + (-logit).exp())
251    } else {
252        0.0
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn absent_model_returns_none() {
262        let data = tempfile::tempdir().unwrap();
263        assert!(ModelLoader::load(data.path()).unwrap().is_none());
264    }
265
266    #[test]
267    fn managed_path_cannot_escape_data_dir() {
268        let data = tempfile::tempdir().unwrap();
269        assert!(is_managed_model_path(data.path(), &model_path(data.path())));
270        assert!(!is_managed_model_path(
271            data.path(),
272            &data.path().join("models/../other.onnx")
273        ));
274    }
275}
lean_ctx/core/triage/model_loader.rs

lean_ctx/core/triage/
model_loader.rs
