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