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