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