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