switchyard_libsy/algorithms/
rand.rs1use 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};
20use crate::core::classifier::{Classification, Classifier, Score};
21use crate::{LibsyError, Result};
22use switchyard_protocol::{ModelId, Request, Response};
23
24pub struct RandomClassifier {
26 targets: Vec<ModelId>,
27 distribution: WeightedIndex<f64>,
28 rng: Mutex<StdRng>,
29}
30
31impl RandomClassifier {
32 pub fn new(
44 targets: Vec<ModelId>,
45 weights: Option<Vec<f64>>,
46 seed: Option<u64>,
47 ) -> Result<Self> {
48 let target_count = targets.len();
49 if target_count == 0 {
50 return Err(LibsyError::NoTargets);
51 }
52 let unique_targets = targets.iter().map(ModelId::as_str).collect::<BTreeSet<_>>();
53 if unique_targets.len() != target_count {
54 return Err(LibsyError::AlgorithmError {
55 message: "random targets must be unique".to_string(),
56 });
57 }
58
59 let weights = weights.unwrap_or_else(|| vec![1.0; target_count]);
60 if weights.len() != target_count {
61 return Err(invalid_weights(format!(
62 "expected {target_count} weights, got {}",
63 weights.len()
64 )));
65 }
66 if weights
67 .iter()
68 .any(|weight| !weight.is_finite() || *weight < 0.0)
69 {
70 return Err(invalid_weights(
71 "weights must be finite and nonnegative".to_string(),
72 ));
73 }
74 if !weights.iter().any(|weight| *weight > 0.0) {
75 return Err(invalid_weights(
76 "at least one weight must be positive".to_string(),
77 ));
78 }
79 let distribution =
80 WeightedIndex::new(weights).map_err(|error| invalid_weights(error.to_string()))?;
81 let rng = match seed {
82 Some(seed) => StdRng::seed_from_u64(seed),
83 None => rand::make_rng(),
84 };
85 Ok(Self {
86 targets,
87 distribution,
88 rng: Mutex::new(rng),
89 })
90 }
91
92 fn select_target(&self) -> ModelId {
93 let mut rng = self.rng.lock();
94 let index = self.distribution.sample(&mut *rng);
95 self.targets[index].clone()
96 }
97}
98
99fn invalid_weights(message: String) -> LibsyError {
100 LibsyError::AlgorithmError {
101 message: format!("invalid random weights: {message}"),
102 }
103}
104
105#[async_trait]
106impl<S> Classifier<S> for RandomClassifier
107where
108 S: Send + 'static,
109{
110 async fn score(
111 &self,
112 _state: &mut S,
113 _request: &mut Request,
114 _driver: Option<&Driver>,
115 ) -> Result<(Classification, Option<Response>)> {
116 Ok((
117 Classification::Scores(vec![Score {
118 confidence: 1.0,
119 target: self.select_target(),
120 }]),
121 None,
122 ))
123 }
124}
125
126pub struct Random {
128 inner: FallThrough<()>,
129}
130
131impl Random {
132 pub fn new(
138 targets: Vec<ModelId>,
139 weights: Option<Vec<f64>>,
140 seed: Option<u64>,
141 ) -> Result<Self> {
142 let classifier = Arc::new(RandomClassifier::new(targets.clone(), weights, seed)?);
143 let inner = FallThrough::<()>::new(targets)
144 .with_name("random")
145 .with_decision_reason(random_decision_reason)
146 .with_classifier(classifier);
147 Ok(Self { inner })
148 }
149}
150
151fn random_decision_reason(_name: &str, winner: &Score) -> String {
152 format!("random routing selected target '{}'", winner.target)
153}
154
155#[async_trait]
156impl Algorithm for Random {
157 fn name(&self) -> &str {
158 "random"
159 }
160
161 async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<Response> {
162 self.inner.execute(driver, request).await
163 }
164}
165
166#[cfg(test)]
167mod tests {
168 use super::*;
169 use std::collections::HashSet;
170
171 use switchyard_protocol::{Metadata, completion_text, text_request};
172
173 use crate::algorithms::util::affinity::AffinityRouter;
174 use crate::core::testing::{echo, test_drive};
175 use switchyard_protocol::Request;
176
177 fn request() -> Request {
178 Request {
179 llm_request: text_request(Some("auto".to_string()), "hi"),
180 raw_request: None,
181 metadata: None,
182 }
183 }
184
185 fn request_for_session(session_id: &str) -> Request {
186 Request {
187 metadata: Some(Metadata {
188 session_id: Some(session_id.to_string()),
189 ..Metadata::default()
190 }),
191 ..request()
192 }
193 }
194
195 fn target_set(names: &[&str]) -> Vec<ModelId> {
196 names.iter().map(|name| ModelId::from(*name)).collect()
197 }
198
199 fn algorithm(names: &[&str], weights: Option<Vec<f64>>, seed: Option<u64>) -> Result<Random> {
200 Random::new(target_set(names), weights, seed)
201 }
202
203 fn shared_algorithm(names: &[&str]) -> Result<Arc<dyn Algorithm>> {
204 Ok(Arc::new(algorithm(names, None, None)?))
205 }
206
207 async fn selected_models(algorithm: Arc<dyn Algorithm>, count: usize) -> Result<Vec<String>> {
208 let mut selected = Vec::with_capacity(count);
209 for _ in 0..count {
210 let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
211 selected.push(
212 response
213 .llm_response
214 .as_agg()
215 .map(completion_text)
216 .unwrap_or_default(),
217 );
218 }
219 Ok(selected)
220 }
221
222 #[tokio::test]
223 async fn single_target_is_always_selected_and_called() -> Result<()> {
224 let algorithm = shared_algorithm(&["only/model"])?;
225 let (trace, response) = test_drive(algorithm, request(), echo()).await?;
226
227 assert_eq!(
228 response
229 .llm_response
230 .as_agg()
231 .map(completion_text)
232 .unwrap_or_default(),
233 "only/model"
234 );
235 assert_eq!(trace.len(), 1);
236 assert_eq!(trace[0].selected_model_id(), "only/model");
237 Ok(())
238 }
239
240 #[tokio::test]
241 async fn selected_target_is_in_the_set_and_matches_the_trace() -> Result<()> {
242 let names = ["a/model", "b/model", "c/model"];
243 let algorithm = shared_algorithm(&names)?;
244
245 for _ in 0..50 {
246 let (trace, response) = test_drive(algorithm.clone(), request(), echo()).await?;
247 let selected = response
248 .llm_response
249 .as_agg()
250 .map(completion_text)
251 .unwrap_or_default();
252 assert!(
253 names.contains(&selected.as_str()),
254 "selected {selected} not in target set"
255 );
256 assert_eq!(trace[0].selected_model_id(), selected.as_str());
257 }
258 Ok(())
259 }
260
261 #[tokio::test]
262 async fn selection_covers_all_targets_over_many_runs() -> Result<()> {
263 let algorithm = shared_algorithm(&["a/model", "b/model"])?;
264 let mut seen = HashSet::new();
265
266 for _ in 0..100 {
267 let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
268 seen.insert(
269 response
270 .llm_response
271 .as_agg()
272 .map(completion_text)
273 .unwrap_or_default(),
274 );
275 }
276
277 assert_eq!(
279 seen.len(),
280 2,
281 "expected both targets to be selected, saw {seen:?}"
282 );
283 Ok(())
284 }
285
286 #[tokio::test]
287 async fn weighted_seeded_selection_is_reproducible() -> Result<()> {
288 let first: Arc<dyn Algorithm> = Arc::new(algorithm(
289 &["a/model", "b/model"],
290 Some(vec![1.0, 3.0]),
291 Some(42),
292 )?);
293 let second: Arc<dyn Algorithm> = Arc::new(algorithm(
294 &["a/model", "b/model"],
295 Some(vec![1.0, 3.0]),
296 Some(42),
297 )?);
298
299 let first_selections = selected_models(first, 1_000).await?;
300 let second_selections = selected_models(second, 1_000).await?;
301 assert_eq!(first_selections, second_selections);
302
303 let second_count = first_selections
304 .iter()
305 .filter(|model| model.as_str() == "b/model")
306 .count();
307 assert!(
308 (700..=800).contains(&second_count),
309 "expected a roughly 25/75 split, selected b/model {second_count} times"
310 );
311 Ok(())
312 }
313
314 #[tokio::test]
315 async fn affinity_reuses_the_initial_random_selection() -> Result<()> {
316 let names = ["a/model", "b/model"];
317 let affinity = Arc::new(AffinityRouter::new());
318 let random = Arc::new(RandomClassifier::new(
319 names.iter().map(|name| ModelId::from(*name)).collect(),
320 None,
321 Some(42),
322 )?);
323 let algorithm: Arc<dyn Algorithm> = Arc::new(
324 FallThrough::<()>::new(target_set(&names))
325 .with_name("affinity_random")
326 .with_processor(affinity.clone())
327 .with_classifier(affinity.clone())
328 .with_classifier(random),
329 );
330
331 let (_, first) =
332 test_drive(algorithm.clone(), request_for_session("session-1"), echo()).await?;
333 let selected = first
334 .llm_response
335 .as_agg()
336 .map(completion_text)
337 .unwrap_or_default();
338
339 let mut state = ();
340 let mut request = request_for_session("session-1");
341 let retained = affinity
342 .score(&mut state, &mut request, None)
343 .await?
344 .0
345 .argmax(false)?;
346 assert_eq!(
347 retained.map(|score| score.target),
348 Some(ModelId::from(selected.clone()))
349 );
350
351 let (_, second) = test_drive(algorithm, request_for_session("session-1"), echo()).await?;
352 assert_eq!(
353 second
354 .llm_response
355 .as_agg()
356 .map(completion_text)
357 .unwrap_or_default(),
358 selected
359 );
360 Ok(())
361 }
362
363 #[test]
364 fn rejects_invalid_weights() {
365 let cases = [
366 (vec![1.0], "expected 2 weights"),
367 (vec![1.0, -1.0], "finite and nonnegative"),
368 (vec![0.0, 0.0], "at least one weight must be positive"),
369 (vec![1.0, f64::INFINITY], "finite and nonnegative"),
370 ];
371
372 for (weights, expected) in cases {
373 let error = algorithm(&["a/model", "b/model"], Some(weights), None)
374 .err()
375 .map(|error| error.to_string())
376 .unwrap_or_default();
377 assert!(error.contains(expected), "unexpected error: {error}");
378 }
379 }
380
381 #[test]
382 fn rejects_invalid_targets() {
383 let error = algorithm(&[], None, None).err();
384 assert!(matches!(error, Some(LibsyError::NoTargets)));
385
386 let error = algorithm(&["same/model", "same/model"], None, None)
387 .err()
388 .map(|error| error.to_string())
389 .unwrap_or_default();
390 assert!(error.contains("random targets must be unique"));
391 }
392
393 #[tokio::test]
394 async fn decision_is_inspectable() -> Result<()> {
395 let algorithm = shared_algorithm(&["only/model"])?;
396 let (trace, _) = test_drive(algorithm, request(), echo()).await?;
397 let decision = &trace[0];
398
399 assert_eq!(decision.selected_model_id(), "only/model");
400 assert!(
401 decision
402 .reasoning()
403 .unwrap_or_default()
404 .contains("only/model")
405 );
406 assert!(decision.is_answer_call());
407 Ok(())
408 }
409}