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

switchyard_libsy/core/
classifier.rs
