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