switchyard_libsy/core/
classifier.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::core::algorithm::{AnswerSummary, Driver};
5use crate::core::state::State;
6use crate::{LibsyError, Result};
7use async_trait::async_trait;
8use switchyard_protocol::{Category, ModelId, Request, Response};
9
10/// One classifier's recommendation of a routing `target`, with a `[0.0, 1.0]` confidence.
11#[derive(Debug, Clone, PartialEq)]
12pub struct Score {
13    /// `[0.0, 1.0]` confidence in `target`.
14    pub confidence: f64,
15    /// The target (model / tier) being recommended.
16    pub target: ModelId,
17    /// The category `target` was drawn from, when the classifier picked one. The rest of
18    /// that category is what the turn falls through on failure. A decision made without
19    /// a category leaves this `None`.
20    pub category: Option<Category>,
21}
22
23/// A classifier's verdict for a request: a set of target [`Score`]s, flagged by how
24/// confident the classifier is that they are decisive.
25pub enum Classification {
26    /// Definite recommendations; [`argmax`](Self::argmax) always yields the top target.
27    Scores(Vec<Score>),
28    /// Recommendations the classifier considers ambiguous; [`argmax`](Self::argmax) yields
29    /// nothing unless the caller opts to ignore ambiguity.
30    Ambiguous(Vec<Score>),
31}
32
33impl Classification {
34    /// The top-scoring [`Score`], or `None` when the classifier abstained (an empty set).
35    ///
36    /// An [`Ambiguous`](Self::Ambiguous) classification also yields `None` unless
37    /// `ignore_ambiguous` is set, in which case it falls back to the plain argmax.
38    /// Errors if any confidence is `NaN` (an unorderable score the caller should surface).
39    pub fn argmax(&self, ignore_ambiguous: bool) -> Result<Option<Score>> {
40        match self {
41            Classification::Scores(scores) => argmax(scores),
42            Classification::Ambiguous(scores) => {
43                if ignore_ambiguous {
44                    argmax(scores)
45                } else {
46                    Ok(None)
47                }
48            }
49        }
50    }
51}
52
53/// The highest-confidence score, or `None` when the set is empty (the classifier abstained).
54/// Ties keep the first. Errors on a `NaN` confidence,
55/// which has no defined ordering.
56fn argmax(scores: &[Score]) -> Result<Option<Score>> {
57    let mut best: Option<&Score> = None;
58    for score in scores.iter() {
59        if score.confidence.is_nan() {
60            return Err(LibsyError::AlgorithmError {
61                message: format!(
62                    "classifier returned NaN confidence for target {:?}",
63                    score.target
64                ),
65            });
66        }
67        match best {
68            Some(cur_best) if score.confidence > cur_best.confidence => best = Some(score),
69            None => best = Some(score),
70            _ => {}
71        }
72    }
73    Ok(best.cloned())
74}
75
76/// Scores targets from the current request and the composition's state.
77#[async_trait]
78pub trait Classifier: Send + Sync {
79    /// Score the classifier's targets given the current state and request.
80    ///
81    /// `driver` lets a classifier inspect runtime models and offload model calls.
82    ///
83    /// `request` is borrowed mutably so a classifier may rewrite it in place — inject a
84    /// system prompt, drop tools, compact history. The edit is not scoped to this call:
85    /// later classifiers in the cascade score the rewritten request, and it is the
86    /// rewritten request that is finally sent to the selected model. Most classifiers
87    /// only read it.
88    async fn score(
89        &self,
90        state: &mut State,
91        request: &mut Request,
92        driver: &Driver,
93    ) -> Result<(Classification, Option<Response>)>;
94
95    /// Whether a stateful composition should request served-answer reports.
96    fn tracks_answers(&self) -> bool {
97        false
98    }
99
100    /// Folds the answer the host served for this session into the composition's state.
101    fn observe_answer(&self, _state: &mut State, _answer: AnswerSummary) {}
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use crate::core::testing::empty_driver;
108    use switchyard_protocol::text_request;
109
110    /// Terse `Score` builder for the assertions below.
111    fn score(target: &str, confidence: f64) -> Score {
112        Score {
113            target: ModelId::from(target),
114            confidence,
115            category: None,
116        }
117    }
118
119    #[test]
120    fn argmax_picks_the_highest_confidence_score() -> Result<()> {
121        let scores = vec![score("weak", 0.2), score("strong", 0.9), score("mid", 0.5)];
122        let best = Classification::Scores(scores).argmax(false)?;
123        assert_eq!(best, Some(score("strong", 0.9)));
124        Ok(())
125    }
126
127    #[test]
128    fn argmax_breaks_ties_by_cascade_order() -> Result<()> {
129        // Equal confidence: the earlier target in cascade order wins the tie.
130        let scores = vec![score("first", 0.7), score("second", 0.7)];
131        let best = Classification::Scores(scores).argmax(false)?;
132        assert_eq!(best.map(|s| s.target), Some(ModelId::from("first")));
133        Ok(())
134    }
135
136    #[test]
137    fn argmax_on_an_empty_set_abstains() -> Result<()> {
138        // No scores means the classifier abstained — no choice to make.
139        assert_eq!(Classification::Scores(vec![]).argmax(false)?, None);
140        assert_eq!(Classification::Ambiguous(vec![]).argmax(true)?, None);
141        Ok(())
142    }
143
144    #[test]
145    fn argmax_errors_on_nan_confidence() {
146        // A NaN confidence has no defined ordering — surface it rather than guess.
147        let scores = vec![score("weak", 0.3), score("strong", f64::NAN)];
148        assert!(matches!(
149            Classification::Scores(scores).argmax(false),
150            Err(LibsyError::AlgorithmError { message })
151                if message == "classifier returned NaN confidence for target \"strong\""
152        ));
153        // A lone NaN errors too, even with nothing to compare it against.
154        assert!(matches!(
155            Classification::Scores(vec![score("only", f64::NAN)]).argmax(false),
156            Err(LibsyError::AlgorithmError { message })
157                if message == "classifier returned NaN confidence for target \"only\""
158        ));
159    }
160
161    #[test]
162    fn ambiguous_without_ignore_makes_no_choice() -> Result<()> {
163        // Ambiguous means "don't pick" unless the caller opts to ignore ambiguity.
164        let scores = vec![score("strong", 0.9)];
165        assert_eq!(Classification::Ambiguous(scores).argmax(false)?, None);
166        Ok(())
167    }
168
169    #[test]
170    fn ambiguous_with_ignore_falls_back_to_argmax() -> Result<()> {
171        let scores = vec![score("weak", 0.3), score("strong", 0.8)];
172        let best = Classification::Ambiguous(scores).argmax(true)?;
173        assert_eq!(best, Some(score("strong", 0.8)));
174        Ok(())
175    }
176
177    #[test]
178    fn scores_variant_ignores_the_ambiguous_flag() -> Result<()> {
179        // A definitive classification always yields its argmax, regardless of the flag.
180        let scores = vec![score("a", 0.4), score("b", 0.6)];
181        let with_ignore = Classification::Scores(scores.clone()).argmax(true)?;
182        let without_ignore = Classification::Scores(scores).argmax(false)?;
183        assert_eq!(with_ignore, without_ignore);
184        assert_eq!(with_ignore, Some(score("b", 0.6)));
185        Ok(())
186    }
187
188    /// Scores the request's requested model at full confidence and records that it ran.
189    struct RecordingClassifier;
190
191    #[async_trait]
192    impl Classifier for RecordingClassifier {
193        async fn score(
194            &self,
195            state: &mut State,
196            request: &mut Request,
197            _driver: &Driver,
198        ) -> Result<(Classification, Option<Response>)> {
199            state
200                .extra
201                .insert("ran".to_string(), crate::core::state::StateValue::Count(1));
202            let target = request.model_id().unwrap_or(ModelId::from("auto"));
203            Ok((
204                Classification::Scores(vec![Score {
205                    target,
206                    confidence: 1.0,
207                    category: None,
208                }]),
209                None,
210            ))
211        }
212    }
213
214    #[tokio::test]
215    async fn classifier_reads_request_and_mutates_state() -> Result<()> {
216        let mut state = State::default();
217        let mut request = Request {
218            llm_request: text_request(Some("strong".to_string()), "hi"),
219            raw_request: None,
220            metadata: None,
221        };
222        let (classification, _) = RecordingClassifier
223            .score(&mut state, &mut request, &empty_driver())
224            .await?;
225        assert_eq!(
226            classification.argmax(false)?.map(|s| s.target),
227            Some(ModelId::from("strong"))
228        );
229        assert_eq!(
230            state.extra.get("ran"),
231            Some(&crate::core::state::StateValue::Count(1))
232        );
233        Ok(())
234    }
235
236    /// Rewrites the request's model, then scores the rewritten value.
237    struct RewritingClassifier;
238
239    #[async_trait]
240    impl Classifier for RewritingClassifier {
241        async fn score(
242            &self,
243            _state: &mut State,
244            request: &mut Request,
245            _driver: &Driver,
246        ) -> Result<(Classification, Option<Response>)> {
247            request.llm_request.model = Some("rewritten".to_string());
248            Ok((
249                Classification::Scores(vec![Score {
250                    target: ModelId::from("rewritten"),
251                    confidence: 1.0,
252                    category: None,
253                }]),
254                None,
255            ))
256        }
257    }
258
259    #[tokio::test]
260    async fn classifier_rewrites_the_request_in_place() -> Result<()> {
261        let mut state = State::default();
262        let mut request = Request {
263            llm_request: text_request(Some("auto".to_string()), "hi"),
264            raw_request: None,
265            metadata: None,
266        };
267
268        RewritingClassifier
269            .score(&mut state, &mut request, &empty_driver())
270            .await?;
271
272        // The rewrite outlives the call: later classifiers in the cascade score this value,
273        // and it is what reaches the model.
274        assert_eq!(request.model_id().as_deref(), Some("rewritten"));
275        Ok(())
276    }
277}
switchyard_libsy/core/classifier.rs

switchyard_libsy/core/
classifier.rs
