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 route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
163        self.inner.execute(driver, request).await
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use std::collections::HashSet;
171
172    use switchyard_protocol::{Metadata, completion_text, text_request};
173
174    use crate::algorithms::util::affinity::AffinityRouter;
175    use crate::core::algorithm::LlmTarget;
176    use crate::core::testing::{echo, test_drive};
177    use switchyard_protocol::Request;
178
179    fn request() -> Request {
180        Request {
181            llm_request: text_request(Some("auto".to_string()), "hi"),
182            raw_request: None,
183            metadata: None,
184        }
185    }
186
187    fn request_for_session(session_id: &str) -> Request {
188        Request {
189            metadata: Some(Metadata {
190                session_id: Some(session_id.to_string()),
191                ..Metadata::default()
192            }),
193            ..request()
194        }
195    }
196
197    fn target_set(names: &[&str]) -> LlmTargetSet {
198        let targets = names
199            .iter()
200            .map(|name| LlmTarget {
201                semantic_name: (*name).to_string(),
202            })
203            .collect();
204        LlmTargetSet::new(targets)
205    }
206
207    fn algorithm(names: &[&str], weights: Option<Vec<f64>>, seed: Option<u64>) -> Result<Random> {
208        Random::new(target_set(names), weights, seed)
209    }
210
211    fn shared_algorithm(names: &[&str]) -> Result<Arc<dyn Algorithm>> {
212        Ok(Arc::new(algorithm(names, None, None)?))
213    }
214
215    async fn selected_models(algorithm: Arc<dyn Algorithm>, count: usize) -> Result<Vec<String>> {
216        let mut selected = Vec::with_capacity(count);
217        for _ in 0..count {
218            let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
219            selected.push(
220                response
221                    .llm_response
222                    .as_agg()
223                    .map(completion_text)
224                    .unwrap_or_default(),
225            );
226        }
227        Ok(selected)
228    }
229
230    #[tokio::test]
231    async fn single_target_is_always_selected_and_called() -> Result<()> {
232        let algorithm = shared_algorithm(&["only/model"])?;
233        let (trace, response) = test_drive(algorithm, request(), echo()).await?;
234
235        assert_eq!(
236            response
237                .llm_response
238                .as_agg()
239                .map(completion_text)
240                .unwrap_or_default(),
241            "only/model"
242        );
243        assert_eq!(trace.len(), 1);
244        assert_eq!(trace[0].selected_model_id(), "only/model");
245        Ok(())
246    }
247
248    #[tokio::test]
249    async fn selected_target_is_in_the_set_and_matches_the_trace() -> Result<()> {
250        let names = ["a/model", "b/model", "c/model"];
251        let algorithm = shared_algorithm(&names)?;
252
253        for _ in 0..50 {
254            let (trace, response) = test_drive(algorithm.clone(), request(), echo()).await?;
255            let selected = response
256                .llm_response
257                .as_agg()
258                .map(completion_text)
259                .unwrap_or_default();
260            assert!(
261                names.contains(&selected.as_str()),
262                "selected {selected} not in target set"
263            );
264            assert_eq!(trace[0].selected_model_id(), selected.as_str());
265        }
266        Ok(())
267    }
268
269    #[tokio::test]
270    async fn selection_covers_all_targets_over_many_runs() -> Result<()> {
271        let algorithm = shared_algorithm(&["a/model", "b/model"])?;
272        let mut seen = HashSet::new();
273
274        for _ in 0..100 {
275            let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
276            seen.insert(
277                response
278                    .llm_response
279                    .as_agg()
280                    .map(completion_text)
281                    .unwrap_or_default(),
282            );
283        }
284
285        // Missing either target after 100 uniform draws has probability about 2^-99.
286        assert_eq!(
287            seen.len(),
288            2,
289            "expected both targets to be selected, saw {seen:?}"
290        );
291        Ok(())
292    }
293
294    #[tokio::test]
295    async fn weighted_seeded_selection_is_reproducible() -> Result<()> {
296        let first: Arc<dyn Algorithm> = Arc::new(algorithm(
297            &["a/model", "b/model"],
298            Some(vec![1.0, 3.0]),
299            Some(42),
300        )?);
301        let second: Arc<dyn Algorithm> = Arc::new(algorithm(
302            &["a/model", "b/model"],
303            Some(vec![1.0, 3.0]),
304            Some(42),
305        )?);
306
307        let first_selections = selected_models(first, 1_000).await?;
308        let second_selections = selected_models(second, 1_000).await?;
309        assert_eq!(first_selections, second_selections);
310
311        let second_count = first_selections
312            .iter()
313            .filter(|model| model.as_str() == "b/model")
314            .count();
315        assert!(
316            (700..=800).contains(&second_count),
317            "expected a roughly 25/75 split, selected b/model {second_count} times"
318        );
319        Ok(())
320    }
321
322    #[tokio::test]
323    async fn affinity_reuses_the_initial_random_selection() -> Result<()> {
324        let names = ["a/model", "b/model"];
325        let affinity = Arc::new(AffinityRouter::new());
326        let random = Arc::new(RandomClassifier::new(
327            names.iter().map(|name| (*name).to_string()).collect(),
328            None,
329            Some(42),
330        )?);
331        let algorithm: Arc<dyn Algorithm> = Arc::new(
332            FallThrough::<()>::new(target_set(&names))
333                .with_name("affinity_random")
334                .with_processor(affinity.clone())
335                .with_classifier(affinity.clone())
336                .with_classifier(random),
337        );
338
339        let (_, first) =
340            test_drive(algorithm.clone(), request_for_session("session-1"), echo()).await?;
341        let selected = first
342            .llm_response
343            .as_agg()
344            .map(completion_text)
345            .unwrap_or_default();
346
347        let mut state = ();
348        let mut request = request_for_session("session-1");
349        let retained = affinity
350            .score(&mut state, &mut request, None)
351            .await?
352            .0
353            .argmax(false)?;
354        assert_eq!(
355            retained.map(|score| score.target),
356            Some(selected.to_string())
357        );
358
359        let (_, second) = test_drive(algorithm, request_for_session("session-1"), echo()).await?;
360        assert_eq!(
361            second
362                .llm_response
363                .as_agg()
364                .map(completion_text)
365                .unwrap_or_default(),
366            selected
367        );
368        Ok(())
369    }
370
371    #[test]
372    fn rejects_invalid_weights() {
373        let cases = [
374            (vec![1.0], "expected 2 weights"),
375            (vec![1.0, -1.0], "finite and nonnegative"),
376            (vec![0.0, 0.0], "at least one weight must be positive"),
377            (vec![1.0, f64::INFINITY], "finite and nonnegative"),
378        ];
379
380        for (weights, expected) in cases {
381            let error = algorithm(&["a/model", "b/model"], Some(weights), None)
382                .err()
383                .map(|error| error.to_string())
384                .unwrap_or_default();
385            assert!(error.contains(expected), "unexpected error: {error}");
386        }
387    }
388
389    #[test]
390    fn rejects_invalid_targets() {
391        let error = algorithm(&[], None, None).err();
392        assert!(matches!(error, Some(LibsyError::NoTargets)));
393
394        let error = algorithm(&["same/model", "same/model"], None, None)
395            .err()
396            .map(|error| error.to_string())
397            .unwrap_or_default();
398        assert!(error.contains("random targets must be unique"));
399    }
400
401    #[tokio::test]
402    async fn decision_is_inspectable() -> Result<()> {
403        let algorithm = shared_algorithm(&["only/model"])?;
404        let (trace, _) = test_drive(algorithm, request(), echo()).await?;
405        let decision = &trace[0];
406
407        assert_eq!(decision.selected_model_id(), "only/model");
408        assert!(
409            decision
410                .reasoning()
411                .unwrap_or_default()
412                .contains("only/model")
413        );
414        assert!(decision.is_answer_call());
415        Ok(())
416    }
417}