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