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