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