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::{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: String,
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 /// Stable tier represented by `selected_model`, when this classifier defines one.
75 fn routing_tier(&self, _selected_model: &str) -> Option<&'static str> {
76 None
77 }
78
79 /// Drops retained routing state when `target` was unavailable for `request`.
80 ///
81 /// Stateless classifiers do not need to implement this hook.
82 fn target_unavailable(&self, _request: &Request, _target: &str) {}
83
84 /// Score the classifier's targets given the current state and request.
85 ///
86 /// When present, `driver` lets a classifier offload model calls. It is `None`
87 /// when the classifier is evaluated outside an algorithm run.
88 ///
89 /// `request` is borrowed mutably so a classifier may rewrite it in place — inject a
90 /// system prompt, drop tools, compact history. The edit is not scoped to this call:
91 /// later classifiers in the cascade score the rewritten request, and it is the
92 /// rewritten request that is finally sent to the selected model. Most classifiers
93 /// only read it.
94 async fn score(
95 &self,
96 state: &mut S,
97 request: &mut Request,
98 driver: Option<&Driver>,
99 ) -> Result<(Classification, Option<Response>)>;
100}
101
102#[cfg(test)]
103#[path = "classifier_tests.rs"]
104mod tests;