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