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, LlmTargetSet};
20use crate::core::classifier::{Classification, Classifier, Score};
21use crate::{LibsyError, Result};
22use switchyard_protocol::{Request, Response};
23
24pub struct RandomClassifier {
26 targets: Vec<String>,
27 distribution: WeightedIndex<f64>,
28 rng: Mutex<StdRng>,
29}
30
31impl RandomClassifier {
32 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
122pub struct Random {
124 inner: FallThrough<()>,
125}
126
127impl Random {
128 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 async fn create_run_task(
163 self: Arc<Self>,
164 driver: Driver,
165 request: Request,
166 ) -> Result<Response> {
167 self.inner.execute(driver, request).await
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174 use std::collections::HashSet;
175
176 use switchyard_protocol::{Metadata, completion_text, text_request};
177
178 use crate::algorithms::util::affinity::AffinityRouter;
179 use crate::core::algorithm::LlmTarget;
180 use crate::core::testing::{echo, test_drive};
181 use switchyard_protocol::{Request, Signals};
182
183 fn request() -> Request {
184 Request {
185 llm_request: text_request(Some("auto".to_string()), "hi"),
186 raw_request: None,
187 metadata: None,
188 }
189 }
190
191 fn request_for_session(session_id: &str) -> Request {
192 Request {
193 metadata: Some(Metadata {
194 session_id: Some(session_id.to_string()),
195 ..Metadata::default()
196 }),
197 ..request()
198 }
199 }
200
201 fn target_set(names: &[&str]) -> LlmTargetSet {
202 let targets = names
203 .iter()
204 .map(|name| LlmTarget {
205 semantic_name: (*name).to_string(),
206 })
207 .collect();
208 LlmTargetSet::new(targets)
209 }
210
211 fn algorithm(names: &[&str], weights: Option<Vec<f64>>, seed: Option<u64>) -> Result<Random> {
212 Random::new(target_set(names), weights, seed)
213 }
214
215 fn shared_algorithm(names: &[&str]) -> Result<Arc<dyn Algorithm>> {
216 Ok(Arc::new(algorithm(names, None, None)?))
217 }
218
219 async fn selected_models(algorithm: Arc<dyn Algorithm>, count: usize) -> Result<Vec<String>> {
220 let mut selected = Vec::with_capacity(count);
221 for _ in 0..count {
222 let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
223 selected.push(
224 response
225 .llm_response
226 .as_agg()
227 .map(completion_text)
228 .unwrap_or_default(),
229 );
230 }
231 Ok(selected)
232 }
233
234 #[tokio::test]
235 async fn single_target_is_always_selected_and_called() -> Result<()> {
236 let algorithm = shared_algorithm(&["only/model"])?;
237 let (trace, response) = test_drive(algorithm, request(), echo()).await?;
238
239 assert_eq!(
240 response
241 .llm_response
242 .as_agg()
243 .map(completion_text)
244 .unwrap_or_default(),
245 "only/model"
246 );
247 assert_eq!(trace.len(), 1);
248 assert_eq!(trace[0].selected_model_id(), "only/model");
249 Ok(())
250 }
251
252 #[tokio::test]
253 async fn selected_target_is_in_the_set_and_matches_the_trace() -> Result<()> {
254 let names = ["a/model", "b/model", "c/model"];
255 let algorithm = shared_algorithm(&names)?;
256
257 for _ in 0..50 {
258 let (trace, response) = test_drive(algorithm.clone(), request(), echo()).await?;
259 let selected = response
260 .llm_response
261 .as_agg()
262 .map(completion_text)
263 .unwrap_or_default();
264 assert!(
265 names.contains(&selected.as_str()),
266 "selected {selected} not in target set"
267 );
268 assert_eq!(trace[0].selected_model_id(), selected.as_str());
269 }
270 Ok(())
271 }
272
273 #[tokio::test]
274 async fn selection_covers_all_targets_over_many_runs() -> Result<()> {
275 let algorithm = shared_algorithm(&["a/model", "b/model"])?;
276 let mut seen = HashSet::new();
277
278 for _ in 0..100 {
279 let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?;
280 seen.insert(
281 response
282 .llm_response
283 .as_agg()
284 .map(completion_text)
285 .unwrap_or_default(),
286 );
287 }
288
289 assert_eq!(
291 seen.len(),
292 2,
293 "expected both targets to be selected, saw {seen:?}"
294 );
295 Ok(())
296 }
297
298 #[tokio::test]
299 async fn weighted_seeded_selection_is_reproducible() -> Result<()> {
300 let first: Arc<dyn Algorithm> = Arc::new(algorithm(
301 &["a/model", "b/model"],
302 Some(vec![1.0, 3.0]),
303 Some(42),
304 )?);
305 let second: Arc<dyn Algorithm> = Arc::new(algorithm(
306 &["a/model", "b/model"],
307 Some(vec![1.0, 3.0]),
308 Some(42),
309 )?);
310
311 let first_selections = selected_models(first, 1_000).await?;
312 let second_selections = selected_models(second, 1_000).await?;
313 assert_eq!(first_selections, second_selections);
314
315 let second_count = first_selections
316 .iter()
317 .filter(|model| model.as_str() == "b/model")
318 .count();
319 assert!(
320 (700..=800).contains(&second_count),
321 "expected a roughly 25/75 split, selected b/model {second_count} times"
322 );
323 Ok(())
324 }
325
326 #[tokio::test]
327 async fn affinity_reuses_the_initial_random_selection() -> Result<()> {
328 let names = ["a/model", "b/model"];
329 let affinity = Arc::new(AffinityRouter::new());
330 let random = Arc::new(RandomClassifier::new(
331 names.iter().map(|name| (*name).to_string()).collect(),
332 None,
333 Some(42),
334 )?);
335 let algorithm: Arc<dyn Algorithm> = Arc::new(
336 FallThrough::<()>::new(target_set(&names))
337 .with_name("affinity_random")
338 .with_processor(affinity.clone())
339 .with_classifier(affinity.clone())
340 .with_classifier(random),
341 );
342
343 let (_, first) =
344 test_drive(algorithm.clone(), request_for_session("session-1"), echo()).await?;
345 let selected = first
346 .llm_response
347 .as_agg()
348 .map(completion_text)
349 .unwrap_or_default();
350
351 let mut state = ();
352 let mut request = request_for_session("session-1");
353 let retained = affinity
354 .score(&mut state, &mut request, None)
355 .await?
356 .0
357 .argmax(false)?;
358 assert_eq!(
359 retained.map(|score| score.target),
360 Some(selected.to_string())
361 );
362
363 let (_, second) = test_drive(algorithm, request_for_session("session-1"), echo()).await?;
364 assert_eq!(
365 second
366 .llm_response
367 .as_agg()
368 .map(completion_text)
369 .unwrap_or_default(),
370 selected
371 );
372 Ok(())
373 }
374
375 #[test]
376 fn rejects_invalid_weights() {
377 let cases = [
378 (vec![1.0], "expected 2 weights"),
379 (vec![1.0, -1.0], "finite and nonnegative"),
380 (vec![0.0, 0.0], "at least one weight must be positive"),
381 (vec![1.0, f64::INFINITY], "finite and nonnegative"),
382 ];
383
384 for (weights, expected) in cases {
385 let error = algorithm(&["a/model", "b/model"], Some(weights), None)
386 .err()
387 .map(|error| error.to_string())
388 .unwrap_or_default();
389 assert!(error.contains(expected), "unexpected error: {error}");
390 }
391 }
392
393 #[test]
394 fn rejects_invalid_targets() {
395 let error = algorithm(&[], None, None).err();
396 assert!(matches!(error, Some(LibsyError::NoTargets)));
397
398 let error = algorithm(&["same/model", "same/model"], None, None)
399 .err()
400 .map(|error| error.to_string())
401 .unwrap_or_default();
402 assert!(error.contains("random targets must be unique"));
403 }
404
405 #[tokio::test]
406 async fn process_signals_is_a_noop() -> Result<()> {
407 let algorithm: Arc<dyn Algorithm> = Arc::new(algorithm(&["only/model"], None, None)?);
408 algorithm.process_signals(Signals {}).await?;
409 Ok(())
410 }
411
412 #[tokio::test]
413 async fn decision_is_inspectable() -> Result<()> {
414 let algorithm = shared_algorithm(&["only/model"])?;
415 let (trace, _) = test_drive(algorithm, request(), echo()).await?;
416 let decision = &trace[0];
417
418 assert_eq!(decision.selected_model_id(), "only/model");
419 assert!(
420 decision
421 .reasoning()
422 .unwrap_or_default()
423 .contains("only/model")
424 );
425 assert!(decision.is_answer_call());
426 Ok(())
427 }
428}