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