Skip to main content

switchyard_libsy/algorithms/
rand.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Random routing as a stateless [`FallThrough`] composition.
5//!
6//! [`RandomClassifier`] selects one target; [`FallThrough`] owns the common
7//! processor/classifier/target-call orchestration.
8
9use std::collections::BTreeSet;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use parking_lot::Mutex;
14use rand::SeedableRng;
15use rand::distr::{Distribution, weighted::WeightedIndex};
16use rand::rngs::StdRng;
17
18use crate::algorithms::fall_through::FallThrough;
19use crate::core::algorithm::{Algorithm, Driver};
20use crate::core::classifier::{Classification, Classifier, Score};
21use crate::{LibsyError, Result};
22use switchyard_protocol::{ModelId, Request, Response};
23
24/// Stateless weighted classifier used by random fall-through routing.
25pub struct RandomClassifier {
26    targets: Vec<ModelId>,
27    distribution: WeightedIndex<f64>,
28    rng: Mutex<StdRng>,
29}
30
31impl RandomClassifier {
32    /// Creates a classifier over ordered target names.
33    ///
34    /// Missing weights default to one per target. Explicit weights are relative,
35    /// follow target order, and need not sum to one. Zero disables a target.
36    /// Missing `seed` uses entropy-backed randomness.
37    ///
38    /// # Errors
39    ///
40    /// Returns an error when targets are empty or duplicated, or when explicit
41    /// weights have the wrong length, are negative or non-finite, or contain no
42    /// positive value.
43    pub fn new(
44        targets: Vec<ModelId>,
45        weights: Option<Vec<f64>>,
46        seed: Option<u64>,
47    ) -> Result<Self> {
48        let target_count = targets.len();
49        if target_count == 0 {
50            return Err(LibsyError::NoTargets);
51        }
52        let unique_targets = targets.iter().map(ModelId::as_str).collect::<BTreeSet<_>>();
53        if unique_targets.len() != target_count {
54            return Err(LibsyError::AlgorithmError {
55                message: "random targets must be unique".to_string(),
56            });
57        }
58
59        let weights = weights.unwrap_or_else(|| vec![1.0; target_count]);
60        if weights.len() != target_count {
61            return Err(invalid_weights(format!(
62                "expected {target_count} weights, got {}",
63                weights.len()
64            )));
65        }
66        if weights
67            .iter()
68            .any(|weight| !weight.is_finite() || *weight < 0.0)
69        {
70            return Err(invalid_weights(
71                "weights must be finite and nonnegative".to_string(),
72            ));
73        }
74        if !weights.iter().any(|weight| *weight > 0.0) {
75            return Err(invalid_weights(
76                "at least one weight must be positive".to_string(),
77            ));
78        }
79        let distribution =
80            WeightedIndex::new(weights).map_err(|error| invalid_weights(error.to_string()))?;
81        let rng = match seed {
82            Some(seed) => StdRng::seed_from_u64(seed),
83            None => rand::make_rng(),
84        };
85        Ok(Self {
86            targets,
87            distribution,
88            rng: Mutex::new(rng),
89        })
90    }
91
92    fn select_target(&self) -> ModelId {
93        let mut rng = self.rng.lock();
94        let index = self.distribution.sample(&mut *rng);
95        self.targets[index].clone()
96    }
97}
98
99fn invalid_weights(message: String) -> LibsyError {
100    LibsyError::AlgorithmError {
101        message: format!("invalid random weights: {message}"),
102    }
103}
104
105#[async_trait]
106impl<S> Classifier<S> for RandomClassifier
107where
108    S: Send + 'static,
109{
110    async fn score(
111        &self,
112        _state: &mut S,
113        _request: &mut Request,
114        _driver: Option<&Driver>,
115    ) -> Result<(Classification, Option<Response>)> {
116        Ok((
117            Classification::Scores(vec![Score {
118                confidence: 1.0,
119                target: self.select_target(),
120            }]),
121            None,
122        ))
123    }
124}
125
126/// Random router implemented as a stateless fall-through composition.
127pub struct Random {
128    inner: FallThrough<()>,
129}
130
131impl Random {
132    /// Creates a router over `targets`.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error when targets or weights are invalid for [`RandomClassifier`].
137    pub fn new(
138        targets: Vec<ModelId>,
139        weights: Option<Vec<f64>>,
140        seed: Option<u64>,
141    ) -> Result<Self> {
142        let classifier = Arc::new(RandomClassifier::new(targets.clone(), weights, seed)?);
143        let inner = FallThrough::<()>::new(targets)
144            .with_name("random")
145            .with_decision_reason(random_decision_reason)
146            .with_classifier(classifier);
147        Ok(Self { inner })
148    }
149}
150
151fn random_decision_reason(_name: &str, winner: &Score) -> String {
152    format!("random routing selected target '{}'", winner.target)
153}
154
155#[async_trait]
156impl Algorithm for Random {
157    fn name(&self) -> &str {
158        "random"
159    }
160
161    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
162        self.inner.execute(driver, request).await
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use std::collections::HashSet;
170
171    use switchyard_protocol::{Metadata, completion_text, text_request};
172
173    use crate::algorithms::util::affinity::AffinityRouter;
174    use crate::core::testing::{echo, test_drive};
175    use switchyard_protocol::Request;
176
177    fn request() -> Request {
178        Request {
179            llm_request: text_request(Some("auto".to_string()), "hi"),
180            raw_request: None,
181            metadata: None,
182        }
183    }
184
185    fn request_for_session(session_id: &str) -> Request {
186        Request {
187            metadata: Some(Metadata {
188                session_id: Some(session_id.to_string()),
189                ..Metadata::default()
190            }),
191            ..request()
192        }
193    }
194
195    fn target_set(names: &[&str]) -> Vec<ModelId> {
196        names.iter().map(|name| ModelId::from(*name)).collect()
197    }
198
199    fn algorithm(names: &[&str], weights: Option<Vec<f64>>, seed: Option<u64>) -> Result<Random> {
200        Random::new(target_set(names), weights, seed)
201    }
202
203    fn shared_algorithm(names: &[&str]) -> Result<Arc<dyn Algorithm>> {
204        Ok(Arc::new(algorithm(names, None, None)?))
205    }
206
207    async fn selected_models(algorithm: Arc<dyn Algorithm>, count: usize) -> Result<Vec<String>> {
208        let mut selected = Vec::with_capacity(count);
209        for _ in 0..count {
210            let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
211            selected.push(
212                response
213                    .llm_response
214                    .as_agg()
215                    .map(completion_text)
216                    .unwrap_or_default(),
217            );
218        }
219        Ok(selected)
220    }
221
222    #[tokio::test]
223    async fn single_target_is_always_selected_and_called() -> Result<()> {
224        let algorithm = shared_algorithm(&["only/model"])?;
225        let (trace, response) = test_drive(algorithm, request(), echo()).await?;
226
227        assert_eq!(
228            response
229                .llm_response
230                .as_agg()
231                .map(completion_text)
232                .unwrap_or_default(),
233            "only/model"
234        );
235        assert_eq!(trace.len(), 1);
236        assert_eq!(trace[0].selected_model_id(), "only/model");
237        Ok(())
238    }
239
240    #[tokio::test]
241    async fn selected_target_is_in_the_set_and_matches_the_trace() -> Result<()> {
242        let names = ["a/model", "b/model", "c/model"];
243        let algorithm = shared_algorithm(&names)?;
244
245        for _ in 0..50 {
246            let (trace, response) = test_drive(algorithm.clone(), request(), echo()).await?;
247            let selected = response
248                .llm_response
249                .as_agg()
250                .map(completion_text)
251                .unwrap_or_default();
252            assert!(
253                names.contains(&selected.as_str()),
254                "selected {selected} not in target set"
255            );
256            assert_eq!(trace[0].selected_model_id(), selected.as_str());
257        }
258        Ok(())
259    }
260
261    #[tokio::test]
262    async fn selection_covers_all_targets_over_many_runs() -> Result<()> {
263        let algorithm = shared_algorithm(&["a/model", "b/model"])?;
264        let mut seen = HashSet::new();
265
266        for _ in 0..100 {
267            let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
268            seen.insert(
269                response
270                    .llm_response
271                    .as_agg()
272                    .map(completion_text)
273                    .unwrap_or_default(),
274            );
275        }
276
277        // Missing either target after 100 uniform draws has probability about 2^-99.
278        assert_eq!(
279            seen.len(),
280            2,
281            "expected both targets to be selected, saw {seen:?}"
282        );
283        Ok(())
284    }
285
286    #[tokio::test]
287    async fn weighted_seeded_selection_is_reproducible() -> Result<()> {
288        let first: Arc<dyn Algorithm> = Arc::new(algorithm(
289            &["a/model", "b/model"],
290            Some(vec![1.0, 3.0]),
291            Some(42),
292        )?);
293        let second: Arc<dyn Algorithm> = Arc::new(algorithm(
294            &["a/model", "b/model"],
295            Some(vec![1.0, 3.0]),
296            Some(42),
297        )?);
298
299        let first_selections = selected_models(first, 1_000).await?;
300        let second_selections = selected_models(second, 1_000).await?;
301        assert_eq!(first_selections, second_selections);
302
303        let second_count = first_selections
304            .iter()
305            .filter(|model| model.as_str() == "b/model")
306            .count();
307        assert!(
308            (700..=800).contains(&second_count),
309            "expected a roughly 25/75 split, selected b/model {second_count} times"
310        );
311        Ok(())
312    }
313
314    #[tokio::test]
315    async fn affinity_reuses_the_initial_random_selection() -> Result<()> {
316        let names = ["a/model", "b/model"];
317        let affinity = Arc::new(AffinityRouter::new());
318        let random = Arc::new(RandomClassifier::new(
319            names.iter().map(|name| ModelId::from(*name)).collect(),
320            None,
321            Some(42),
322        )?);
323        let algorithm: Arc<dyn Algorithm> = Arc::new(
324            FallThrough::<()>::new(target_set(&names))
325                .with_name("affinity_random")
326                .with_processor(affinity.clone())
327                .with_classifier(affinity.clone())
328                .with_classifier(random),
329        );
330
331        let (_, first) =
332            test_drive(algorithm.clone(), request_for_session("session-1"), echo()).await?;
333        let selected = first
334            .llm_response
335            .as_agg()
336            .map(completion_text)
337            .unwrap_or_default();
338
339        let mut state = ();
340        let mut request = request_for_session("session-1");
341        let retained = affinity
342            .score(&mut state, &mut request, None)
343            .await?
344            .0
345            .argmax(false)?;
346        assert_eq!(
347            retained.map(|score| score.target),
348            Some(ModelId::from(selected.clone()))
349        );
350
351        let (_, second) = test_drive(algorithm, request_for_session("session-1"), echo()).await?;
352        assert_eq!(
353            second
354                .llm_response
355                .as_agg()
356                .map(completion_text)
357                .unwrap_or_default(),
358            selected
359        );
360        Ok(())
361    }
362
363    #[test]
364    fn rejects_invalid_weights() {
365        let cases = [
366            (vec![1.0], "expected 2 weights"),
367            (vec![1.0, -1.0], "finite and nonnegative"),
368            (vec![0.0, 0.0], "at least one weight must be positive"),
369            (vec![1.0, f64::INFINITY], "finite and nonnegative"),
370        ];
371
372        for (weights, expected) in cases {
373            let error = algorithm(&["a/model", "b/model"], Some(weights), None)
374                .err()
375                .map(|error| error.to_string())
376                .unwrap_or_default();
377            assert!(error.contains(expected), "unexpected error: {error}");
378        }
379    }
380
381    #[test]
382    fn rejects_invalid_targets() {
383        let error = algorithm(&[], None, None).err();
384        assert!(matches!(error, Some(LibsyError::NoTargets)));
385
386        let error = algorithm(&["same/model", "same/model"], None, None)
387            .err()
388            .map(|error| error.to_string())
389            .unwrap_or_default();
390        assert!(error.contains("random targets must be unique"));
391    }
392
393    #[tokio::test]
394    async fn decision_is_inspectable() -> Result<()> {
395        let algorithm = shared_algorithm(&["only/model"])?;
396        let (trace, _) = test_drive(algorithm, request(), echo()).await?;
397        let decision = &trace[0];
398
399        assert_eq!(decision.selected_model_id(), "only/model");
400        assert!(
401            decision
402                .reasoning()
403                .unwrap_or_default()
404                .contains("only/model")
405        );
406        assert!(decision.is_answer_call());
407        Ok(())
408    }
409}