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

switchyard_libsy/core/
classifier.rs
