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