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, 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
24pub type RandomDecision = FallThroughDecision;
26
27pub struct RandomClassifier {
29 targets: Vec<String>,
30 distribution: WeightedIndex<f64>,
31 rng: Mutex<StdRng>,
32}
33
34impl RandomClassifier {
35 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
125pub struct Random {
127 inner: FallThrough<()>,
128}
129
130impl Random {
131 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 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}